#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
export_onnx.py
- Export YOLOv8-(seg) 'best.pt' to ONNX
- Assemble all required companion files into a deploy bundle:
    model.onnx, names.json, label_map.json, dataset.yaml, train.yaml,
    metrics.json, summary_eval.json, calibration_temp.json, val_predictions.json,
    postprocess_config.json, model_card.json, CHECKSUMS.sha256
- Optional: dynamic axes, simplify, half/fp16, quantize
- Optional: ZIP the bundle

Example:
python ml/scripts/export_onnx.py \
    --artifacts_dir artifacts/v2025YYYYMMDD-HHMMSS \
    --out_dir deploy/v2025YYYYMMDD-HHMMSS \
    --dynamic --simplify --opset 12 --zip
"""

import argparse
import json
import os
import shutil
import sys
import hashlib
from pathlib import Path
from datetime import datetime
from typing import Optional, List

# --- Optional imports (used if available) ---
try:
    import onnx
except Exception:
    onnx = None

try:
    from onnxruntime.quantization import quantize_dynamic, QuantType
except Exception:
    quantize_dynamic, QuantType = None, None

try:
    from ultralytics import YOLO
except Exception as e:
    print("Ultralytics YOLOv8 required. Please `pip install ultralytics`.", file=sys.stderr)
    raise

try:
    import yaml
except Exception:
    print("Please `pip install pyyaml`.", file=sys.stderr)
    raise


# ---------------------------- IO/Utils ---------------------------------------

def read_json(p: Path):
    if p.exists():
        return json.loads(p.read_text(encoding="utf-8"))
    return None

def write_json(p: Path, obj):
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding="utf-8")

def read_yaml(p: Path):
    return yaml.safe_load(p.read_text(encoding="utf-8")) if p.exists() else None

def copy_if_exists(src: Path, dst: Path):
    if src.exists():
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)

def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 16), b""):
            h.update(chunk)
    return h.hexdigest()

def _glob_any(paths: List[Path], pattern: str) -> List[Path]:
    out = []
    for p in paths:
        try:
            out.extend(list(p.glob(pattern)))
        except Exception:
            pass
    return out


# ---------------------------- Names / Metadata --------------------------------

def build_names_json(artifacts_dir: Path) -> list:
    # Prefer names.json; otherwise, read from dataset.yaml
    names = read_json(artifacts_dir / "names.json")
    if names is None:
        ds = read_yaml(artifacts_dir / "dataset.yaml") or {}
        names = ds.get("names", [])
    # normalisieren auf Liste
    if isinstance(names, dict):
    # Keys can be strings like "0","1",...
        keys = [int(k) for k in names.keys()] if names else []
        arr = [None] * (max(keys) + 1 if keys else 0)
        for k, v in names.items():
            arr[int(k)] = v
        names = arr
    if not isinstance(names, list):
        names = list(names) if names is not None else []
    names = [("" if n is None else str(n)) for n in names]
    return names


# ---------------------------- ONNX Export/Verify/Quant ------------------------

def onnx_export(
    best_pt: Path,
    save_dir: Path,
    dynamic: bool,
    opset: int,
    simplify: bool,
    half: bool
) -> Path:
    """
    Run Ultralytics export and copy the first found *.onnx to save_dir/model.onnx.
    Search robustly in:
        - model.trainer.save_dir
        - current working dir
        - directory of best_pt
    """
    save_dir.mkdir(parents=True, exist_ok=True)

    # Sanity for opset
    if opset < 11:
        print("[WARN] opset < 11 is unusual; forcing 11.", file=sys.stderr)
        opset = 11

    model = YOLO(str(best_pt))
    try:
        model.export(format="onnx", dynamic=dynamic, opset=opset, simplify=simplify, half=half)
    except Exception as e:
        print(f"[ERROR] ONNX export failed: {e}", file=sys.stderr)
        raise

    # Search candidate locations
    candidates = []
    try:
        candidates.append(Path(model.trainer.save_dir))
    except Exception:
        pass
    candidates.extend([Path.cwd(), best_pt.parent])

    onnx_files = _glob_any(candidates, "*.onnx")
    if not onnx_files:
        raise FileNotFoundError("No *.onnx found after export (trainer folder, CWD, artifacts).")

    out_onnx = save_dir / "model.onnx"
    shutil.copy2(onnx_files[0], out_onnx)
    return out_onnx

def onnx_verify(onnx_path: Path):
    if onnx is None:
        print("[WARN] onnx not installed; skipping checker.")
        return
    try:
        m = onnx.load(str(onnx_path))
        onnx.checker.check_model(m)
        print("[INFO] ONNX Model check OK.")
    except Exception as e:
        print(f"[WARN] ONNX checker warning/error: {e}", file=sys.stderr)

def onnx_quantize_dynamic(in_path: Path, out_path: Path) -> Optional[Path]:
    if quantize_dynamic is None:
        print("[WARN] onnxruntime-quantization not available; skipping quantize.")
        return None
    try:
        quantize_dynamic(
            model_input=str(in_path),
            model_output=str(out_path),
            weight_type=QuantType.QInt8
        )
        return out_path
    except Exception as e:
        print(f"[WARN] Dynamic quantization failed: {e}", file=sys.stderr)
        return None


# ---------------------------- Serving/Bundle Metadaten ------------------------

def compose_postprocess_config(artifacts_dir: Path) -> dict:
    # Default thresholds; can be overridden in serving
    cfg = {
        "confidence_threshold": 0.25,
        "iou_threshold": 0.50,
        "max_detections": 300,
        "calibration": None
    }
    cal = read_json(artifacts_dir / "calibration_temp.json")
    if cal and "temperature" in cal:
        cfg["calibration"] = {
            "method": cal.get("method", "temperature_scaling"),
            "temperature": float(cal["temperature"])
        }
    return cfg

def compose_model_card(artifacts_dir: Path, names: list, onnx_path: Path, onnx_q_path: Optional[Path]) -> dict:
    ds = read_yaml(artifacts_dir / "dataset.yaml") or {}
    metrics = read_json(artifacts_dir / "metrics.json") or {}
    evalsum = read_json(artifacts_dir / "summary_eval.json") or {}
    label_map = read_json(artifacts_dir / "label_map.json") or {}

    return {
        "model_card_version": "1.0",
        "created_utc": datetime.utcnow().isoformat() + "Z",
        "artifacts_dir": str(artifacts_dir.resolve()),
        "onnx": {
            "path": str(onnx_path.name),
            "quantized_path": (onnx_q_path.name if onnx_q_path else None)
        },
        "dataset": {
            "path": ds.get("path"),
            "splits": {
                "train": ds.get("train"),
                "val": ds.get("val"),
                "test": ds.get("test")
            },
            "classes": names
        },
        "metrics": {
            "ultralytics": metrics,
            "evaluation": evalsum
        },
        "mapping": {
            "label_map": label_map
        },
    "notes": "YOLOv8-seg export; calibration via temperature scaling (if present)."
    }

def build_checksums(out_dir: Path):
    lines = []
    for p in sorted(out_dir.rglob("*")):
        if p.is_file():
            if p.name.endswith(".zip"):  # don't include the ZIP itself in checksums
                continue
            digest = sha256_file(p)
            rel = p.relative_to(out_dir)
            lines.append(f"{digest}  {rel.as_posix()}")
    (out_dir / "CHECKSUMS.sha256").write_text("\n".join(lines) + "\n", encoding="utf-8")

def make_zip(out_dir: Path):
    base = out_dir.parent / (out_dir.name + ".zip")
    if base.exists():
        base.unlink()
    shutil.make_archive(str(base.with_suffix("")), "zip", root_dir=str(out_dir))
    print(f"[INFO] ZIP: {base}")


# ---------------------------- Main -------------------------------------------

def main():
    ap = argparse.ArgumentParser(description="Export YOLOv8 best.pt to ONNX and assemble deploy bundle.")
    ap.add_argument("--artifacts_dir", required=True, type=Path, help="Path to artifacts/<version>")
    ap.add_argument("--out_dir", required=True, type=Path, help="Deploy target directory (prefer empty/non-existent)")
    ap.add_argument("--opset", type=int, default=12)
    ap.add_argument("--dynamic", action="store_true")
    ap.add_argument("--simplify", action="store_true")
    ap.add_argument("--half", action="store_true", help="fp16 export (if supported)")
    ap.add_argument("--quantize", action="store_true", help="additional dynamic INT8 quantization")
    ap.add_argument("--zip", action="store_true", help="create ZIP of the deploy directory")
    args = ap.parse_args()

    artifacts_dir = args.artifacts_dir.resolve()
    out_dir = args.out_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    # Accept several common locations
    candidates = [
        artifacts_dir / "best.pt",
        artifacts_dir / "weights" / "best.pt",
    ]

    best_pt = next((p for p in candidates if p.exists()), None)

    # As a last resort, search typical Ultralytics runs dirs
    if best_pt is None:
        from glob import glob
        runs = []
        runs += glob(str(artifacts_dir / "weights" / "*.pt"))
        runs += glob(str(artifacts_dir.parent / "runs-seg" / "*" / "weights" / "best.pt"))
        if runs:
            best_pt = Path(sorted(runs)[-1])

    if best_pt is None or not best_pt.exists():
        print(f"[ERROR] best.pt not found in artifacts. Checked: {', '.join(map(str, candidates))}", file=sys.stderr)
        sys.exit(1)


    # 1) ONNX Export
    onnx_path = onnx_export(
        best_pt=best_pt,
        save_dir=out_dir,
        dynamic=bool(args.dynamic),
        opset=int(args.opset),
        simplify=bool(args.simplify),
        half=bool(args.half)
    )
    onnx_verify(onnx_path)

    # 2) Optional: dynamic quantization
    onnx_q_path = None
    if args.quantize:
        q_path = out_dir / "model.int8.onnx"
        q = onnx_quantize_dynamic(onnx_path, q_path)
        if q:
            onnx_q_path = q

    # 3) Collect companion files
    names = build_names_json(artifacts_dir)
    write_json(out_dir / "names.json", names)

    copy_if_exists(artifacts_dir / "label_map.json", out_dir / "label_map.json")
    copy_if_exists(artifacts_dir / "dataset.yaml", out_dir / "dataset.yaml")
    copy_if_exists(artifacts_dir / "train.yaml", out_dir / "train.yaml")
    copy_if_exists(artifacts_dir / "metrics.json", out_dir / "metrics.json")
    copy_if_exists(artifacts_dir / "summary_eval.json", out_dir / "summary_eval.json")
    copy_if_exists(artifacts_dir / "val_predictions.json", out_dir / "val_predictions.json")
    copy_if_exists(artifacts_dir / "calibration_temp.json", out_dir / "calibration_temp.json")

    # 4) Postprocess-Config (Thresholds + Calibration)
    post_cfg = compose_postprocess_config(artifacts_dir)
    write_json(out_dir / "postprocess_config.json", post_cfg)

    # 5) Model Card
    model_card = compose_model_card(artifacts_dir, names, onnx_path, onnx_q_path)
    write_json(out_dir / "model_card.json", model_card)

    # 6) Checksums
    build_checksums(out_dir)

    # 7) ZIP (optional)
    if args.zip:
        make_zip(out_dir)

    print("\n=== Export complete ===")
    print(f"Deploy bundle: {out_dir}")
    if args.zip:
        print(f"ZIP: {out_dir.parent / (out_dir.name + '.zip')}")


if __name__ == "__main__":
    main()