"""빌드/온라인 머신 전용: Hailo yolov8s_pose HEF 다운로드(폐쇄망 번들용). MIT (Hailo Model Zoo).

Mirrors edge/tools/fetch_movenet.py: pure stdlib, no deps, idempotent (re-running with the
file already in place verifies and returns without re-downloading).

AIR GAP (rules/common/build-deployment.md): this runs on the BUILD machine or on an
INTERNET-CONNECTED Pi (deploy/setup_online.sh [6/8]) — never on the air-gapped target. The
offline bundle carries the resulting .hef across the gap (.gitignore ignores *.hef, so it is
a build artifact, never a committed one).

REFUSES TO GUESS THE BOARD. A Hailo-8 HEF does not load on a Hailo-8L and vice versa:
    [HailoRT] [error] HEF format is not compatible with device. Device arch: HAILO8L, HEF arch: HAILO8
    [HailoRT] [error] CHECK_SUCCESS failed with status=HAILO_INVALID_HEF(26)
so the arch is read from the board (`hailortcli fw-control identify`) or must be passed
explicitly with --arch. With neither, this exits non-zero rather than install a coin-flip.

WHY v2.14.0: it is what the official hailo-rpi5-examples download_resources.sh pins, and its
HEF header is version 1 — parsed by every HailoRT from 4.18 through 4.23, so one file works on
both Bookworm (hailo-all 4.20) and Trixie (hailo-all 5.1.1 -> hailort 4.23). Newer model-zoo
releases (v2.15.0+) emit HEF header version 2, which HailoRT 4.18/4.19 reject outright.

CONTRACT (verified 2026-07-16 by decoding the HEF protobuf itself, BOTH arches — matches
edge/adapters/hailo_pose.py and tests/test_hailo_pose.py exactly):
    input_layer1     640x640x3  uint8 RGB; a `normalization1` layer does 0-255 scaling ON CHIP
    output_layer1..9            3 strides x {box DFL (H,W,64), score (H,W,1, sigmoid on chip),
                                keypoints (H,W,51)} @ 80x80, 40x40, 20x20
    NO on-chip NMS (zero nms/hailo_nms markers anywhere in either file) -> the host-side DFL
    decode + NMS in hailo_pose.py is the correct and necessary post-processing.
"""
import argparse
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.request

MZ_VERSION = "v2.14.0"
BASE = "https://hailo-model-zoo.s3.eu-west-2.amazonaws.com/ModelZoo/Compiled"
MODEL = "yolov8s_pose"

# Hailo publishes NO checksum for these objects: .sha256 / .md5 / SHA256SUMS sidecars all
# return HTTP 403 and bucket listing is denied, so there is no upstream manifest to verify
# against. The digests below are therefore NOT vendor-published — they were computed from the
# bytes these URLs served on 2026-07-16 and are pinned so that a later download which differs
# (mirror swap, truncation, tampering) fails loudly instead of installing wrong weights.
# To re-pin deliberately: curl the URL, sha256sum it, update BOTH sha256 and size here.
#
# The S3 ETag is a usable cross-check for hailo8 ONLY (single-part upload -> ETag == MD5
# 9c0542f0009ce952e16f9d32febebe01). hailo8l is a MULTIPART upload (ETag ends "-2"), so its
# ETag is a digest-of-digests, not an MD5 — do not "fix" it to compare equal.
HEFS = {
    "hailo8": {  # 26 TOPS "AI HAT+"
        "sha256": "4f06be93349a21fcc28ee0a746ca84f1c0e17aee99bfd1f59596a477d5160f99",
        "size": 10294713,
    },
    "hailo8l": {  # 13 TOPS "AI Kit"
        "sha256": "7293a6b88a2a7bf9ea07819e92057d67e947300f8d0e87f49d50bb060fb959b1",
        "size": 22658846,
    },
}

DEFAULT_DEST = os.path.join(os.path.dirname(__file__), "..", "..", "models", f"{MODEL}.hef")


def url_for(arch):
    return f"{BASE}/{MZ_VERSION}/{arch}/{MODEL}.hef"


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _arch_from_enum_name(name):
    """Map a hailo_platform DeviceArchitectureTypes member NAME to our HEF key.
    -> 'hailo8' | 'hailo8l' | None. Same first-match 8L-before-8 discipline as the hailortcli
    regex ('HAILO8L' contains 'HAILO8'), so it also folds silicon revisions ('HAILO8_A0' /
    'HAILO8_B0') onto 'hailo8', and returns None for any non-8 family (Hailo-10/15) we have no
    HEF for. Never guesses."""
    n = (name or "").upper()
    if "HAILO8L" in n:
        return "hailo8l"
    if "HAILO8" in n:
        return "hailo8"
    return None


