#!/usr/bin/env python3
"""Export one registered YOLOv8 detector into the strict edge ONNX contract."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

# Ultralytics normally tries to install missing export dependencies at runtime.
# That would bypass uv.lock and make deployment artifacts irreproducible.
os.environ["YOLO_AUTOINSTALL"] = "false"

from cmvr_edge_ai.detection.models.mobile_phone import MOBILE_PHONE_LABELS
from cmvr_edge_ai.detection.models.people_talking import PEOPLE_TALKING_LABELS
from cmvr_edge_ai.detection.models.ppe_yolo import (
    CONSTRUCTION_PPE_LABELS,
    PPE_6CLASS_LABELS,
)
from cmvr_edge_ai.detection.models.yolo_onnx import OnnxYoloModel


@dataclass(frozen=True, slots=True)
class ExportDefinition:
    source_model_id: str
    target_model_id: str
    source: Path
    output: Path
    labels: tuple[str, ...]


EXPORT_DEFINITIONS = {
    "construction-ppe-yolov8@2": ExportDefinition(
        source_model_id="construction-ppe-yolov8@1",
        target_model_id="construction-ppe-yolov8@2",
        source=Path("models/detection/construction-ppe-yolov8/v1/best.pt"),
        output=Path("models/detection/construction-ppe-yolov8/v2/model.onnx"),
        labels=CONSTRUCTION_PPE_LABELS,
    ),
    "ppe-6classes-yolov8n@2": ExportDefinition(
        source_model_id="ppe-6classes-yolov8n@1",
        target_model_id="ppe-6classes-yolov8n@2",
        source=Path("models/detection/ppe-6classes-yolov8n/v1/best.pt"),
        output=Path("models/detection/ppe-6classes-yolov8n/v2/model.onnx"),
        labels=PPE_6CLASS_LABELS,
    ),
    "people-talking-yolov8x@2": ExportDefinition(
        source_model_id="people-talking-yolov8x@1",
        target_model_id="people-talking-yolov8x@2",
        source=Path("models/detection/people-talking-yolov8x/v1/best.pt"),
        output=Path("models/detection/people-talking-yolov8x/v2/model.onnx"),
        labels=PEOPLE_TALKING_LABELS,
    ),
    "yolov8n-mobile-phone@2": ExportDefinition(
        source_model_id="yolov8n-mobile-phone@1",
        target_model_id="yolov8n-mobile-phone@2",
        source=Path(
            "models/detection/yolov8n-mobile-phone/yolov8n-mobile-phone.pt"
        ),
        output=Path("models/detection/yolov8n-mobile-phone/v2/model.onnx"),
        labels=MOBILE_PHONE_LABELS,
    ),
}


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model-id",
        required=True,
        choices=tuple(sorted(EXPORT_DEFINITIONS)),
        help="registered @2 ONNX model to build",
    )
    parser.add_argument(
        "--source",
        type=Path,
        help="override the registered trusted .pt source path",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="override the registered v2/model.onnx destination",
    )
    parser.add_argument(
        "--manifest",
        type=Path,
        help="manifest destination (default: manifest.json beside the ONNX file)",
    )
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument(
        "--simplify",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="replace an existing validated output and manifest",
    )
    return parser.parse_args(argv)


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


def _ordered_names(value: Any) -> tuple[str, ...]:
    if isinstance(value, dict):
        normalized = {int(key): str(label) for key, label in value.items()}
        indexes = sorted(normalized)
        if indexes != list(range(len(indexes))):
            raise ValueError("source checkpoint class IDs must start at zero")
        return tuple(normalized[index] for index in indexes)
    if isinstance(value, (list, tuple)):
        return tuple(str(label) for label in value)
    raise ValueError(f"source checkpoint has invalid class names: {value!r}")


def _sanitize_onnx_metadata(
    onnx: Any,
    path: Path,
    *,
    definition: ExportDefinition,
    source_sha256: str,
) -> Any:
    """Remove build-host paths/timestamps and add stable CMVR identity."""

    model = onnx.load(str(path))
    metadata = {
        str(item.key): str(item.value)
        for item in getattr(model, "metadata_props", ())
    }
    metadata.pop("date", None)
    metadata["description"] = (
        f"CMVR static ONNX export for {definition.target_model_id}"
    )
    metadata["cmvr_model_id"] = definition.target_model_id
    metadata["cmvr_source_sha256"] = source_sha256
    onnx.helper.set_model_props(model, metadata)
    onnx.save_model(model, str(path))
    return model


def _require_export_dependencies(*, simplify: bool) -> dict[str, Any]:
    try:
        import onnx
        import onnxruntime
        import torch
        import ultralytics
        from ultralytics import YOLO

        onnxslim = None
        if simplify:
            import onnxslim
    except ImportError as exc:
        raise RuntimeError(
            "ONNX export dependencies are incomplete; run uv with the "
            "onnx-export-cpu extra"
        ) from exc
    return {
        "onnx": onnx,
        "onnxruntime": onnxruntime,
        "onnxslim": onnxslim,
        "torch": torch,
        "ultralytics": ultralytics,
        "YOLO": YOLO,
    }


def _validate_paths(
    *,
    source: Path,
    output: Path,
    manifest: Path,
    force: bool,
) -> None:
    if not source.is_file():
        raise FileNotFoundError(f"source checkpoint does not exist: {source}")
    if source.suffix.lower() != ".pt":
        raise ValueError(f"source checkpoint must end in .pt: {source}")
    if output.suffix.lower() != ".onnx":
        raise ValueError(f"output artifact must end in .onnx: {output}")
    if output.resolve() == source.resolve():
        raise ValueError("source and output paths must be different")
    if manifest.resolve() in {source.resolve(), output.resolve()}:
        raise ValueError("manifest path must differ from source and output")
    existing = [path for path in (output, manifest) if path.exists()]
    if existing and not force:
        raise FileExistsError(
            "refusing to replace existing export path(s) without --force: "
            + ", ".join(str(path) for path in existing)
        )
    with source.open("rb") as artifact:
        prefix = artifact.read(128)
    if b"version https://git-lfs.github.com/spec/v1" in prefix:
        raise ValueError(f"source checkpoint is a Git LFS pointer: {source}")


def _write_manifest(path: Path, payload: dict[str, Any]) -> None:
    with path.open("x", encoding="utf-8") as stream:
        json.dump(payload, stream, ensure_ascii=False, indent=2, sort_keys=True)
        stream.write("\n")


def export_model(args: argparse.Namespace) -> tuple[Path, Path]:
    definition = EXPORT_DEFINITIONS[args.model_id]
    source = (args.source or definition.source).expanduser()
    output = (args.output or definition.output).expanduser()
    manifest = (
        args.manifest.expanduser()
        if args.manifest is not None
        else output.with_name("manifest.json")
    )
    if args.imgsz < 32 or args.imgsz % 32:
        raise ValueError("imgsz must be at least 32 and divisible by 32")
    if args.opset < 12:
        raise ValueError("opset must be at least 12")
    _validate_paths(
        source=source,
        output=output,
        manifest=manifest,
        force=args.force,
    )
    dependencies = _require_export_dependencies(simplify=args.simplify)
    output.parent.mkdir(parents=True, exist_ok=True)
    manifest.parent.mkdir(parents=True, exist_ok=True)

    source_sha256 = _sha256(source)
    with tempfile.TemporaryDirectory(
        dir=output.parent,
        prefix=".onnx-export-",
    ) as temporary:
        staging = Path(temporary)
        staged_source = staging / "source.pt"
        shutil.copy2(source, staged_source)
        model = dependencies["YOLO"](str(staged_source), task="detect")
        if model.task != "detect":
            raise ValueError(f"source checkpoint task must be detect, got {model.task!r}")
        actual_labels = _ordered_names(model.names)
        if actual_labels != definition.labels:
            raise ValueError(
                "source checkpoint labels do not match the registered model: "
                f"expected {definition.labels!r}, got {actual_labels!r}"
            )

        exported = model.export(
            format="onnx",
            imgsz=args.imgsz,
            batch=1,
            dynamic=False,
            simplify=args.simplify,
            opset=args.opset,
            nms=False,
            half=False,
            device="cpu",
        )
        staged_output = Path(str(exported)).resolve()
        if not staged_output.is_file() or staged_output.suffix.lower() != ".onnx":
            raise RuntimeError(
                f"Ultralytics did not produce the expected ONNX file: {exported!r}"
            )
        try:
            staged_output.relative_to(staging.resolve())
        except ValueError as exc:
            raise RuntimeError(
                f"export escaped the atomic staging directory: {staged_output}"
            ) from exc

        onnx = dependencies["onnx"]
        exported_model = _sanitize_onnx_metadata(
            onnx,
            staged_output,
            definition=definition,
            source_sha256=source_sha256,
        )
        onnx.checker.check_model(exported_model)
        runtime = OnnxYoloModel(
            {
                "weights": str(staged_output),
                "providers": ["CPUExecutionProvider"],
                "imgsz": args.imgsz,
            },
            expected_labels=definition.labels,
        )
        runtime.load()
        runtime.close()
        output_sha256 = _sha256(staged_output)

        staged_manifest = staging / "manifest.json"
        onnxslim = dependencies["onnxslim"]
        manifest_payload = {
            "schema_version": "cmvr.detection-model-manifest/v1",
            "source_model_id": definition.source_model_id,
            "model_id": definition.target_model_id,
            "backend": "onnxruntime-yolov8",
            "source": str(source),
            "source_sha256": source_sha256,
            "artifact": str(output),
            "artifact_sha256": output_sha256,
            "task": "detect",
            "labels": list(definition.labels),
            "input": {
                "layout": "NCHW",
                "shape": [1, 3, args.imgsz, args.imgsz],
                "pixel_format": "RGB",
                "normalization": "uint8 / 255",
                "letterbox_color": [114, 114, 114],
            },
            "export": {
                "format": "onnx",
                "opset": args.opset,
                "batch": 1,
                "dynamic": False,
                "nms": False,
                "half": False,
                "simplify": bool(args.simplify),
            },
            "tools": {
                "torch": dependencies["torch"].__version__,
                "ultralytics": dependencies["ultralytics"].__version__,
                "onnx": dependencies["onnx"].__version__,
                "onnxruntime": dependencies["onnxruntime"].__version__,
                "onnxslim": (
                    None if onnxslim is None else onnxslim.__version__
                ),
            },
        }
        _write_manifest(staged_manifest, manifest_payload)

        # Each file is published with an atomic same-filesystem replacement.
        # The manifest is published last and carries the artifact digest, so a
        # process interrupted between replacements fails closed on hash check
        # instead of accepting a mismatched artifact pair.
        os.replace(staged_output, output)
        os.replace(staged_manifest, manifest)

    return output, manifest


def main(argv: list[str] | None = None) -> int:
    output, manifest = export_model(parse_args(argv))
    print(f"exported ONNX artifact: {output}")
    print(f"wrote model manifest: {manifest}")
    return 0


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