package org.raddatz.familienarchiv.ocr;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.raddatz.familienarchiv.ocr.TrainingHistoryResponse;
import org.raddatz.familienarchiv.ocr.TrainingInfoResponse;
import org.raddatz.familienarchiv.exception.DomainException;
import org.raddatz.familienarchiv.exception.ErrorCode;
import org.raddatz.familienarchiv.ocr.OcrTrainingRun;
import org.raddatz.familienarchiv.ocr.SenderModel;
import org.raddatz.familienarchiv.ocr.TrainingStatus;
import org.raddatz.familienarchiv.ocr.OcrTrainingRunRepository;
import org.raddatz.familienarchiv.person.PersonService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockQueryService;
import org.slf4j.MDC;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionTemplate;

import java.io.ByteArrayOutputStream;
import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;

@Service
@RequiredArgsConstructor
@Slf4j
public class OcrTrainingService {

    private final OcrTrainingRunRepository trainingRunRepository;
    private final TrainingDataExportService trainingDataExportService;
    private final SegmentationTrainingExportService segmentationTrainingExportService;
    private final OcrClient ocrClient;
    private final OcrHealthClient ocrHealthClient;
    private final TranscriptionBlockQueryService transcriptionBlockQueryService;
    private final TransactionTemplate txTemplate;
    private final PersonService personService;
    private final SenderModelService senderModelService;