def _arch_via_hailortcli(run):
    """First signal: parse `hailortcli fw-control identify`. -> 'hailo8'|'hailo8l'|None."""
    try:
        p = run(["hailortcli", "fw-control", "identify"],
                capture_output=True, text=True, timeout=30)
    except (OSError, subprocess.SubprocessError):
        return None
    if p.returncode != 0:
        return None
    # HAILO8L must precede HAILO8 in the alternation: regex alternation is first-match, so
    # (HAILO8|HAILO8L) would match the "HAILO8" prefix of "HAILO8L" and mis-detect an AI Kit
    # as an AI HAT+ — the exact HAILO_INVALID_HEF(26) this function exists to prevent.
    m = re.search(r"Device Architecture\s*:\s*(HAILO8L|HAILO8)\b", p.stdout or "", re.I)
    return m.group(1).lower() if m else None


def _arch_via_python_api():  # pragma: no cover (device-only: needs hailo_platform + a board)
    """Second signal: read the arch through the hailo_platform bindings instead of parsing
    hailortcli text. Runs the SAME firmware identify underneath, so it does NOT rescue a board
    whose firmware will not answer — but it (a) works when hailortcli is absent from PATH
    (hailo-all ships the Python bindings independently) and (b) returns a typed enum, so it is
    not defeated by the \\b regex on develop-firmware strings like 'HAILO8_B0'. Never raises:
    any missing symbol/attr or absent board -> None.
    UNVERIFIED: the API names (Device.scan / control.identify().device_architecture) are from
    Hailo's pyhailort and are not runnable in this repo. The broad except makes a wrong name
    safe — it yields None (refuse-to-guess), never a crash or a guess."""
    try:
        from hailo_platform import Device
        ids = Device.scan()
        if not ids:
            return None
        with Device(ids[0]) as dev:
            arch = dev.control.identify().device_architecture
        return _arch_from_enum_name(getattr(arch, "name", str(arch)))
    except Exception:
        return None


def detect_arch(run=subprocess.run, python_api=_arch_via_python_api):
    """Read the arch off the board itself. -> 'hailo8' | 'hailo8l' | None.

    Two independent signals, in order: (1) `hailortcli fw-control identify` text, (2) the
    hailo_platform bindings. Both read the board's ARCHITECTURE, which is reported even when
    the serial/part/product identity fields are unprogrammed — a Hailo-8 commonly shows
    'Device Architecture: HAILO8' with 'Serial Number: <N/A>'. So a serial-less board is NOT
    why this returns None. None still means "no board answered": no hailortcli AND no
    bindings, no /dev/hailo0 (driver installed, not rebooted), or an arch we have no HEF for.
    None is never papered over with a default — an operator who knows the board bypasses
    detection entirely with --arch / HAILO_ARCH."""
    arch = _arch_via_hailortcli(run)
    if arch is not None:
        return arch
    # hailortcli missing/non-zero, or its Device Architecture line absent/unparseable. The
    # in-process bindings survive a missing CLI and the \b regex fragility; still None if no
    # board answers.
    return python_api()


