"""ov_models.py — OpenVINO runtime wrappers for Marquee.

Models used (all CPU, all from Open Model Zoo):
  - face-detection-retail-0004      : SSD face detector (for star card crops)
  - person-detection-retail-0013    : SSD person detector (for tracking)

Dropped from original:
  - landmarks-regression-retail-0009  (only needed for re-id alignment)
  - face-reidentification-retail-0095 (replaced by IoU tracker in faces.py)

[OpenVINO 101] Models ship as .xml (graph) + .bin (weights). Compile once
for CPU, infer many times. Never touches GPU — no ZeroGPU contention.
"""

import urllib.request
from pathlib import Path

import cv2
import numpy as np
import openvino as ov

MODELS_DIR = Path(__file__).parent / "models"
PRECISION   = "FP16"

# Models we actually need now
MODEL_NAMES = [
    "face-detection-retail-0004",    # face crops for star card UI
    "person-detection-retail-0013",  # person bboxes for IoU tracking
]

_OMZ_BASE = ("https://storage.openvinotoolkit.org/repositories/open_model_zoo/"
             "2023.0/models_bin/1")

_core = ov.Core()


def download_models():
    for name in MODEL_NAMES:
        dst  = MODELS_DIR / name
        xml  = dst / f"{name}.xml"
        binf = dst / f"{name}.bin"
        if xml.exists() and binf.exists():
            continue
        dst.mkdir(parents=True, exist_ok=True)
        for target in (xml, binf):
            url  = f"{_OMZ_BASE}/{name}/{PRECISION}/{target.name}"
            part = target.with_suffix(target.suffix + ".part")
            print(f"[ov] downloading {target.name} ...")
            urllib.request.urlretrieve(url, part)
            part.replace(target)


def _find_xml(name: str) -> str:
    hits = list(MODELS_DIR.rglob(f"{name}.xml"))
    if not hits:
        raise FileNotFoundError(
            f"{name}.xml not found under {MODELS_DIR}. Run download_models().")
    return str(hits[0])


class OVModel:
    """Compile once, infer many. Handles NCHW resize + BGR."""

    def __init__(self, name: str):
        model          = _core.read_model(_find_xml(name))
        self.compiled  = _core.compile_model(model, "CPU")
        self.input     = self.compiled.input(0)
        _, _, self.h, self.w = self.input.shape

    def _prep(self, bgr: np.ndarray) -> np.ndarray:
        img = cv2.resize(bgr, (self.w, self.h))
        return img.transpose(2, 0, 1)[None].astype(np.float32)  # HWC->NCHW

    def infer(self, bgr: np.ndarray):
        out = self.compiled({self.input: self._prep(bgr)})
        return out[self.compiled.output(0)]