#!/usr/bin/env python3
"""Recognize cached Surya boxes from the same isolated handwriting layer."""

from __future__ import annotations

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

import numpy as np
import pyarrow.parquet as pq
from paddleocr import TextRecognition
from PIL import Image

from bina02_template_layers import TemplateCatalog, index_from_id


ROWS = 6_669
LTR_RUN = re.compile(r"[a-zA-Z0-9 :*./%+-]")


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("--detections-dir", type=Path, required=True)
    parser.add_argument("--model-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--variant", required=True)
    parser.add_argument("--device", default="gpu:0")
    parser.add_argument("--crop-buffer", type=int, default=512)
    parser.add_argument("--recognition-batch-size", type=int, default=64)
    parser.add_argument("--pad-x-ratio", type=float, default=0.01)
    parser.add_argument("--pad-y-ratio", type=float, default=0.08)
    parser.add_argument("--row-factor", type=float, default=0.70)
    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 pred_reverse(text: str) -> str:
    segments: list[str] = []
    current = ""
    for character in text:
        if LTR_RUN.search(character):
            current += character
        else:
            if current:
                segments.append(current)
                current = ""
            segments.append(character)
    if current:
        segments.append(current)
    return "".join(reversed(segments))


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"))
    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("benchmark row load contract failed")
    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 center(line: dict[str, Any]) -> tuple[float, float]:
    x0, y0, x1, y1 = (float(value) for value in line["bbox"])
    return (x0 + x1) / 2.0, (y0 + y1) / 2.0


def height(line: dict[str, Any]) -> float:
    return max(float(line["bbox"][3]) - float(line["bbox"][1]), 1.0)


def group_rows(lines: list[dict[str, Any]], factor: float) -> list[list[dict[str, Any]]]:
    ordered = sorted(lines, key=lambda line: center(line)[1])
    rows: list[list[dict[str, Any]]] = []
    for line in ordered:
        if not rows:
            rows.append([line])
            continue
        previous = rows[-1]
        mean_y = sum(center(item)[1] for item in previous) / len(previous)
        mean_height = sum(height(item) for item in previous) / len(previous)
        if abs(center(line)[1] - mean_y) <= factor * max(height(line), mean_height):
            previous.append(line)
        else:
            rows.append([line])
    for row in rows:
        row.sort(key=lambda line: center(line)[0], reverse=True)
    return rows


def grouped_lines(
    lines: list[dict[str, Any]],
    width: int,
    height_value: int,
    factor: float,
) -> list[dict[str, Any]]:
    aspect = width / max(height_value, 1)
    if 1.03 <= aspect <= 1.42:
        midpoint = width / 2.0
        groups = group_rows(
            [line for line in lines if center(line)[0] >= midpoint],
            factor,
        ) + group_rows(
            [line for line in lines if center(line)[0] < midpoint],
            factor,
        )
    else:
        groups = group_rows(lines, factor)
    return [
        {
            "bbox": [
                min(float(line["bbox"][0]) for line in row),
                min(float(line["bbox"][1]) for line in row),
                max(float(line["bbox"][2]) for line in row),
                max(float(line["bbox"][3]) for line in row),
            ],
            "confidence": sum(float(line.get("confidence", 0.0)) for line in row)
            / len(row),
            "member_count": len(row),
        }
        for row in groups
    ]


def crop(
    image: Image.Image,
    bbox: list[float],
    pad_x_ratio: float,
    pad_y_ratio: float,
) -> np.ndarray:
    x0, y0, x1, y1 = (float(value) for value in bbox)
    width = max(x1 - x0, 1.0)
    height_value = max(y1 - y0, 1.0)
    pad_x = max(2, int(width * pad_x_ratio))
    pad_y = max(2, int(height_value * pad_y_ratio))
    box = (
        max(0, int(x0) - pad_x),
        max(0, int(y0) - pad_y),
        min(image.width, int(x1) + pad_x),
        min(image.height, int(y1) + pad_y),
    )
    if box[2] <= box[0] or box[3] <= box[1]:
        raise ValueError(f"invalid crop: {bbox}")
    return np.asarray(image.crop(box).convert("RGB"))


def recognize(
    model: TextRecognition,
    crops: list[np.ndarray],
    batch_size: int,
) -> list[dict[str, Any]]:
    if not crops:
        return []
    last_error: Exception | None = None
    for attempt in range(1, 4):
        try:
            results = list(model.predict(input=crops, batch_size=batch_size))
            if len(results) != len(crops):
                raise RuntimeError("recognizer returned wrong result count")
            output = []
            for result in results:
                payload = result.json
                if callable(payload):
                    payload = payload()
                payload = payload.get("res", payload)
                raw = str(payload.get("rec_text", ""))
                output.append(
                    {
                        "text": pred_reverse(raw),
                        "raw_text": raw,
                        "rec_score": float(payload.get("rec_score", 0.0)),
                    }
                )
            return output
        except Exception as exc:
            last_error = exc
            traceback.print_exc()
            time.sleep(attempt)
    raise RuntimeError(f"recognition failed after retries: {last_error}")


