#!/usr/bin/env python3
"""Split a PDF, identify each BOS form by QR, and route pages to OMRChecker."""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from omr_service.form_routing import identify_form  # noqa: E402
from omr_service.input_validation import render_pdf_pages  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("pdf", type=Path)
    parser.add_argument("--omrchecker", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--dpi", type=int, default=300)
    args = parser.parse_args()

    pdf = args.pdf.resolve()
    output = args.output.resolve()
    omr_main = args.omrchecker.resolve() / "main.py"
    if not pdf.is_file() or not omr_main.is_file():
        raise FileNotFoundError(pdf if not pdf.is_file() else omr_main)
    if output.exists():
        raise FileExistsError(f"Refusing to replace existing output: {output}")

    rendered_dir = output / "rendered-pages"
    rendered_dir.mkdir(parents=True)

    started = time.perf_counter()
    page_map, validation_issues = render_pdf_pages(pdf, rendered_dir, args.dpi)

    groups: dict[str, list[dict[str, object]]] = defaultdict(list)
    route_by_form = {}
    routing_issues = []
    routed_pages = []

    for item in page_map:
        image_path = rendered_dir / str(item["file_id"])
        route, issue = identify_form(image_path, ROOT)
        if issue:
            routing_issues.append({"page": item["page"], **issue.as_dict()})
            continue
        assert route is not None
        groups[route.form_id].append(item)
        route_by_form[route.form_id] = route
        routed_pages.append(
            {
                **item,
                "form_id": route.form_id,
                "profile_name": route.profile_name,
                "rotation_clockwise": route.rotation_clockwise,
            }
        )

    runs = []
    return_code = 0
    for form_id, items in groups.items():
        route = route_by_form[form_id]
        form_root = output / "forms" / form_id
        input_dir = form_root / "input"
        omr_output = form_root / "omr-output"
        input_dir.mkdir(parents=True)
        omr_output.mkdir()

        shutil.copy2(route.template, input_dir / "template.json")
        shutil.copy2(route.config, input_dir / "config.json")
        shutil.copy2(route.marker, input_dir / "omr_marker.jpg")
        for item in items:
            shutil.copy2(rendered_dir / str(item["file_id"]), input_dir / str(item["file_id"]))

        run_started = time.perf_counter()
        completed = subprocess.run(
            [sys.executable, str(omr_main), "-i", str(input_dir), "-o", str(omr_output)],
            cwd=args.omrchecker.resolve(),
            text=True,
            capture_output=True,
        )
        elapsed = time.perf_counter() - run_started
        (form_root / "console.stdout.txt").write_text(completed.stdout)
        (form_root / "console.stderr.txt").write_text(completed.stderr)
        runs.append(
            {
                "form_id": form_id,
                "profile_name": route.profile_name,
                "pages": len(items),
                "seconds": round(elapsed, 3),
                "omr_return_code": completed.returncode,
            }
        )
        if completed.returncode:
            return_code = completed.returncode

    elapsed = time.perf_counter() - started
    summary = {
        "source_pdf": str(pdf),
        "dpi": args.dpi,
        "rendered_pages": len(page_map),
        "routed_pages": len(routed_pages),
        "total_seconds": round(elapsed, 3),
        "status": "rejected" if not routed_pages else "partial" if validation_issues or routing_issues else "ok",
        "validation_issues": [issue.as_dict() for issue in validation_issues],
        "routing_issues": routing_issues,
        "pages": routed_pages,
        "runs": runs,
    }
    (output / "batch-summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))

    # QR failures are intentionally non-zero: BOS must not guess the form profile.
    if validation_issues or routing_issues:
        return 2
    return return_code


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