# Copyright (C) 2026 NPO NT
# SPDX-License-Identifier: AGPL-3.0-only

"""Evaluate synthetic registry values against local OCR without exporting them."""

from __future__ import annotations

import argparse
import json
import re
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any

from mika_deid.detectors import _box_for_span, _build_line_map, _line_groups, detect_tokens
from mika_deid.loader import load_document
from mika_deid.models import Box, Finding, WordToken
from mika_deid.ocr import TesseractOcr


NEUTRAL_DOCUMENT_ID = re.compile(r"document-\d{4}", re.IGNORECASE)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Check synthetic registry names and birth dates against OCR output"
    )
    parser.add_argument("documents_dir", type=Path)
    parser.add_argument("private_truth_json", type=Path)
    parser.add_argument("output_json", type=Path)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--max-pages", type=int, default=0, help="0 means all pages")
    parser.add_argument(
        "--confirmed-synthetic",
        action="store_true",
        help="Required acknowledgement that all input values and documents are synthetic",
    )
    return parser.parse_args()


def normalize_words(value: str) -> list[str]:
    normalized = value.upper().replace("Ё", "Е")
    return re.findall(r"[A-ZА-Я0-9]+", normalized)


def date_digit_variants(iso_date: str) -> set[str]:
    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", iso_date):
        raise ValueError("birth_date must use YYYY-MM-DD")
    year, month, day = iso_date.split("-")
    variants = {
        f"{day}{month}{year}",
        f"{int(day)}{int(month)}{year}",
        f"{year}{month}{day}",
        f"{year}{int(month)}{int(day)}",
    }
    return variants


def _span_box_for_words(
    line_text: str,
    mapping: list[tuple[int, int, WordToken]],
    expected_words: list[str],
) -> Box | None:
    normalized_tokens = [normalize_words(token.text) for _, _, token in mapping]
    flattened = [parts[0] if len(parts) == 1 else "".join(parts) for parts in normalized_tokens]
    expected = [word for word in expected_words if word]
    if not expected:
        return None
    for start_index in range(len(flattened)):
        if flattened[start_index] != expected[0]:
            continue
        cursor = start_index
        matched_tokens: list[WordToken] = []
        for word in expected:
            while cursor < len(flattened) and not flattened[cursor]:
                cursor += 1
            if cursor >= len(flattened) or flattened[cursor] != word:
                break
            matched_tokens.append(mapping[cursor][2])
            cursor += 1
        else:
            box = matched_tokens[0].box
            for token in matched_tokens[1:]:
                box = box.union(token.box)
            return box
    compact_line = "".join(normalize_words(line_text))
    compact_expected = "".join(expected)
    if compact_expected not in compact_line:
        return None
    # OCR may fuse multiple name parts into one token. Use the whole matching line
    # as a conservative region for the classification-overlap check.
    return _box_for_span(mapping, 0, len(line_text))


def _date_box(
    line_text: str,
    mapping: list[tuple[int, int, WordToken]],
    variants: set[str],
) -> Box | None:
    for start_index in range(len(mapping)):
        for end_index in range(start_index + 1, min(len(mapping), start_index + 5) + 1):
            tokens = [mapping[index][2] for index in range(start_index, end_index)]
            digits = "".join(re.findall(r"\d", " ".join(token.text for token in tokens)))
            if digits not in variants:
                continue
            box = tokens[0].box
            for token in tokens[1:]:
                box = box.union(token.box)
            return box
    line_digits = "".join(re.findall(r"\d", line_text))
    if line_digits in variants:
        return _box_for_span(mapping, 0, len(line_text))
    return None


def _overlaps_category(findings: list[Finding], category: str, expected: Box) -> bool:
    return any(
        finding.category == category and finding.box.overlap_ratio(expected) >= 0.5
        for finding in findings
    )


def _evaluate_one(
    source_path: Path,
    truth: dict[str, str],
    max_pages: int,
) -> dict[str, Any]:
    started = time.monotonic()
    document = load_document(source_path)
    ocr = TesseractOcr()
    name_words = normalize_words(truth["person_name"])
    birth_variants = date_digit_variants(truth["birth_date"])
    name_recognized = False
    birth_date_recognized = False
    name_classified = False
    birth_date_classified = False
    birth_date_as_event = False
    pages_scanned = 0

    for page in document.pages[: max_pages or None]:
        result = ocr.recognize(page.image, page.index)
        pages_scanned += 1
        findings = detect_tokens(result.tokens, page.index)
        for line_tokens in _line_groups(result.tokens):
            line_text, mapping = _build_line_map(line_tokens)
            if not name_recognized or not name_classified:
                name_box = _span_box_for_words(line_text, mapping, name_words)
                if name_box is not None:
                    name_recognized = True
                    name_classified = name_classified or _overlaps_category(
                        findings, "person_name", name_box
                    )
            if not birth_date_recognized or not birth_date_classified:
                birth_box = _date_box(line_text, mapping, birth_variants)
                if birth_box is not None:
                    birth_date_recognized = True
                    birth_date_classified = birth_date_classified or _overlaps_category(
                        findings, "birth_date", birth_box
                    )
                    birth_date_as_event = birth_date_as_event or _overlaps_category(
                        findings, "event_date", birth_box
                    )
        # Once both exact registry values have been recognized, later duplicate
        # occurrences do not add value-level evidence and would only repeat OCR.
        if name_recognized and birth_date_recognized:
            break

    return {
        "document_id": source_path.stem,
        "status": "ok",
        "page_count": len(document.pages),
        "pages_scanned": pages_scanned,
        "name_value_recognized": name_recognized,
        "birth_date_value_recognized": birth_date_recognized,
        "name_classified_person_name": name_classified,
        "birth_date_classified_birth_date": birth_date_classified,
        "birth_date_classified_event_date": birth_date_as_event,
        "duration_seconds": round(time.monotonic() - started, 3),
    }


