#!/usr/bin/env python3
"""Resumable Surya detection on template-subtracted handwriting pages."""

from __future__ import annotations

import argparse
import fcntl
import io
import json
import os
import time
import traceback
from pathlib import Path
from typing import Any

import numpy as np
import pyarrow.parquet as pq
from PIL import Image
from surya.detection import DetectionPredictor

from bina02_template_layers import TemplateCatalog


ROWS = 6_669


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--rank", type=int, required=True)
    parser.add_argument("--world-size", type=int, required=True)
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--background-root", type=Path, required=True)
    parser.add_argument("--ocr-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--variant", required=True)
    parser.add_argument("--limit", type=int, default=ROWS)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--progress-every", type=int, default=20)
    return parser.parse_args()


def atomic_json(path: Path, payload: dict[str, Any]) -> None:
    temporary = path.with_suffix(path.suffix + f".{os.getpid()}.tmp")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def selected_indices(limit: int) -> list[int]:
    if not 1 <= limit <= ROWS:
        raise ValueError(f"limit must be in [1, {ROWS}]")
    if limit == ROWS:
        return list(range(ROWS))
    return sorted(
        {
            int(round(value))
            for value in np.linspace(0, ROWS - 1, num=limit)
        }
    )


def load_rows(data: Path, wanted: set[int]) -> dict[int, dict[str, Any]]:
    rows: dict[int, dict[str, Any]] = {}
    offset = 0
    paths = sorted((data / "data").glob("handwriting-*.parquet"))
    if len(paths) != 14:
        raise RuntimeError(f"expected 14 handwriting shards, found {len(paths)}")
    for path in paths:
        count = pq.ParquetFile(path).metadata.num_rows
        local = sorted(index - offset for index in wanted if offset <= index < offset + count)
        if local:
            table = pq.read_table(path, columns=["image", "label"])
            for index in local:
                rows[offset + index] = table.slice(index, 1).to_pylist()[0]
        offset += count
    if offset != ROWS or rows.keys() != wanted:
        raise RuntimeError(
            f"benchmark row contract failed: total={offset}, loaded={len(rows)}, wanted={len(wanted)}"
        )
    return rows


def decode_image(row: dict[str, Any]) -> Image.Image:
    value = row["image"]
    encoded = value.get("bytes") if isinstance(value, dict) else value
    if encoded is None and isinstance(value, dict) and value.get("path"):
        encoded = Path(value["path"]).read_bytes()
    if not encoded:
        raise ValueError("dataset row has no image bytes")
    return Image.open(io.BytesIO(encoded)).convert("RGB")


def line_payload(result: Any) -> list[dict[str, Any]]:
    lines = []
    for item in getattr(result, "bboxes", []):
        bbox = getattr(item, "bbox", None)
        if bbox is None or len(bbox) != 4:
            continue
        polygon = getattr(item, "polygon", None)
        lines.append(
            {
                "bbox": [float(value) for value in bbox],
                "confidence": float(getattr(item, "confidence", 0.0)),
                "polygon": (
                    [[float(value) for value in point] for point in polygon]
                    if polygon is not None
                    else None
                ),
            }
        )
    return lines


def load_done(path: Path) -> set[str]:
    if not path.is_file():
        return set()
    with path.open(encoding="utf-8") as handle:
        return {
            str(row["id"])
            for line in handle
            if line.strip()
            for row in [json.loads(line)]
        }


def merge_if_complete(
    output_dir: Path,
    world_size: int,
    selected: list[int],
    variant: str,
) -> None:
    with (output_dir / "merge.lock").open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        statuses = []
        for rank in range(world_size):
            path = output_dir / f"status-rank{rank:02d}.json"
            if not path.is_file():
                return
            value = json.loads(path.read_text(encoding="utf-8"))
            if value.get("state") != "complete":
                return
            statuses.append(value)
        rows: dict[str, dict[str, Any]] = {}
        duplicates = []
        for rank in range(world_size):
            path = output_dir / f"detections-rank{rank:02d}.jsonl"
            with path.open(encoding="utf-8") as handle:
                for line in handle:
                    row = json.loads(line)
                    row_id = str(row["id"])
                    if row_id in rows:
                        duplicates.append(row_id)
                    rows[row_id] = row
        expected = [f"handwriting:{index}" for index in selected]
        missing = sorted(set(expected) - rows.keys())
        extras = sorted(rows.keys() - set(expected))
        if duplicates or missing or extras:
            raise RuntimeError(
                f"detection merge failed: duplicates={duplicates[:3]} missing={missing[:3]} extras={extras[:3]}"
            )
        temporary = output_dir / "detections.jsonl.tmp"
        with temporary.open("w", encoding="utf-8") as handle:
            for row_id in expected:
                handle.write(json.dumps(rows[row_id], ensure_ascii=False) + "\n")
        os.replace(temporary, output_dir / "detections.jsonl")
        atomic_json(
            output_dir / "run-summary.json",
            {
                "state": "complete",
                "variant": variant,
                "selected_indices": selected,
                "rows": len(rows),
                "failed_pages": sum(row["failures"] for row in statuses),
                "template_match_failures": sum(
                    row["template_match_failures"] for row in statuses
                ),
                "detected_lines": sum(row["detected_lines"] for row in statuses),
                "completed_unix": time.time(),
            },
        )


