package org.raddatz.familienarchiv.ocr;

import org.raddatz.familienarchiv.ocr.ScriptType;
import org.springframework.lang.Nullable;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.function.Consumer;

public interface OcrClient {
    List<OcrBlockResult> extractBlocks(String pdfUrl, ScriptType scriptType);

    /**
     * A pre-drawn annotation region to use as guidance for OCR.
     * When regions are provided, the OCR engine crops to each region and
     * runs recognition only within that area, skipping full-page layout detection.
     */
    record OcrRegion(String annotationId, int pageNumber,
                     double x, double y, double width, double height) {}

    /**
     * Send a training ZIP to the OCR service for fine-tuning the Kurrent model.
     *
     * @param trainingDataZip raw ZIP bytes produced by TrainingDataExportService
     * @return training result metrics (loss, accuracy, epochs)
     */
    TrainingResult trainModel(byte[] trainingDataZip);

    record TrainingResult(Double loss, Double accuracy, Double cer, Integer epochs) {}

    /**
     * Send a segmentation training ZIP to the OCR service for fine-tuning the blla model.
     *
     * @param trainingDataZip raw ZIP bytes produced by SegmentationTrainingExportService
     * @return training result metrics
     */
    TrainingResult segtrainModel(byte[] trainingDataZip);

    /**
     * Fine-tune the Kurrent model for a specific sender.
     *
     * @param trainingDataZip raw ZIP bytes produced by TrainingDataExportService.exportForSender()
     * @param outputModelPath where to save the trained model (e.g. /app/models/sender_{uuid}.mlmodel)
     * @return training result metrics
     */
    TrainingResult trainSenderModel(byte[] trainingDataZip, String outputModelPath);

    /**
     * Stream OCR results page-by-page via NDJSON, optionally using a sender-specific model.
     * The default implementation synthesizes events from extractBlocks() for backward compatibility.
     * Implementations that support real streaming (e.g. RestClientOcrClient) override this.
     *
     * @param regions         optional list of pre-drawn annotation regions; when non-null,
     *                        the OCR service runs in guided mode (crop + recognize per region)
     * @param senderModelPath optional path to a per-sender model file; null means use base model
     */
    default void streamBlocks(String pdfUrl, ScriptType scriptType,
                               List<OcrRegion> regions,
                               @Nullable String senderModelPath,
                               Consumer<OcrStreamEvent> handler) {
        List<OcrBlockResult> allBlocks = extractBlocks(pdfUrl, scriptType);

        LinkedHashMap<Integer, List<OcrBlockResult>> byPage = new LinkedHashMap<>();
        for (OcrBlockResult block : allBlocks) {
            byPage.computeIfAbsent(block.pageNumber(), k -> new ArrayList<>()).add(block);
        }

        int totalPages = byPage.isEmpty() ? 0 : byPage.keySet().stream().mapToInt(i -> i).max().orElse(0) + 1;
        handler.accept(new OcrStreamEvent.Start(totalPages));

        for (var entry : byPage.entrySet()) {
            handler.accept(new OcrStreamEvent.Page(entry.getKey(), entry.getValue()));
        }

        handler.accept(new OcrStreamEvent.Done(allBlocks.size(), 0));
    }

    default void streamBlocks(String pdfUrl, ScriptType scriptType,
                               List<OcrRegion> regions, Consumer<OcrStreamEvent> handler) {
        streamBlocks(pdfUrl, scriptType, regions, null, handler);
    }
}
