package org.raddatz.familienarchiv.ocr;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.raddatz.familienarchiv.document.annotation.CreateAnnotationDTO;
import org.raddatz.familienarchiv.document.annotation.DocumentAnnotation;
import org.raddatz.familienarchiv.ocr.OcrJobDocumentRepository;
import org.raddatz.familienarchiv.ocr.OcrJobRepository;
import org.raddatz.familienarchiv.document.Document;
import org.raddatz.familienarchiv.document.DocumentStatus;
import org.raddatz.familienarchiv.document.annotation.AnnotationService;
import org.raddatz.familienarchiv.document.DocumentService;
import org.raddatz.familienarchiv.document.transcription.TranscriptionService;
import org.raddatz.familienarchiv.filestorage.FileService;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

@Component
@RequiredArgsConstructor
@Slf4j
public class OcrAsyncRunner {

    private static final String OCR_ANNOTATION_COLOR = "#00C7B1";

    private final OcrClient ocrClient;
    private final DocumentService documentService;
    private final TranscriptionService transcriptionService;
    private final AnnotationService annotationService;
    private final FileService fileService;
    private final OcrJobRepository ocrJobRepository;
    private final OcrJobDocumentRepository ocrJobDocumentRepository;
    private final OcrProgressService ocrProgressService;
    private final SenderModelService senderModelService;

    @Async
    public void runSingleDocument(UUID jobId, UUID documentId, UUID userId) {
        runSingleDocument(jobId, documentId, userId, false);
    }

    @Async
    public void runSingleDocument(UUID jobId, UUID documentId, UUID userId, boolean useExistingAnnotations) {
        OcrJob job = ocrJobRepository.findById(jobId).orElse(null);
        if (job == null) return;

        job.setStatus(OcrJobStatus.RUNNING);
        updateProgress(job, "PREPARING");

        OcrJobDocument jobDoc = ocrJobDocumentRepository.findByJobIdAndDocumentId(jobId, documentId)
                .orElse(null);
        if (jobDoc != null) {
            jobDoc.setStatus(OcrDocumentStatus.RUNNING);
            ocrJobDocumentRepository.save(jobDoc);
        }

        Document doc = documentService.getDocumentById(documentId);

        try {
            updateProgress(job, "LOADING");

            List<OcrClient.OcrRegion> regions = null;
            if (useExistingAnnotations) {
                regions = annotationService.listAnnotations(documentId).stream()
                        .map(a -> new OcrClient.OcrRegion(
                                a.getId().toString(), a.getPageNumber(),
                                a.getX(), a.getY(), a.getWidth(), a.getHeight()))
                        .toList();
            } else {
                clearExistingBlocks(documentId);
            }

            String pdfUrl = fileService.generatePresignedUrl(doc.getFilePath());

            String senderModelPath = null;
            if (doc.getSender() != null && doc.getScriptType() == ScriptType.HANDWRITING_KURRENT) {
                senderModelPath = senderModelService.maybeGetModelPath(doc.getSender().getId()).orElse(null);
            }

            AtomicInteger blockCounter = new AtomicInteger(0);
            AtomicInteger currentPage = new AtomicInteger(0);
            AtomicInteger skippedPages = new AtomicInteger(0);
            AtomicInteger totalPages = new AtomicInteger(0);

            final String finalSenderModelPath = senderModelPath;
            ocrClient.streamBlocks(pdfUrl, doc.getScriptType(), regions, finalSenderModelPath, event -> {
                switch (event) {
                    case OcrStreamEvent.Start start -> {
                        totalPages.set(start.totalPages());
                        if (jobDoc != null) {
                            jobDoc.setTotalPages(start.totalPages());
                            ocrJobDocumentRepository.save(jobDoc);
                        }
                    }
                    case OcrStreamEvent.Preprocessing preprocessing -> {
                        updateProgress(job, "PREPROCESSING_PAGE:" + preprocessing.pageNumber()
                                + ":" + totalPages.get());
                    }
                    case OcrStreamEvent.Page page -> {
                        for (OcrBlockResult block : page.blocks()) {
                            createSingleBlock(documentId, block, userId,
                                    doc.getFileHash(), blockCounter.getAndIncrement());
                        }
                        currentPage.incrementAndGet();
                        if (jobDoc != null) {
                            jobDoc.setCurrentPage(currentPage.get());
                            ocrJobDocumentRepository.save(jobDoc);
                        }
                        updateProgress(job, "ANALYZING_PAGE:" + currentPage.get()
                                + ":" + totalPages.get() + ":" + blockCounter.get());
                    }
                    case OcrStreamEvent.Error error -> {
                        log.warn("OCR page {} failed for document {}: {}",
                                error.pageNumber(), documentId, error.message());
                        skippedPages.incrementAndGet();
                        currentPage.incrementAndGet();
                        if (jobDoc != null) {
                            jobDoc.setCurrentPage(currentPage.get());
                            ocrJobDocumentRepository.save(jobDoc);
                        }
                    }
                    case OcrStreamEvent.Done done -> {
                        if (jobDoc != null) {
                            jobDoc.setCurrentPage(totalPages.get());
                            ocrJobDocumentRepository.save(jobDoc);
                        }
                    }
                }
            });

            job.setStatus(OcrJobStatus.DONE);
            job.setProcessedDocuments(1);
            updateProgress(job, "DONE:" + blockCounter.get() + ":" + skippedPages.get());
            if (jobDoc != null) {
                jobDoc.setStatus(OcrDocumentStatus.DONE);
                ocrJobDocumentRepository.save(jobDoc);
            }
        } catch (Exception e) {
            log.error("OCR processing failed for document {}", documentId, e);
            job.setStatus(OcrJobStatus.FAILED);
            job.setErrorCount(1);
            updateProgress(job, "ERROR");
            if (jobDoc != null) {
                jobDoc.setStatus(OcrDocumentStatus.FAILED);
                jobDoc.setErrorMessage(e.getMessage());
                ocrJobDocumentRepository.save(jobDoc);
            }
        }
    }

