package org.raddatz.familienarchiv.ocr;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
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.ocr.SenderModelRepository;
import org.raddatz.familienarchiv.person.PersonService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionBlockQueryService;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.scheduling.annotation.Async;
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.List;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;

@Service
@RequiredArgsConstructor
@Slf4j
public class SenderModelService {

    private final SenderModelRepository senderModelRepository;
    private final TranscriptionBlockQueryService transcriptionBlockQueryService;
    private final OcrTrainingRunRepository trainingRunRepository;
    private final OcrClient ocrClient;
    private final TransactionTemplate txTemplate;
    private final TrainingDataExportService trainingDataExportService;
    private final PersonService personService;

    // Self-reference through the Spring proxy so @Async is honoured on self-calls.
    @Lazy
    @Autowired
    private SenderModelService self;

    @Value("${ocr.sender-model.activation-threshold:100}")
    private int activationThreshold;

    @Value("${ocr.sender-model.retrain-delta:50}")
    private int retrainDelta;

    /** Returns the model path if a trained sender model exists for this person. */
    public Optional<String> maybeGetModelPath(UUID personId) {
        return senderModelRepository.findByPersonId(personId)
                .map(SenderModel::getModelPath);
    }

    public List<SenderModel> getAllSenderModels() {
        return senderModelRepository.findAll();
    }

