#!/usr/bin/env python3
"""Write deterministic per-image predictions from a CartoLegend point detector."""

from __future__ import annotations

import argparse
import json
from datetime import datetime, timezone
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import DataLoader

from train_cartolegend_point_detector import (
    PointSymbolDataset,
    build_model,
    collate,
    load_unique_records,
    sha256,
)
from cartolegend_point_detector_artifact import load_point_detector_artifact


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


def portable_path(path: Path) -> str:
    resolved = path.expanduser().resolve()
    try:
        return "project://" + resolved.relative_to(ROOT).as_posix()
    except ValueError:
        return resolved.name


def load_prediction_records(path: Path, image_field: str) -> list[dict]:
    if image_field == "images":
        return load_unique_records(path)
    records = []
    seen: set[Path] = set()
    for line_number, line in enumerate(path.read_text().splitlines(), start=1):
        if not line.strip():
            continue
        row = json.loads(line)
        value = str(row.get(image_field) or "")
        image = Path(value)
        image = (image if image.is_absolute() else ROOT / image).resolve()
        if not image.is_file():
            raise FileNotFoundError(f"line {line_number}: {image}")
        if image in seen:
            continue
        seen.add(image)
        with Image.open(image) as source:
            width, height = source.size
        records.append(
            {
                "image": str(image),
                "width": width,
                "height": height,
                "boxes": [],
                "text_labels": [],
                "group": str(row.get("source_id") or image.stem),
            }
        )
    return records


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--config", type=Path, default=None)
    parser.add_argument(
        "--allow-legacy-pt",
        action="store_true",
        help="Explicitly allow a legacy .pt checkpoint through weights_only=True.",
    )
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--image-field", default="images")
    parser.add_argument("--workers", type=int, default=2)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    checkpoint_path = args.checkpoint.expanduser().resolve()
    input_path = args.input.expanduser().resolve()
    output_path = args.output.expanduser().resolve()
    artifact = load_point_detector_artifact(
        checkpoint_path,
        config_path=args.config,
        device="cpu",
        allow_legacy_pt=args.allow_legacy_pt,
    )
    config = artifact.model_config
    model = build_model(
        False,
        int(config["min_size"]),
        int(config["max_size"]),
        int(config.get("trainable_backbone_layers", 6)),
        str(config.get("architecture", "mobilenet")),
    )
    model.load_state_dict(artifact.state_dict, strict=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device).eval()

    records = load_prediction_records(input_path, args.image_field)
    loader = DataLoader(
        PointSymbolDataset(records, augment=False),
        batch_size=1,
        shuffle=False,
        num_workers=args.workers,
        collate_fn=collate,
        pin_memory=device.type == "cuda",
        persistent_workers=args.workers > 0,
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w") as output_file, torch.inference_mode():
        for record, (images, _targets) in zip(records, loader, strict=True):
            prediction = model([images[0].to(device)])[0]
            boxes = prediction["boxes"].detach().cpu().tolist()
            scores = prediction["scores"].detach().cpu().tolist()
            row = {
                "schema": "cartolegend_point_detector_predictions_v1",
                "image": portable_path(Path(record["image"])),
                "width": record["width"],
                "height": record["height"],
                "detections": [
                    {
                        "symbol_bbox": [round(float(value), 4) for value in box],
                        "score": round(float(score), 8),
                    }
                    for box, score in zip(boxes, scores, strict=True)
                ],
            }
            output_file.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")

    manifest = {
        "schema": "cartolegend_point_detector_prediction_manifest_v1",
        "generated_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "checkpoint": portable_path(checkpoint_path),
        "checkpoint_sha256": artifact.sha256,
        "checkpoint_epoch": artifact.metadata.get("completed_epochs"),
        "input": portable_path(input_path),
        "input_sha256": sha256(input_path),
        "output": portable_path(output_path),
        "output_sha256": sha256(output_path),
        "images": len(records),
        "image_field": args.image_field,
        "device": str(device),
        "config": config,
    }
    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 0


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