#!/usr/bin/env python3
"""Export an Ultralytics checkpoint and its metadata for one Rockchip chip."""

import argparse
import os
from pathlib import Path
import shutil
import tempfile


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model", help="Checkpoint path or official model name, e.g. yolo11n-seg.pt")
    parser.add_argument("--chip", required=True, choices=("rk3576", "rk3588"))
    parser.add_argument("--output", required=True, type=Path, help="New output directory")
    parser.add_argument("--imgsz", type=int, help="Square input size; defaults depend on the task")
    parser.add_argument("--precision", choices=("fp16", "int8", "mixed"), default="fp16")
    parser.add_argument("--data", type=Path, help="Calibration YAML, image list, or directory; detection INT8 needs a dataset YAML")
    parser.add_argument("--calibration-limit", type=int, default=100,
                        help="Maximum calibration images for mixed precision (default: 100)")
    args = parser.parse_args()

    destination = args.output.resolve()
    if destination.exists():
        parser.error("Output already exists; choose a new directory")
    if args.precision != "fp16" and not args.data:
        parser.error("Quantization requires --data with representative calibration images")
    if args.imgsz and (args.imgsz < 32 or args.imgsz > 4096 or args.imgsz % 32):
        parser.error("--imgsz must be a multiple of 32 between 32 and 4096")
    if not 1 <= args.calibration_limit <= 10000:
        parser.error("--calibration-limit must be between 1 and 10000")
    checkpoint = Path(args.model).expanduser()
    checkpoint = checkpoint.resolve() if checkpoint.is_file() else None
    dataset = args.data.resolve() if args.data else None

    # Conversion runs in an existing toolkit environment; do not install packages for the user.
    os.environ["YOLO_AUTOINSTALL"] = "false"
    from ultralytics import YOLO
    from ultralytics.utils import YAML

    previous_directory = Path.cwd()
    try:
        # Exporters write intermediate files beside the checkpoint. Keep the original untouched.
        with tempfile.TemporaryDirectory(prefix="nvr-export-") as scratch:
            os.chdir(scratch)
            if checkpoint:
                shutil.copy2(checkpoint, checkpoint.name)
            model = YOLO(checkpoint.name if checkpoint else args.model)
            size = args.imgsz or {"classify": 224, "semantic": 512, "depth": 384}.get(model.task, 640)
            optimized = (args.precision == "mixed" or model.task in {"semantic", "depth"}
                         or (args.precision == "int8" and model.task != "detect"))
            if optimized:
                from rknn_export import export_optimized
                binary, metadata = export_optimized(
                    model, args.chip, size, dataset, args.calibration_limit,
                    args.precision)
            else:
                settings = dict(format="rknn", name=args.chip, imgsz=size, batch=1,
                                quantize=8 if args.precision == "int8" else 16,
                                device="cpu", simplify=True, opset=19)
                if dataset:
                    settings["data"] = str(dataset)
                exported = Path(model.export(**settings))
                binaries = list(exported.glob("*.rknn"))
                if len(binaries) != 1:
                    raise RuntimeError("Expected one RKNN model in the export directory")
                binary = binaries[0]
                metadata = YAML.load(exported / "metadata.yaml")
            metadata["precision"] = args.precision
            metadata["target_platform"] = args.chip
            metadata["coordinates"] = ("normalized" if args.precision == "int8"
                                       and model.task == "detect" else "pixels")
            metadata["preprocess"] = "center_crop" if model.task == "classify" else "letterbox"
            destination.mkdir(parents=True)
            shutil.copy2(binary, destination / "model.rknn")
            YAML.save(destination / "metadata.yaml", metadata)
    finally:
        os.chdir(previous_directory)
    print(f"Exported {model.task} for {args.chip}: {destination}")


if __name__ == "__main__":
    main()