    private void assertNoRunningTraining() {
        if (trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING).isPresent()) {
            throw DomainException.conflict(ErrorCode.TRAINING_ALREADY_RUNNING,
                    "A training run is already in progress");
        }
    }

    // Not safe for horizontal scaling: training reloads the Kraken model in-process on the
    // Python OCR service after each run. The DB-level RUNNING constraint (V30 partial unique
    // index) prevents concurrent training API calls, but cannot prevent two OCR service replicas
    // from diverging on model state. Deploy as a single instance only. See ADR-001.
    public OcrTrainingRun triggerTraining(UUID triggeredBy) {
        // Short transaction: guard check + create RUNNING row, then commit immediately.
        // The DB connection is released before the OCR HTTP call, which can take several minutes.
        OcrTrainingRun run = Objects.requireNonNull(txTemplate.execute(status -> {
            assertNoRunningTraining();

            var eligibleBlocks = trainingDataExportService.queryEligibleBlocks();
            if (eligibleBlocks.size() < 5) {
                throw DomainException.badRequest(ErrorCode.VALIDATION_ERROR,
                        "At least 5 eligible blocks are required to start training (found " + eligibleBlocks.size() + ")");
            }

            long documentCount = eligibleBlocks.stream()
                    .map(b -> b.getDocumentId())
                    .distinct()
                    .count();

            OcrTrainingRun newRun = OcrTrainingRun.builder()
                    .status(TrainingStatus.RUNNING)
                    .blockCount(eligibleBlocks.size())
                    .documentCount((int) documentCount)
                    .modelName("german_kurrent")
                    .triggeredBy(triggeredBy)
                    .build();
            return trainingRunRepository.save(newRun);
        }));

        String runId = run.getId().toString();
        MDC.put("trainingRunId", runId);
        log.info("Started training run {} with {} blocks from {} documents",
                runId, run.getBlockCount(), run.getDocumentCount());

        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            trainingDataExportService.exportToZip().writeTo(baos);
            byte[] zipBytes = baos.toByteArray();

            log.info("[trainingRun={}] Sending {} bytes to OCR service", runId, zipBytes.length);
            OcrClient.TrainingResult result = ocrClient.trainModel(zipBytes);

            return Objects.requireNonNull(txTemplate.execute(status -> {
                run.setStatus(TrainingStatus.DONE);
                run.setCompletedAt(Instant.now());
                run.setCer(result.cer());
                run.setLoss(result.loss());
                run.setAccuracy(result.accuracy());
                run.setEpochs(result.epochs());
                OcrTrainingRun updated = trainingRunRepository.save(run);
                log.info("[trainingRun={}] Training completed — cer={} epochs={}", runId, result.cer(), result.epochs());
                return updated;
            }));
        } catch (Exception e) {
            return Objects.requireNonNull(txTemplate.execute(status -> {
                run.setStatus(TrainingStatus.FAILED);
                run.setErrorMessage(e.getMessage());
                run.setCompletedAt(Instant.now());
                OcrTrainingRun failed = trainingRunRepository.save(run);
                log.error("[trainingRun={}] Training failed: {}", runId, e.getMessage(), e);
                return failed;
            }));
        } finally {
            MDC.remove("trainingRunId");
        }
    }

    public OcrTrainingRun triggerSegTraining(UUID triggeredBy) {
        // Same pattern as triggerTraining: narrow transactions around DB writes only.
        OcrTrainingRun run = Objects.requireNonNull(txTemplate.execute(status -> {
            assertNoRunningTraining();

            var segBlocks = segmentationTrainingExportService.querySegmentationBlocks();
            if (segBlocks.size() < 5) {
                throw DomainException.badRequest(ErrorCode.VALIDATION_ERROR,
                        "At least 5 eligible segments are required to start training (found " + segBlocks.size() + ")");
            }

            long documentCount = segBlocks.stream()
                    .map(b -> b.getDocumentId())
                    .distinct()
                    .count();

            OcrTrainingRun newRun = OcrTrainingRun.builder()
                    .status(TrainingStatus.RUNNING)
                    .blockCount(segBlocks.size())
                    .documentCount((int) documentCount)
                    .modelName("blla")
                    .triggeredBy(triggeredBy)
                    .build();
            return trainingRunRepository.save(newRun);
        }));

        String runId = run.getId().toString();
        MDC.put("trainingRunId", runId);
        log.info("Started segmentation training run {} with {} segments from {} documents",
                runId, run.getBlockCount(), run.getDocumentCount());

        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            segmentationTrainingExportService.exportToZip().writeTo(baos);
            byte[] zipBytes = baos.toByteArray();

            log.info("[trainingRun={}] Sending {} bytes to OCR service for segtrain", runId, zipBytes.length);
            OcrClient.TrainingResult result = ocrClient.segtrainModel(zipBytes);

            return Objects.requireNonNull(txTemplate.execute(status -> {
                run.setStatus(TrainingStatus.DONE);
                run.setCompletedAt(Instant.now());
                run.setCer(result.cer());
                run.setLoss(result.loss());
                run.setAccuracy(result.accuracy());
                run.setEpochs(result.epochs());
                OcrTrainingRun updated = trainingRunRepository.save(run);
                log.info("[trainingRun={}] Segmentation training completed — cer={} epochs={}", runId, result.cer(), result.epochs());
                return updated;
            }));
        } catch (Exception e) {
            return Objects.requireNonNull(txTemplate.execute(status -> {
                run.setStatus(TrainingStatus.FAILED);
                run.setErrorMessage(e.getMessage());
                run.setCompletedAt(Instant.now());
                OcrTrainingRun failed = trainingRunRepository.save(run);
                log.error("[trainingRun={}] Segmentation training failed: {}", runId, e.getMessage(), e);
                return failed;
            }));
        } finally {
            MDC.remove("trainingRunId");
        }
    }

    public TrainingInfoResponse getTrainingInfo() {
        var eligibleBlocks = trainingDataExportService.queryEligibleBlocks();
        int availableDocuments = (int) eligibleBlocks.stream()
                .map(b -> b.getDocumentId())
                .distinct()
                .count();

        int totalOcrBlocks = (int) transcriptionBlockQueryService.count();
        int availableSegBlocks = segmentationTrainingExportService.querySegmentationBlocks().size();

        List<OcrTrainingRun> recentRuns = trainingRunRepository.findTop20ByOrderByCreatedAtDesc();
        OcrTrainingRun lastRun = recentRuns.isEmpty() ? null : recentRuns.get(0);

        List<SenderModel> senderModels = senderModelService.getAllSenderModels();

        List<UUID> allPersonIds = senderModels.stream()
                .map(SenderModel::getPersonId)
                .distinct()
                .toList();
        Map<String, String> personNames = new HashMap<>();
        if (!allPersonIds.isEmpty()) {
            personService.getAllById(allPersonIds)
                    .forEach(p -> personNames.put(p.getId().toString(), p.getDisplayName()));
        }

        return new TrainingInfoResponse(
                eligibleBlocks.size(),
                totalOcrBlocks,
                availableDocuments,
                availableSegBlocks,
                ocrHealthClient.isHealthy(),
                lastRun,
                recentRuns,
                personNames,
                senderModels
        );
    }

    public TrainingHistoryResponse getGlobalTrainingHistory() {
        List<OcrTrainingRun> runs = trainingRunRepository.findByPersonIdIsNullOrderByCreatedAtDesc();
        return new TrainingHistoryResponse(runs, Map.of());
    }

    public TrainingHistoryResponse getSenderTrainingHistory(UUID personId) {
        String personName = personService.getById(personId).getDisplayName();
        List<OcrTrainingRun> runs = trainingRunRepository.findByPersonIdOrderByCreatedAtDesc(personId);
        return new TrainingHistoryResponse(runs, Map.of(personId.toString(), personName));
    }

    @EventListener(ApplicationReadyEvent.class)
    @Transactional
    public void recoverOrphanedRuns() {
        var cutoff = Instant.now().minusSeconds(3600);
        trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING).ifPresent(run -> {
            if (run.getCreatedAt().isBefore(cutoff)) {
                run.setStatus(TrainingStatus.FAILED);
                run.setErrorMessage("Abgebrochen: Dienst wurde neugestartet");
                run.setCompletedAt(Instant.now());
                trainingRunRepository.save(run);
                log.warn("Recovered orphaned training run {} (marked FAILED on startup)", run.getId());
            }
        });
    }

}
