#!/usr/bin/env python3
"""Extract tiled OCR from source-resolution crops and project it to model pixels."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from PIL import Image

from extract_cartolegend_ocr import parse_tsv


ROOT = Path(__file__).resolve().parents[1]
Image.MAX_IMAGE_PIXELS = None


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def resolve_path(value: str | Path) -> Path:
    path = Path(value).expanduser()
    return (path if path.is_absolute() else ROOT / path).resolve()


def row_image(row: dict[str, Any], image_field: str) -> Path:
    if image_field == "images":
        values = row.get("images") or []
        if len(values) != 1:
            raise ValueError("each input row must contain exactly one image")
        value = values[0]
    else:
        value = str(row.get(image_field) or "")
        if not value:
            raise ValueError(f"input row is missing {image_field}")
    image = resolve_path(value)
    if not image.is_file():
        raise FileNotFoundError(image)
    return image


def safe_name(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("_") or "image"


def axis_starts(length: int, tile_size: int, overlap: int) -> list[int]:
    if length <= 0 or tile_size <= 0:
        raise ValueError("length and tile_size must be positive")
    if overlap < 0 or overlap >= tile_size:
        raise ValueError("overlap must be in [0, tile_size)")
    if length <= tile_size:
        return [0]
    step = tile_size - overlap
    starts = list(range(0, max(1, length - tile_size + 1), step))
    final = length - tile_size
    if starts[-1] != final:
        starts.append(final)
    return starts


def tile_boxes(width: int, height: int, tile_size: int, overlap: int) -> list[list[int]]:
    return [
        [x, y, min(width, x + tile_size), min(height, y + tile_size)]
        for y in axis_starts(height, tile_size, overlap)
        for x in axis_starts(width, tile_size, overlap)
    ]


def project_bbox(
    bbox: list[int | float],
    stage_size: tuple[int, int],
    target_size: tuple[int, int],
) -> list[int]:
    stage_width, stage_height = stage_size
    target_width, target_height = target_size
    if len(bbox) != 4 or stage_width <= 0 or stage_height <= 0:
        raise ValueError("invalid projection geometry")
    x1 = round(float(bbox[0]) * target_width / stage_width)
    y1 = round(float(bbox[1]) * target_height / stage_height)
    x2 = round(float(bbox[2]) * target_width / stage_width)
    y2 = round(float(bbox[3]) * target_height / stage_height)
    x1 = min(max(0, x1), max(0, target_width - 1))
    y1 = min(max(0, y1), max(0, target_height - 1))
    x2 = min(max(x1 + 1, x2), target_width)
    y2 = min(max(y1 + 1, y2), target_height)
    return [x1, y1, x2, y2]


def source_crop(
    row: dict[str, Any], target: Path
) -> tuple[Path, list[int] | None, str]:
    bbox = row.get("crop_bbox")
    candidates = [row.get("crop_source_image"), row.get("original_image")]
    if (
        isinstance(bbox, list)
        and len(bbox) == 4
        and all(isinstance(value, (int, float)) for value in bbox)
        and float(bbox[0]) < float(bbox[2])
        and float(bbox[1]) < float(bbox[3])
    ):
        for value in candidates:
            if not value:
                continue
            candidate = resolve_path(str(value))
            if candidate.is_file():
                return candidate, [round(float(value)) for value in bbox], "source_crop"
    return target, None, "target_upscale"


def stage_row_image(
    row: dict[str, Any],
    target: Path,
    output: Path,
    requested_scale: float,
    max_dimension: int,
) -> dict[str, Any]:
    with Image.open(target) as image:
        target_size = (image.width, image.height)
    source, requested_bbox, mode = source_crop(row, target)
    with Image.open(source) as opened:
        source_size = (opened.width, opened.height)
        if requested_bbox is None:
            clipped_bbox = [0, 0, opened.width, opened.height]
        else:
            clipped_bbox = [
                min(max(0, requested_bbox[0]), max(0, opened.width - 1)),
                min(max(0, requested_bbox[1]), max(0, opened.height - 1)),
                min(max(1, requested_bbox[2]), opened.width),
                min(max(1, requested_bbox[3]), opened.height),
            ]
            if clipped_bbox[0] >= clipped_bbox[2] or clipped_bbox[1] >= clipped_bbox[3]:
                raise ValueError(f"crop_bbox does not intersect source image: {source}")
        crop = opened.crop(tuple(clipped_bbox)).convert("RGB")
    effective_scale = min(
        float(requested_scale),
        float(max_dimension) / max(target_size),
    )
    effective_scale = max(1.0, effective_scale)
    stage_size = (
        max(1, round(target_size[0] * effective_scale)),
        max(1, round(target_size[1] * effective_scale)),
    )
    if crop.size != stage_size:
        crop = crop.resize(stage_size, Image.Resampling.LANCZOS)
    output.parent.mkdir(parents=True, exist_ok=True)
    crop.save(output, compress_level=3)
    return {
        "mode": mode,
        "target_image": str(target),
        "target_size": list(target_size),
        "source_image": str(source),
        "source_size": list(source_size),
        "source_crop_bbox": clipped_bbox,
        "requested_scale": requested_scale,
        "effective_scale": effective_scale,
        "stage_image": str(output),
        "stage_size": list(stage_size),
    }


def run_tesseract(
    tile_path: Path, psm: int, timeout_seconds: int
) -> tuple[int, str, str]:
    try:
        result = subprocess.run(
            [
                "tesseract",
                str(tile_path),
                "stdout",
                "--oem",
                "1",
                "--psm",
                str(psm),
                "tsv",
            ],
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
            check=False,
        )
        return result.returncode, result.stdout, result.stderr.strip()[:500]
    except subprocess.TimeoutExpired as error:
        stderr = (error.stderr or "") if isinstance(error.stderr, str) else ""
        return 124, "", (stderr + " tesseract timeout").strip()[:500]


def offset_and_project(
    values: list[dict[str, Any]],
    tile_bbox: list[int],
    tile_id: str,
    stage_size: tuple[int, int],
    target_size: tuple[int, int],
) -> list[dict[str, Any]]:
    projected = []
    for value in values:
        local = value["bbox"]
        stage_bbox = [
            local[0] + tile_bbox[0],
            local[1] + tile_bbox[1],
            local[2] + tile_bbox[0],
            local[3] + tile_bbox[1],
        ]
        projected.append(
            {
                **value,
                "bbox": project_bbox(stage_bbox, stage_size, target_size),
                "stage_bbox": stage_bbox,
                "tile_id": tile_id,
                "tile_bbox": tile_bbox,
            }
        )
    return projected


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--stage-dir", type=Path, required=True)
    parser.add_argument("--image-field", default="images")
    parser.add_argument("--scale", type=float, default=4.0)
    parser.add_argument("--max-dimension", type=int, default=5200)
    parser.add_argument("--tile-size", type=int, default=1400)
    parser.add_argument("--tile-overlap", type=int, default=220)
    parser.add_argument("--psm", type=int, action="append", default=[])
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--timeout-seconds", type=int, default=120)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    input_path = args.input.expanduser().resolve()
    output_path = args.output.expanduser().resolve()
    stage_dir = args.stage_dir.expanduser().resolve()
    psm_modes = sorted(set(args.psm or [6, 11]))
    if args.scale < 1.0 or args.max_dimension <= 0 or args.workers <= 0:
        raise ValueError("invalid scale, max dimension, or worker count")
    rows = read_jsonl(input_path)
    targets = [row_image(row, args.image_field) for row in rows]
    if len(set(targets)) != len(targets):
        raise ValueError("input image paths must be unique")
    version = subprocess.run(
        ["tesseract", "--version"], capture_output=True, text=True, check=True
    ).stdout.splitlines()[0]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    stage_dir.mkdir(parents=True, exist_ok=True)
    failures: list[dict[str, Any]] = []
    artifact_rows: list[dict[str, Any]] = []
    hash_cache: dict[Path, str] = {}

    with output_path.open("w") as output_stream:
        for row_index, (row, target) in enumerate(zip(rows, targets, strict=True), start=1):
            stage_path = stage_dir / (
                f"{row_index:04d}_{safe_name(target.stem)}__source_ocr.png"
            )
            stage = stage_row_image(
                row,
                target,
                stage_path,
                args.scale,
                args.max_dimension,
            )
            stage_size = tuple(stage["stage_size"])
            target_size = tuple(stage["target_size"])
            boxes = tile_boxes(
                stage_size[0],
                stage_size[1],
                args.tile_size,
                args.tile_overlap,
            )
            tasks: dict[Any, tuple[str, list[int], int]] = {}
            all_words: list[dict[str, Any]] = []
            all_lines: list[dict[str, Any]] = []
            mode_rows: list[dict[str, Any]] = []
            with tempfile.TemporaryDirectory(prefix="cartolegend_source_ocr_") as raw_tmp:
                temporary = Path(raw_tmp)
                with Image.open(stage_path) as staged_image:
                    for tile_index, box in enumerate(boxes, start=1):
                        tile_id = f"tile_{tile_index:03d}"
                        tile_path = temporary / f"{tile_id}.png"
                        staged_image.crop(tuple(box)).save(tile_path, compress_level=1)
                        for psm in psm_modes:
                            tasks[(tile_id, psm)] = (str(tile_path), box, psm)
                with ThreadPoolExecutor(max_workers=args.workers) as executor:
                    futures = {
                        executor.submit(
                            run_tesseract,
                            Path(tile_path),
                            psm,
                            args.timeout_seconds,
                        ): (tile_id, box, psm)
                        for (tile_id, psm), (tile_path, box, _mode) in tasks.items()
                    }
                    for future in as_completed(futures):
                        tile_id, box, psm = futures[future]
                        returncode, stdout, stderr = future.result()
                        words, lines = parse_tsv(stdout, psm)
                        all_words.extend(
                            offset_and_project(
                                words, box, tile_id, stage_size, target_size
                            )
                        )
                        all_lines.extend(
                            offset_and_project(
                                lines, box, tile_id, stage_size, target_size
                            )
                        )
                        mode_rows.append(
                            {
                                "tile_id": tile_id,
                                "tile_bbox": box,
                                "psm": psm,
                                "returncode": returncode,
                                "words": len(words),
                                "lines": len(lines),
                                "stderr": stderr,
                            }
                        )
                        if returncode:
                            failures.append(
                                {
                                    "image": str(target),
                                    "tile_id": tile_id,
                                    "psm": psm,
                                    "returncode": returncode,
                                }
                            )
            all_words.sort(
                key=lambda value: (
                    value["bbox"][1],
                    value["bbox"][0],
                    value["psm"],
                    value["tile_id"],
                )
            )
            all_lines.sort(
                key=lambda value: (
                    value["bbox"][1],
                    value["bbox"][0],
                    value["psm"],
                    value["tile_id"],
                )
            )
            source_path = Path(stage["source_image"])
            for path in (target, source_path, stage_path):
                hash_cache.setdefault(path, sha256(path))
            output_row = {
                "schema": "cartolegend_source_projected_ocr_v1",
                "image": str(target),
                "image_sha256": hash_cache[target],
                "source_image": str(source_path),
                "source_image_sha256": hash_cache[source_path],
                "source_crop_bbox": stage["source_crop_bbox"],
                "source_mode": stage["mode"],
                "target_size": stage["target_size"],
                "stage_image": str(stage_path),
                "stage_image_sha256": hash_cache[stage_path],
                "stage_size": stage["stage_size"],
                "effective_scale": stage["effective_scale"],
                "tiles": len(boxes),
                "psm_modes": psm_modes,
                "runs": sorted(mode_rows, key=lambda value: (value["tile_id"], value["psm"])),
                "words": all_words,
                "lines": all_lines,
                "model_outputs_are_proposals_not_annotations": True,
                "training_allowed": False,
            }
            output_stream.write(
                json.dumps(output_row, sort_keys=True, separators=(",", ":")) + "\n"
            )
            artifact_rows.append(
                {
                    key: output_row[key]
                    for key in (
                        "image",
                        "image_sha256",
                        "source_image",
                        "source_image_sha256",
                        "source_crop_bbox",
                        "source_mode",
                        "stage_image",
                        "stage_image_sha256",
                        "target_size",
                        "stage_size",
                        "effective_scale",
                        "tiles",
                    )
                }
            )
            print(
                f"source-ocr={row_index}/{len(rows)} image={target.name} "
                f"tiles={len(boxes)} words={len(all_words)} lines={len(all_lines)}",
                flush=True,
            )

    manifest = {
        "schema": "cartolegend_source_projected_ocr_manifest_v1",
        "generated_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "tesseract_version": version,
        "input": str(input_path),
        "input_sha256": sha256(input_path),
        "output": str(output_path),
        "output_sha256": sha256(output_path),
        "rows": len(rows),
        "settings": {
            "image_field": args.image_field,
            "requested_scale": args.scale,
            "max_dimension": args.max_dimension,
            "tile_size": args.tile_size,
            "tile_overlap": args.tile_overlap,
            "psm_modes": psm_modes,
            "workers": args.workers,
            "timeout_seconds": args.timeout_seconds,
        },
        "artifacts": artifact_rows,
        "failures": failures,
        "model_outputs_are_proposals_not_annotations": True,
        "training_allowed": False,
    }
    manifest_path = output_path.with_suffix(output_path.suffix + ".manifest.json")
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(json.dumps(manifest, indent=2, sort_keys=True))
    return 1 if failures else 0


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