def main() -> None:
    args = parse_args()
    selected = selected_indices(args.limit)
    assigned = selected[args.rank :: args.world_size]
    output_path = args.output_dir / f"detections-rank{args.rank:02d}.jsonl"
    status_path = args.output_dir / f"status-rank{args.rank:02d}.json"
    args.output_dir.mkdir(parents=True, exist_ok=True)
    done = load_done(output_path)
    wanted = {
        index for index in assigned if f"handwriting:{index}" not in done
    }
    rows = load_rows(args.data, wanted) if wanted else {}
    catalog = TemplateCatalog(args.background_root, args.ocr_root)
    predictor = DetectionPredictor.local(device="cuda")
    started = time.monotonic()
    completed = 0
    failures = 0
    match_failures = 0
    detected_lines = 0
    pending: list[tuple[int, dict[str, Any], Image.Image, dict[str, Any]]] = []

    with output_path.open("a", encoding="utf-8", buffering=1) as output:
        def flush() -> None:
            nonlocal completed, failures, detected_lines
            if not pending:
                return
            images = [item[2] for item in pending]
            try:
                results: list[Any | Exception] = list(predictor(images))
                if len(results) != len(images):
                    raise RuntimeError("detector returned wrong result count")
            except Exception:
                traceback.print_exc()
                results = []
                for image in images:
                    try:
                        results.append(list(predictor([image]))[0])
                    except Exception as exc:
                        results.append(exc)
            for (index, row, image, layer), result in zip(pending, results):
                lines = []
                error = None
                if isinstance(result, Exception):
                    failures += 1
                    error = f"{type(result).__name__}: {result}"
                else:
                    lines = line_payload(result)
                    detected_lines += len(lines)
                payload = {
                    "id": f"handwriting:{index}",
                    "split": "handwriting",
                    "reference": str(row.get("label") or ""),
                    "width": image.width,
                    "height": image.height,
                    "variant": args.variant,
                    "lines": lines,
                    **layer,
                }
                if error:
                    payload["error"] = error
                output.write(json.dumps(payload, ensure_ascii=False) + "\n")
                completed += 1
            pending.clear()

        for index in assigned:
            if index not in wanted:
                continue
            row = rows[index]
            try:
                image = decode_image(row)
                isolated, match, metrics = catalog.isolate(image, args.variant)
                layer = {
                    "template_id": match.template.id,
                    "template_file": match.template.file,
                    "template_text": match.template.template_text,
                    "template_placeholder_count": match.template.template_text.count(
                        "{{HANDWRITING}}"
                    ),
                    "match_median_abs_delta": match.median_abs_delta,
                    "match_p90_abs_delta": match.p90_abs_delta,
                    "mask_threshold": metrics["threshold"],
                    "mask_kept_fraction": metrics["kept_fraction"],
                }
                pending.append((index, row, isolated, layer))
                if len(pending) >= args.batch_size:
                    flush()
            except Exception as exc:
                traceback.print_exc()
                failures += 1
                match_failures += 1
                completed += 1
                output.write(
                    json.dumps(
                        {
                            "id": f"handwriting:{index}",
                            "split": "handwriting",
                            "reference": str(row.get("label") or ""),
                            "width": 0,
                            "height": 0,
                            "variant": args.variant,
                            "lines": [],
                            "error": f"{type(exc).__name__}: {exc}",
                        },
                        ensure_ascii=False,
                    )
                    + "\n"
                )
        flush()

    elapsed = max(time.monotonic() - started, 1e-9)
    atomic_json(
        status_path,
        {
            "state": "complete",
            "rank": args.rank,
            "assigned": len(assigned),
            "completed": completed,
            "failures": failures,
            "template_match_failures": match_failures,
            "detected_lines": detected_lines,
            "rate_pages_per_second": completed / elapsed,
            "eta_seconds": 0,
            "updated_unix": time.time(),
        },
    )
    merge_if_complete(args.output_dir, args.world_size, selected, args.variant)


if __name__ == "__main__":
    main()