# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Device tests of the Python API (needs one Blackhole chip as device 0, the released weights in the HF cache).

1. ``test_api_equals_server``: boots the HTTP server (``uvicorn hamer.server.app:app`` subprocess, the ``tt-model.yaml``
   serve env), takes ``/predict`` for the 8 demo hands of ``media/``, stops it, then
   builds ``HamerModel.from_pretrained(device_id=0)`` and requires ``hand.to_dict()`` == the server response (all keys
   except ``timing_ms`` / ``trace_active``) for every hand: the same numbers at the server's rounding.
2. ``test_api_call_time``: warm repeated calls, fresh crop every call (the 8 demo crops in turn), median / min of
   - ``model.tt(u8)``            the bare ``TtHamer.__call__`` (= the published "model call" / served call),
   - ``model.run_crop(u8)``      the same call through the API (lock + inference_mode),
   - ``model(frame, bbox)``      the full API call (host crop + resize, model call, post-processing),
   and requires ``run_crop`` not slower than the bare call (+3 %).  Results go to ``.work/api_timing.json``.

    source tools/chipenv.sh $CHIP models/hamer-p150/.venv
    cd models/hamer-p150 && timeout -s INT 1200 python -m pytest code/tests/test_api_device.py -q -s
"""
from __future__ import annotations

import base64
import json
import os
import statistics
import time
from pathlib import Path

import pytest

if os.environ.get("TT_VISIBLE_DEVICES") == "none" or os.environ.get("HAMER_USE_TT") == "0":
    pytest.skip("device test (TT_VISIBLE_DEVICES=none / HAMER_USE_TT=0)", allow_module_level=True)

import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402

REPO = Path(__file__).resolve().parents[2]
MEDIA = REPO / "media"
PUBLISHED_MODEL_CALL_MS = (3.480, 3.499)  # OPT_REPORT.md round 15 headline, chip 5: served TtHamer.__call__ median
ITERS = int(os.environ.get("HAMER_API_ITERS", "200"))


def _hands():
    boxes = json.loads((MEDIA / "bboxes.json").read_text())
    return [(name, e["bbox"], bool(e.get("is_right", 1))) for name in sorted(boxes) for e in boxes[name]]


def _serve_env():
    from hamer.served_config import SERVED_DISPATCH, SERVED_ENV, WEIGHTS_REVISION

    env = dict(SERVED_ENV, HAMER_DISPATCH=SERVED_DISPATCH, HAMER_WARMUP_ITERS="3", TT_MESH_SHAPE="1x1",
               TT_DEVICE_ID="0", TT_WEIGHTS_REVISION=WEIGHTS_REVISION)
    for k, v in env.items():
        os.environ[k] = v
    return env


def _free_port() -> int:
    import socket

    with socket.socket() as so:
        so.bind(("127.0.0.1", 0))
        return so.getsockname()[1]


@pytest.fixture(scope="module")
def server_outputs():
    """The HTTP server as deployed (uvicorn subprocess, serve env), queried for every demo hand, then stopped."""
    import signal
    import subprocess
    import sys
    import urllib.request

    pytest.importorskip("fastapi")
    pytest.importorskip("uvicorn")
    env = dict(os.environ, **_serve_env())
    port = _free_port()
    log = REPO / ".work" / "api_test_server.log"
    log.parent.mkdir(exist_ok=True)
    with open(log, "w") as fh:
        proc = subprocess.Popen([sys.executable, "-m", "uvicorn", "--host", "127.0.0.1", "--port", str(port),
                                 "--lifespan", "on", "hamer.server.app:app"], env=env, stdout=fh, stderr=subprocess.STDOUT)
    url = f"http://127.0.0.1:{port}"

    def get(path, body=None):
        req = urllib.request.Request(url + path, data=None if body is None else json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        return json.loads(urllib.request.urlopen(req, timeout=60).read())

    try:
        for _ in range(600):
            if proc.poll() is not None:
                raise RuntimeError(f"server exited ({proc.returncode}); see {log}")
            if "Application startup complete" in log.read_text():
                break
            time.sleep(1)
        else:
            raise RuntimeError(f"server did not become ready; see {log}")
        outs = {}
        for name, bbox, right in _hands():
            img = base64.b64encode((MEDIA / name).read_bytes()).decode()
            out = get("/predict", {"image": img, "bbox": bbox, "is_right": right})
            assert out["trace_active"] is True
            out.pop("timing_ms")
            out.pop("trace_active")
            outs[(name, tuple(bbox), right)] = out
        info = get("/info")
    finally:
        proc.send_signal(signal.SIGINT)
        try:
            proc.wait(timeout=120)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    return outs, info


@pytest.fixture(scope="module")
def model(server_outputs):
    from hamer import HamerModel

    m = HamerModel.from_pretrained(device_id=0, verbose=True)
    yield m
    m.close()


def test_api_equals_server(server_outputs, model):
    outs, info = server_outputs
    assert model.trace_active and model.info["grid"] == "12x10" and model.info["dispatch"] == "eth"
    print(f"\n[api] info: {json.dumps({k: v for k, v in model.info.items() if k != 'fused'})}")
    print(f"[api] server warm-up: {info['warmup']}")
    n = 0
    for (name, bbox, right), want in outs.items():
        hand = model(MEDIA / name, bbox=list(bbox), is_right=right)
        got = hand.to_dict()
        diff = [k for k in want if got.get(k) != want[k]]
        assert not diff and set(got) == set(want), f"{name} {bbox}: keys differ from the server: {diff}"
        n += 1
    # the same hands through the list form (one decode per frame)
    by_frame = {}
    for (name, bbox, right) in outs:
        by_frame.setdefault(name, []).append((list(bbox), right))
    for name, items in by_frame.items():
        hands = model(MEDIA / name, bbox=[b for b, _ in items], is_right=[r for _, r in items])
        for (b, r), h in zip(items, hands):
            assert h.to_dict() == outs[(name, tuple(b), r)]
    print(f"[api] {n} demo hands: HamerModel output == server /predict output (all keys, server rounding)")


def _stats(v):
    return {"median_ms": round(statistics.median(v), 4), "min_ms": round(min(v), 4), "n": len(v)}


def test_api_call_time(model):
    from hamer.preprocess import preprocess_pil

    hands = _hands()
    frames = {name: Image.open(MEDIA / name).convert("RGB") for name, _, _ in hands}
    crops = [preprocess_pil(frames[n], b, is_right=r, raw=True)[0] for n, b, r in hands]
    tt = model.tt

    def loop(fn, iters=ITERS):
        ts = []
        for i in range(iters):
            t0 = time.perf_counter()
            fn(i)
            ts.append((time.perf_counter() - t0) * 1e3)
        return ts

    with torch.inference_mode():
        loop(lambda i: tt(crops[i % len(crops)]), 20)  # warm
    res = {}
    # alternate the two measurements (3 rounds each) so host-load drift hits both equally
    bare, api = [], []
    for _ in range(3):
        with torch.inference_mode():
            bare += loop(lambda i: tt(crops[i % len(crops)]))
        api += loop(lambda i: model.run_crop(crops[i % len(crops)]))
    res["bare_TtHamer_call"] = _stats(bare)
    res["api_run_crop"] = _stats(api)
    full, inner = [], []
    for i in range(ITERS):
        n, b, r = hands[i % len(hands)]
        t0 = time.perf_counter()
        h = model(frames[n], bbox=b, is_right=r)
        full.append((time.perf_counter() - t0) * 1e3)
        inner.append(h.timing_ms["model"])
    res["api_call_full_frame"] = _stats(full)
    res["api_call_full_frame_model_stage"] = _stats(inner)
    res["published_model_call_ms"] = PUBLISHED_MODEL_CALL_MS
    res["grid"], res["dispatch"] = model.info["grid"], model.info["dispatch"]
    res["chip"] = os.environ.get("TT_CHIP")
    res["host_load"] = [round(v, 1) for v in os.getloadavg()]
    print("\n[api] timing " + json.dumps(res))
    out = REPO / ".work"
    out.mkdir(exist_ok=True)
    (out / "api_timing.json").write_text(json.dumps(res, indent=2))
    assert res["api_run_crop"]["median_ms"] <= res["bare_TtHamer_call"]["median_ms"] * 1.03