    public OcrTrainingRun triggerManualSenderTraining(UUID personId) {
        personService.getById(personId);
        long correctedLines = transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId);
        boolean runNow = runOrQueueSenderTraining(personId, (int) correctedLines);
        TrainingStatus targetStatus = runNow ? TrainingStatus.RUNNING : TrainingStatus.QUEUED;
        OcrTrainingRun run = trainingRunRepository.findFirstByPersonIdAndStatus(personId, targetStatus)
                .orElseThrow(() -> DomainException.internal(
                        ErrorCode.OCR_TRAINING_CONFLICT,
                        "Expected " + targetStatus + " run for person " + personId));
        if (runNow) {
            self.runSenderTraining(personId);
        }
        return run;
    }

    @Async
    public void runSenderTraining(UUID personId) {
        long correctedLines = transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId);
        triggerSenderTraining(personId, (int) correctedLines);
    }

    /**
     * Called after every MANUAL block save for HANDWRITING_KURRENT documents.
     * Checks activation and retrain thresholds; enqueues or starts sender training when met.
     */
    @Async
    public void checkAndTriggerTraining(UUID personId) {
        long correctedLines = transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId);
        Optional<SenderModel> existing = senderModelRepository.findByPersonId(personId);

        boolean shouldActivate = existing.isEmpty() && correctedLines >= activationThreshold;
        boolean shouldRetrain = existing.isPresent()
                && (correctedLines - existing.get().getCorrectedLinesAtTraining()) >= retrainDelta;

        if (!shouldActivate && !shouldRetrain) {
            return;
        }

        log.info("Sender training threshold met for person {} (correctedLines={}, activate={}, retrain={})",
                personId, correctedLines, shouldActivate, shouldRetrain);

        boolean runNow = runOrQueueSenderTraining(personId, (int) correctedLines);
        if (runNow) {
            triggerSenderTraining(personId, (int) correctedLines);
        }
    }

    /**
     * Atomically checks the queue state and either creates a RUNNING row (returns true) or a
     * QUEUED row (returns false). All three operations — idle check, duplicate-queue guard, and
     * RUNNING row creation — happen in one transaction, eliminating the race window that would
     * otherwise exist between the check and a separate RUNNING row creation.
     */
    @Transactional
    public boolean runOrQueueSenderTraining(UUID personId, int correctedLines) {
        if (trainingRunRepository.existsByPersonIdAndStatus(personId, TrainingStatus.QUEUED)) {
            log.info("Sender training already queued for person {} — skipping duplicate trigger", personId);
            return false;
        }

        if (trainingRunRepository.findFirstByStatus(TrainingStatus.RUNNING).isPresent()) {
            int blockCount = (int) transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId);
            trainingRunRepository.save(OcrTrainingRun.builder()
                    .status(TrainingStatus.QUEUED)
                    .personId(personId)
                    .blockCount(blockCount)
                    .documentCount(0)
                    .modelName("sender_" + personId)
                    .build());
            log.info("Queued sender training for person {} — training already running", personId);
            return false;
        }

        long blockCount = transcriptionBlockQueryService.countManualKurrentBlocksByPerson(personId);
        trainingRunRepository.save(OcrTrainingRun.builder()
                .status(TrainingStatus.RUNNING)
                .personId(personId)
                .blockCount((int) blockCount)
                .documentCount(0)
                .modelName("sender_" + personId)
                .build());
        return true;
    }

    /**
     * Executes sender training synchronously. Caller must run this on a background thread.
     * The RUNNING row is expected to already exist — created atomically by
     * runOrQueueSenderTraining (for new runs) or by promoteNextQueuedRun (for promoted runs).
     */
    public void triggerSenderTraining(UUID personId, int correctedLines) {
        String outputModelPath = "/app/models/sender_" + personId + ".mlmodel";

        OcrTrainingRun run = Objects.requireNonNull(txTemplate.execute(status ->
                trainingRunRepository.findFirstByPersonIdAndStatus(personId, TrainingStatus.RUNNING)
                        .orElseThrow(() -> DomainException.internal(
                                ErrorCode.OCR_TRAINING_CONFLICT,
                                "Expected RUNNING row for person " + personId + " but none found"))));

        String runId = run.getId().toString();
        MDC.put("trainingRunId", runId);
        log.info("Started sender training run {} for person {}", runId, personId);

        try {
            byte[] zipBytes = exportSenderData(personId);
            log.info("[trainingRun={}] Sending {} bytes to OCR service for sender training", runId, zipBytes.length);
            OcrClient.TrainingResult result = ocrClient.trainSenderModel(zipBytes, outputModelPath);

            txTemplate.execute(status -> {
                SenderModel model = senderModelRepository.findByPersonId(personId)
                        .orElseGet(() -> SenderModel.builder().personId(personId).build());
                model.setModelPath(outputModelPath);
                model.setCer(result.cer());
                model.setAccuracy(result.accuracy());
                model.setCorrectedLinesAtTraining(correctedLines);
                senderModelRepository.save(model);

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

    private byte[] exportSenderData(UUID personId) throws java.io.IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        trainingDataExportService.exportForSender(personId).writeTo(baos);
        return baos.toByteArray();
    }

    /**
     * Promotes the oldest QUEUED sender run to RUNNING and triggers its training.
     * Called in the finally block of triggerSenderTraining, creating a sequential chain:
     * each run promotes the next only after it fully completes (success or failure).
     *
     * This is intentionally tail-recursive via the @Async thread: the same thread holds the
     * full queue drain, serialising all sender training runs naturally without an external
     * scheduler. With N queued runs the thread stays occupied for N sequential trainings —
     * acceptable because the @Async executor is dedicated to long-running background work.
     */
    private void promoteNextQueuedRun() {
        Optional<OcrTrainingRun> queuedOpt = txTemplate.execute(status ->
                trainingRunRepository.findFirstByStatusOrderByCreatedAtAsc(TrainingStatus.QUEUED)
                        .map(queued -> {
                            queued.setStatus(TrainingStatus.RUNNING);
                            return trainingRunRepository.save(queued);
                        }));

        if (queuedOpt != null && queuedOpt.isPresent()) {
            OcrTrainingRun promoted = queuedOpt.get();
            log.info("Promoting queued sender training run {} for person {}", promoted.getId(), promoted.getPersonId());
            long freshCount = transcriptionBlockQueryService.countManualKurrentBlocksByPerson(promoted.getPersonId());
            triggerSenderTraining(promoted.getPersonId(), (int) freshCount);
        }
    }
}