    private void updateProgress(OcrJob job, String message) {
        job.setProgressMessage(message);
        ocrJobRepository.save(job);
    }

    @Async
    public void runBatch(UUID jobId, UUID userId) {
        OcrJob job = ocrJobRepository.findById(jobId).orElse(null);
        if (job == null) return;

        job.setStatus(OcrJobStatus.RUNNING);
        ocrJobRepository.save(job);

        List<OcrJobDocument> jobDocs = ocrJobDocumentRepository.findByJobIdOrderByCreatedAtAsc(jobId);

        for (OcrJobDocument jobDoc : jobDocs) {
            Document doc = documentService.getDocumentById(jobDoc.getDocumentId());

            if (doc.getStatus() == DocumentStatus.PLACEHOLDER) {
                jobDoc.setStatus(OcrDocumentStatus.SKIPPED);
                ocrJobDocumentRepository.save(jobDoc);
                job.setSkippedCount(job.getSkippedCount() + 1);
                ocrJobRepository.save(job);
                ocrProgressService.emit(jobId, "document", Map.of(
                        "documentId", jobDoc.getDocumentId(),
                        "status", "SKIPPED",
                        "processed", job.getProcessedDocuments(),
                        "total", job.getTotalDocuments()));
                continue;
            }

            jobDoc.setStatus(OcrDocumentStatus.RUNNING);
            ocrJobDocumentRepository.save(jobDoc);

            try {
                processDocument(jobDoc.getDocumentId(), doc, userId);
                jobDoc.setStatus(OcrDocumentStatus.DONE);
                job.setProcessedDocuments(job.getProcessedDocuments() + 1);
            } catch (Exception e) {
                log.error("OCR batch: failed document {}", jobDoc.getDocumentId(), e);
                jobDoc.setStatus(OcrDocumentStatus.FAILED);
                jobDoc.setErrorMessage(e.getMessage());
                job.setErrorCount(job.getErrorCount() + 1);
            }

            ocrJobDocumentRepository.save(jobDoc);
            ocrJobRepository.save(job);

            ocrProgressService.emit(jobId, "document", Map.of(
                    "documentId", jobDoc.getDocumentId(),
                    "status", jobDoc.getStatus().name(),
                    "processed", job.getProcessedDocuments(),
                    "total", job.getTotalDocuments()));
        }

        job.setStatus(OcrJobStatus.DONE);
        ocrJobRepository.save(job);

        ocrProgressService.emit(jobId, "done", Map.of(
                "processed", job.getProcessedDocuments(),
                "errors", job.getErrorCount(),
                "skipped", job.getSkippedCount()));
        ocrProgressService.complete(jobId);
    }

    void processDocument(UUID documentId, Document doc, UUID userId) {
        clearExistingBlocks(documentId);

        String pdfUrl = fileService.generatePresignedUrl(doc.getFilePath());

        String senderModelPath = null;
        if (doc.getSender() != null && doc.getScriptType() == ScriptType.HANDWRITING_KURRENT) {
            senderModelPath = senderModelService.maybeGetModelPath(doc.getSender().getId()).orElse(null);
        }

        final AtomicReference<List<OcrBlockResult>> blocksRef = new AtomicReference<>();
        final String finalSenderModelPath = senderModelPath;
        ocrClient.streamBlocks(pdfUrl, doc.getScriptType(), null, finalSenderModelPath, event -> {
            switch (event) {
                case OcrStreamEvent.Page page -> {
                    blocksRef.compareAndSet(null, new ArrayList<>());
                    blocksRef.get().addAll(page.blocks());
                }
                default -> {}
            }
        });

        List<OcrBlockResult> blocks = blocksRef.get() != null ? blocksRef.get() : List.of();
        createTranscriptionBlocks(documentId, blocks, userId, doc.getFileHash());
    }

    private void clearExistingBlocks(UUID documentId) {
        transcriptionService.deleteAllBlocksByDocument(documentId);
    }

    private void createTranscriptionBlocks(UUID documentId, List<OcrBlockResult> blocks,
                                            UUID userId, String fileHash) {
        for (int i = 0; i < blocks.size(); i++) {
            createSingleBlock(documentId, blocks.get(i), userId, fileHash, i);
        }
    }

    void createSingleBlock(UUID documentId, OcrBlockResult block,
                           UUID userId, String fileHash, int sortOrder) {
        if (block.annotationId() != null) {
            // Guided mode — annotation already exists; upsert the text block only
            transcriptionService.upsertGuidedBlock(
                    documentId, UUID.fromString(block.annotationId()), block.text(), userId);
        } else {
            // Normal mode — create a new annotation and a new OCR block
            CreateAnnotationDTO annotationDTO = new CreateAnnotationDTO(
                    block.pageNumber(), block.x(), block.y(),
                    block.width(), block.height(), OCR_ANNOTATION_COLOR);

            DocumentAnnotation annotation = annotationService.createOcrAnnotation(
                    documentId, annotationDTO, userId, fileHash, block.polygon());

            transcriptionService.createOcrBlock(documentId, annotation.getId(),
                    block.text(), sortOrder, userId);
        }
    }
}