def merge_if_complete(output_dir: Path, detections_dir: Path, world_size: int) -> None:
    with (output_dir / "merge.lock").open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        for rank in range(world_size):
            path = output_dir / f"status-rank{rank:02d}.json"
            if not path.is_file() or json.loads(path.read_text()).get("state") != "complete":
                return
        expected = [
            json.loads(line)["id"]
            for line in (detections_dir / "detections.jsonl").read_text(
                encoding="utf-8"
            ).splitlines()
            if line.strip()
        ]
        rows = {}
        for rank in range(world_size):
            path = output_dir / f"recognized-rank{rank:02d}.jsonl"
            for line in path.read_text(encoding="utf-8").splitlines():
                if line.strip():
                    row = json.loads(line)
                    rows[row["id"]] = row
        if set(rows) != set(expected):
            raise RuntimeError("recognition merge contract failed")
        temporary = output_dir / "recognized-lines.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 / "recognized-lines.jsonl")
        atomic_json(
            output_dir / "run-summary.json",
            {"state": "complete", "rows": len(rows), "completed_unix": time.time()},
        )


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    detections_path = args.detections_dir / f"detections-rank{args.rank:02d}.jsonl"
    detections = [
        json.loads(line)
        for line in detections_path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    indices = {index_from_id(row["id"]) for row in detections}
    rows = load_rows(args.data, indices)
    catalog = TemplateCatalog(args.background_root, args.ocr_root)
    model = TextRecognition(model_dir=str(args.model_dir), device=args.device)
    output_path = args.output_dir / f"recognized-rank{args.rank:02d}.jsonl"
    status_path = args.output_dir / f"status-rank{args.rank:02d}.json"
    started = time.monotonic()
    failures = 0
    recognized_lines = 0
    page_buffer: list[dict[str, Any]] = []
    crop_buffer: list[np.ndarray] = []

    with output_path.open("w", encoding="utf-8", buffering=1) as output:
        def flush() -> None:
            nonlocal failures, recognized_lines
            if not page_buffer:
                return
            try:
                values = recognize(model, crop_buffer, args.recognition_batch_size)
                recognition_error = None
            except Exception as exc:
                values = [
                    {"text": "", "raw_text": "", "rec_score": 0.0}
                    for _ in crop_buffer
                ]
                recognition_error = f"{type(exc).__name__}: {exc}"
            for page in page_buffer:
                if recognition_error and not page.get("error"):
                    page["error"] = recognition_error
                if page.get("error"):
                    failures += 1
                for line in page["lines"]:
                    result_index = line.pop("_result_index", None)
                    if result_index is not None:
                        line.update(values[result_index])
                        recognized_lines += 1
                output.write(json.dumps(page, ensure_ascii=False) + "\n")
            page_buffer.clear()
            crop_buffer.clear()

        for detection in detections:
            index = index_from_id(detection["id"])
            page = {
                key: detection.get(key)
                for key in (
                    "id",
                    "reference",
                    "width",
                    "height",
                    "variant",
                    "template_id",
                    "template_file",
                    "template_text",
                    "template_placeholder_count",
                    "match_median_abs_delta",
                    "match_p90_abs_delta",
                    "mask_threshold",
                    "mask_kept_fraction",
                )
            }
            page["lines"] = []
            if detection.get("error"):
                page["error"] = detection["error"]
            try:
                image = decode_image(rows[index])
                template = catalog.by_id[str(detection["template_id"])]
                isolated, _ = catalog.isolate_with_template(
                    image,
                    template,
                    args.variant,
                )
                for source in grouped_lines(
                    detection.get("lines", []),
                    isolated.width,
                    isolated.height,
                    args.row_factor,
                ):
                    line = dict(source)
                    line["_result_index"] = len(crop_buffer)
                    crop_buffer.append(
                        crop(
                            isolated,
                            source["bbox"],
                            args.pad_x_ratio,
                            args.pad_y_ratio,
                        )
                    )
                    page["lines"].append(line)
            except Exception as exc:
                traceback.print_exc()
                page["error"] = f"{type(exc).__name__}: {exc}"
            page_buffer.append(page)
            if len(crop_buffer) >= args.crop_buffer:
                flush()
        flush()

    elapsed = max(time.monotonic() - started, 1e-9)
    atomic_json(
        status_path,
        {
            "state": "complete",
            "rank": args.rank,
            "assigned": len(detections),
            "failures": failures,
            "recognized_lines": recognized_lines,
            "rate_pages_per_second": len(detections) / elapsed,
            "eta_seconds": 0,
        },
    )
    merge_if_complete(args.output_dir, args.detections_dir, args.world_size)


if __name__ == "__main__":
    main()