def fetch(arch, dest, force=False, log=print):
    """Download (or verify in place) the yolov8s_pose HEF for `arch`. -> dest path."""
    spec = HEFS[arch]
    url = url_for(arch)

    if os.path.exists(dest) and not force:
        if sha256_of(dest) == spec["sha256"]:
            log(f"{dest} already present and matches the pinned {arch} digest; nothing to do")
            return dest
        # Covers the nastiest case: a previous run left the OTHER arch's HEF (or a truncated
        # one) here. Silently keeping it is how a board ends up loading a HEF it will reject.
        # ASCII only in printed strings: a Windows build machine's console is often cp949/cp1252,
        # where an em-dash raises UnicodeEncodeError. That would crash the tamper-detection path
        # itself, replacing "digest mismatch" with a traceback about a dash.
        log(f"WARNING: {dest} exists but does not match the pinned {arch} digest - replacing it")

    dest_dir = os.path.dirname(dest) or "."
    os.makedirs(dest_dir, exist_ok=True)
    log(f"downloading {MODEL} HEF for {arch} ({spec['size']} bytes, model-zoo {MZ_VERSION})")
    log(f"  {url}")
    log(f"  -> {dest}")

    tmp = None
    try:
        fd, tmp = tempfile.mkstemp(dir=dest_dir, prefix=f".{MODEL}.", suffix=".part")
        os.close(fd)
        # Staged through a temp file in the SAME directory, then os.replace()'d: an
        # interrupted download must never leave a half-written .hef at `dest`, because the
        # installer's next step only tests -f and would happily install the truncation.
        with urllib.request.urlopen(url, timeout=60) as r, open(tmp, "wb") as f:
            shutil.copyfileobj(r, f)

        size = os.path.getsize(tmp)
        digest = sha256_of(tmp)
        if size != spec["size"] or digest != spec["sha256"]:
            raise SystemExit(
                f"ERROR: downloaded {MODEL}.hef for {arch} does not match the pinned digest.\n"
                f"  url     : {url}\n"
                f"  expected: {spec['size']} bytes, sha256 {spec['sha256']}\n"
                f"  got     : {size} bytes, sha256 {digest}\n"
                f"  Refusing to install it. If upstream legitimately re-published this model,\n"
                f"  re-pin HEFS['{arch}'] in {os.path.basename(__file__)} deliberately."
            )
        os.replace(tmp, dest)
        tmp = None
    finally:
        if tmp and os.path.exists(tmp):
            os.remove(tmp)

    log(f"done: {dest} (sha256 verified)")
    return dest


def main(argv=None):
    ap = argparse.ArgumentParser(
        description="Download the Hailo yolov8s_pose HEF for THIS board "
                    "(build/online machine only - never the air-gapped target).")
    ap.add_argument("--arch", choices=sorted(HEFS),
                    help="target arch. Omit to read it from the board via `hailortcli "
                         "fw-control identify`, or set the HAILO_ARCH env var. Pass "
                         "explicitly to cross-fetch on a build machine with no board.")
    ap.add_argument("--dest", default=DEFAULT_DEST, help=f"output path (default: {DEFAULT_DEST})")
    ap.add_argument("--force", action="store_true", help="re-download even if the digest matches")
    args = ap.parse_args(argv)

    # Operator override for a board auto-detect cannot read (e.g. a Hailo-8 whose firmware
    # reports no identity fields, or a Pi where hailortcli is not yet on PATH). Precedence:
    # --arch flag > HAILO_ARCH env > board auto-detect. argparse already validated the flag
    # against HEFS; validate the env value against the same allow-list — an unknown value is a
    # typo, not a request to fall back.
    env_arch = os.environ.get("HAILO_ARCH", "").strip().lower() or None
    if env_arch is not None and env_arch not in HEFS:
        sys.exit(
            f"ERROR: HAILO_ARCH={os.environ['HAILO_ARCH']!r} is not a supported arch.\n"
            f"  Use exactly one of: {', '.join(sorted(HEFS))} (lowercase)."
        )

    detected = detect_arch()
    override = args.arch or env_arch
    arch = override or detected

    if arch is None:
        sys.exit(
            "ERROR: cannot determine the Hailo arch, and refusing to guess.\n"
            "  A Hailo-8 HEF does not load on a Hailo-8L (HAILO_INVALID_HEF(26)), so there is\n"
            "  no safe default here.\n"
            "\n"
            "  On the Pi this usually means the board is not answering yet:\n"
            "    hailortcli fw-control identify     # expect 'Device Architecture: HAILO8'\n"
            "    dpkg -s hailo-all >/dev/null 2>&1 || sudo apt-get install -y hailo-all\n"
            "    ls -l /dev/hailo0                  # absent -> the driver needs a REBOOT\n"
            "\n"
            "  If the board answers but prints no 'Device Architecture' line (some Hailo-8\n"
            "  boards report no serial/identity), name the arch yourself:\n"
            "    HAILO_ARCH=hailo8  python -m edge.tools.fetch_hailo_hef   # 26 TOPS AI HAT+\n"
            "    HAILO_ARCH=hailo8l python -m edge.tools.fetch_hailo_hef   # 13 TOPS AI Kit"
        )

    if override and detected and override != detected:
        # Legitimate when cross-fetching for a fleet; disastrous when it is a typo. Say it
        # loudly, but let the explicit override win — the operator stated an intent.
        src = "--arch" if args.arch else "HAILO_ARCH"
        print(f"WARNING: {src} {override} overrides the detected board ({detected}). "
              f"This HEF will NOT load on this machine's board.", file=sys.stderr)

    fetch(arch, os.path.abspath(args.dest), force=args.force)


if __name__ == "__main__":
    main()