def _safe_error(document_id: str, exc: BaseException) -> dict[str, Any]:
    return {
        "document_id": document_id,
        "status": "error",
        "error_type": type(exc).__name__,
        "error_message_exported": False,
    }


def main() -> int:
    args = parse_args()
    if not args.confirmed_synthetic:
        raise SystemExit("Refusing value-level evaluation without --confirmed-synthetic")
    if args.workers < 1 or args.workers > 4:
        raise SystemExit("--workers must be between 1 and 4")
    if args.max_pages < 0:
        raise SystemExit("--max-pages cannot be negative")

    documents_dir = args.documents_dir.expanduser().resolve()
    truth_path = args.private_truth_json.expanduser().resolve()
    output_path = args.output_json.expanduser().resolve()
    if output_path.exists():
        raise SystemExit("Output JSON already exists")
    truth_payload = json.loads(truth_path.read_text(encoding="utf-8"))
    if truth_payload.get("corpus_classification") != "CONFIRMED_FULLY_SYNTHETIC":
        raise SystemExit("Private truth JSON is not marked fully synthetic")
    truth_records = truth_payload.get("records") or []
    if not truth_records:
        raise SystemExit("Private truth JSON contains no records")

    started = time.monotonic()
    records: list[dict[str, Any]] = []
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        futures = {}
        for truth in truth_records:
            document_id = str(truth.get("document_id") or "")
            if not NEUTRAL_DOCUMENT_ID.fullmatch(document_id):
                raise SystemExit("Private truth JSON contains a non-neutral document_id")
            source_path = documents_dir / f"{document_id}.pdf"
            if not source_path.is_file():
                raise SystemExit(f"Missing source PDF for {document_id}")
            futures[executor.submit(_evaluate_one, source_path, truth, args.max_pages)] = document_id
        completed = 0
        for future in as_completed(futures):
            document_id = futures[future]
            try:
                record = future.result()
            except BaseException as exc:
                record = _safe_error(document_id, exc)
            records.append(record)
            completed += 1
            print(
                json.dumps(
                    {
                        "completed": completed,
                        "total": len(futures),
                        "document_id": document_id,
                        "status": record["status"],
                        "pages_scanned": record.get("pages_scanned"),
                    },
                    ensure_ascii=True,
                ),
                flush=True,
            )

    records.sort(key=lambda record: record["document_id"])
    successful = [record for record in records if record["status"] == "ok"]
    summary = {
        "schema": "mika-deid-registry-value-evaluation-v1",
        "corpus_classification": "CONFIRMED_FULLY_SYNTHETIC",
        "document_count": len(records),
        "successful_documents": len(successful),
        "failed_documents": len(records) - len(successful),
        "name_value_recognized_documents": sum(
            record["name_value_recognized"] for record in successful
        ),
        "birth_date_value_recognized_documents": sum(
            record["birth_date_value_recognized"] for record in successful
        ),
        "name_classified_person_name_documents": sum(
            record["name_classified_person_name"] for record in successful
        ),
        "birth_date_classified_birth_date_documents": sum(
            record["birth_date_classified_birth_date"] for record in successful
        ),
        "birth_date_classified_event_date_documents": sum(
            record["birth_date_classified_event_date"] for record in successful
        ),
        "pages_scanned": sum(record["pages_scanned"] for record in successful),
        "elapsed_seconds": round(time.monotonic() - started, 3),
        "metric_boundary": (
            "Value-level OCR/classification coverage against a synthetic registry; "
            "not coordinate-level recall or precision and not proof that every PDF contains every registry value."
        ),
        "privacy": {
            "registry_values_exported": False,
            "ocr_text_exported": False,
            "original_filenames_exported": False,
        },
        "records": records,
    }
    output_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    output_path.chmod(0o600)
    print(json.dumps({key: value for key, value in summary.items() if key != "records"}, ensure_ascii=True))
    return 0 if summary["failed_documents"] == 0 else 2


if __name__ == "__main__":
    raise SystemExit(main())
