#!/bin/bash
set -euo pipefail

# Downloads Kraken HTR models, runs OCR evaluation on sample PDFs, and activates the winner.
#
# Usage:
#   ./scripts/download-kraken-models.sh                # download models + run evaluation
#   ./scripts/download-kraken-models.sh --activate 1   # activate model 1 (german_kurrent_manu_9)
#   ./scripts/download-kraken-models.sh --activate 2   # activate model 2 (kurrent-de)
#   ./scripts/download-kraken-models.sh --eval-only    # re-run evaluation (models already downloaded)

COMPOSE_SERVICE="ocr-service"
MODEL_DIR="/app/models"
ACTIVE_MODEL="$MODEL_DIR/german_kurrent.mlmodel"

MODEL_1_NAME="german_kurrent_manu_9"
MODEL_1_DESC="19th-century German administrative Kurrent (HTR-United)"
MODEL_1_PATH="$MODEL_DIR/$MODEL_1_NAME.mlmodel"

MODEL_2_NAME="kurrent-de"
MODEL_2_DESC="Broad German Kurrent coverage (HTR-United)"
MODEL_2_PATH="$MODEL_DIR/$MODEL_2_NAME.mlmodel"

EVAL_PDFS=("Eu-0693.pdf" "Eu-0692.pdf" "W-0150.pdf" "W-0575.pdf")
IMPORT_DIR="./import"
OUTPUT_BASE="./ocr-model-evaluation"

# ─── Functions ────────────────────────────────────────────────────────────────

download_models() {
    echo "Downloading Kraken HTR models into the ocr_models volume..."
    echo ""

    echo "Model 1: $MODEL_1_NAME"
    echo "  $MODEL_1_DESC"
    docker compose run --rm "$COMPOSE_SERVICE" kraken get "$MODEL_1_NAME" -o "$MODEL_1_PATH"
    echo ""

    echo "Model 2: $MODEL_2_NAME"
    echo "  $MODEL_2_DESC"
    docker compose run --rm "$COMPOSE_SERVICE" kraken get "$MODEL_2_NAME" -o "$MODEL_2_PATH"
    echo ""

    echo "Both models downloaded."
}

run_evaluation() {
    echo ""
    echo "═══════════════════════════════════════════════════════"
    echo "  Running OCR evaluation on ${#EVAL_PDFS[@]} documents"
    echo "═══════════════════════════════════════════════════════"
    echo ""

    # Create output directories on the host
    local dir_1="$OUTPUT_BASE/$MODEL_1_NAME"
    local dir_2="$OUTPUT_BASE/$MODEL_2_NAME"
    mkdir -p "$dir_1" "$dir_2"

    for pdf in "${EVAL_PDFS[@]}"; do
        local src="$IMPORT_DIR/$pdf"
        local basename="${pdf%.pdf}"

        if [[ ! -f "$src" ]]; then
            echo "SKIP: $src not found"
            continue
        fi

        echo "──── $pdf ────"

        # Model 1
        echo "  Model 1: $MODEL_1_NAME ..."
        docker compose run --rm \
            -v "$(cd "$IMPORT_DIR" && pwd):/eval-input:ro" \
            -v "$(cd "$dir_1" && pwd):/eval-output" \
            "$COMPOSE_SERVICE" \
            kraken -i "/eval-input/$pdf" "/eval-output/${basename}.txt" segment -bl ocr -m "$MODEL_1_PATH" \
            2>/dev/null || echo "  ⚠ Model 1 failed on $pdf"

        # Model 2
        echo "  Model 2: $MODEL_2_NAME ..."
        docker compose run --rm \
            -v "$(cd "$IMPORT_DIR" && pwd):/eval-input:ro" \
            -v "$(cd "$dir_2" && pwd):/eval-output" \
            "$COMPOSE_SERVICE" \
            kraken -i "/eval-input/$pdf" "/eval-output/${basename}.txt" segment -bl ocr -m "$MODEL_2_PATH" \
            2>/dev/null || echo "  ⚠ Model 2 failed on $pdf"

        echo ""
    done

    echo "═══════════════════════════════════════════════════════"
    echo "  Evaluation complete. Results:"
    echo ""
    echo "  Model 1 ($MODEL_1_NAME):"
    for f in "$dir_1"/*.txt; do
        [[ -f "$f" ]] && echo "    $(basename "$f"): $(wc -l < "$f") lines, $(wc -c < "$f") bytes"
    done
    echo ""
    echo "  Model 2 ($MODEL_2_NAME):"
    for f in "$dir_2"/*.txt; do
        [[ -f "$f" ]] && echo "    $(basename "$f"): $(wc -l < "$f") lines, $(wc -c < "$f") bytes"
    done
    echo ""
    echo "  Compare the outputs side by side:"
    echo "    diff $dir_1/Eu-0693.txt $dir_2/Eu-0693.txt"
    echo ""
    echo "  Then activate the better model:"
    echo "    ./scripts/download-kraken-models.sh --activate 1   # $MODEL_1_NAME"
    echo "    ./scripts/download-kraken-models.sh --activate 2   # $MODEL_2_NAME"
    echo "═══════════════════════════════════════════════════════"
}

activate_model() {
    local choice="$1"
    case "$choice" in
        1)
            echo "Activating model 1: $MODEL_1_NAME"
            docker compose run --rm "$COMPOSE_SERVICE" cp "$MODEL_1_PATH" "$ACTIVE_MODEL"
            ;;
        2)
            echo "Activating model 2: $MODEL_2_NAME"
            docker compose run --rm "$COMPOSE_SERVICE" cp "$MODEL_2_PATH" "$ACTIVE_MODEL"
            ;;
        *)
            echo "Error: --activate expects 1 or 2"
            exit 1
            ;;
    esac

    echo "Active model is now: $ACTIVE_MODEL"
    echo "Restart the OCR service to load the new model:"
    echo "  docker compose restart ocr-service"
}

# ─── Main ─────────────────────────────────────────────────────────────────────

case "${1:-}" in
    --activate)
        activate_model "${2:-}"
        ;;
    --eval-only)
        run_evaluation
        ;;
    *)
        download_models
        run_evaluation
        ;;
esac