# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Host-only tests of the Python API (``hamer.HamerModel``): no device, no ttnn tensors, no weight download.

The API runs with ``backend="cpu"`` on a seeded random-weight reference (and a synthetic MANO pkl with the real array
shapes).  Its outputs are compared with the HTTP server's own ``/predict`` handler (``hamer.server.app.predict``,
called in-process with the same reference as its model): the published API must return the server's numbers.

    cd models/hamer-p150 && TT_VISIBLE_DEVICES=none HAMER_USE_TT=0 .venv/bin/python -m pytest code/tests/test_api_host.py -q
"""
from __future__ import annotations

import base64
import os
import pickle
import subprocess
import sys
from pathlib import Path

os.environ.setdefault("HAMER_USE_TT", "0")  # never open a chip from a host test

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

import hamer  # noqa: E402
from hamer import HamerModel, HandPrediction  # noqa: E402
from hamer.api import _check_bbox, to_pil  # noqa: E402
from hamer.reference import build_reference  # noqa: E402

REPO = Path(__file__).resolve().parents[2]
MEDIA = REPO / "media"
SAMPLE = MEDIA / "sample1.jpg"
BOX_R, BOX_L = [20, 170, 295, 260], [130, 340, 320, 510]

torch.set_grad_enabled(False)


@pytest.fixture(scope="module")
def ref():
    return build_reference(seed=0)


@pytest.fixture(scope="module")
def mano_pkl(tmp_path_factory):
    g = np.random.default_rng(0)
    w = g.random((778, 16))
    w /= w.sum(1, keepdims=True)
    jr = np.zeros((16, 778))
    for j in range(16):
        jr[j, g.choice(778, 20, replace=False)] = 1 / 20
    d = {"v_template": g.random((778, 3)) * 0.1, "shapedirs": g.normal(0, 1e-3, (778, 3, 10)),
         "posedirs": g.normal(0, 1e-3, (778, 3, 135)), "J_regressor": jr, "weights": w,
         "kintree_table": np.array([[4294967295, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14], list(range(16))]),
         "f": g.integers(0, 778, (1538, 3))}
    p = tmp_path_factory.mktemp("mano") / "MANO_RIGHT.pkl"
    with open(p, "wb") as f:
        pickle.dump(d, f)
    return p


@pytest.fixture(scope="module")
def mano(mano_pkl):
    from hamer.mano import MANO

    return MANO(mano_pkl).eval()


@pytest.fixture(scope="module")
def cpu_model(ref):
    m = HamerModel.from_reference(ref, backend="cpu")
    yield m
    m.close()


@pytest.fixture(scope="module")
def cpu_model_mano(ref, mano):
    m = HamerModel.from_reference(ref, backend="cpu", mano=mano)
    yield m
    m.close()


# --------------------------------------------------------------------------- #
# Package surface
# --------------------------------------------------------------------------- #
def test_import_does_not_load_ttnn():
    code = "import sys, hamer; from hamer import HamerModel, HandPrediction, Hamer; assert 'ttnn' not in sys.modules"
    env = {k: v for k, v in os.environ.items() if k not in ("TT_METAL_HOME", "TT_METAL_RUNTIME_ROOT")}
    r = subprocess.run([sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=120)
    assert r.returncode == 0, r.stderr


def test_public_names():
    for n in ("HamerModel", "HandPrediction", "open_device", "load_reference", "apply_served_config",
              "Hamer", "HamerConfig", "build_reference", "sample_input", "build_from_reference",
              "load_hamer_checkpoint"):
        assert n in hamer.__all__ and getattr(hamer, n) is not None
    with pytest.raises(AttributeError):
        hamer.does_not_exist  # noqa: B018


def test_package_data_present():
    pkg = Path(hamer.__file__).parent
    assert (pkg / "kernels" / "hgemv" / "compute.cpp").is_file()
    assert (pkg / "native" / "resize.c").is_file()


def test_no_machine_paths_in_runtime_code():
    pkg = Path(hamer.__file__).parent
    bad = [str(p) for p in pkg.rglob("*.py") if "/home/" in p.read_text()]
    assert not bad, bad


def test_served_config_matches_manifest():
    yaml = pytest.importorskip("yaml")
    man = REPO / "tt-model.yaml"
    if not man.is_file():
        pytest.skip("tt-model.yaml not shipped here")
    from hamer import served_config as sc

    env = dict(yaml.safe_load(man.read_text())["serve"]["env"])
    assert env.pop("TT_WEIGHTS_REVISION") == sc.WEIGHTS_REVISION
    assert env.pop("HAMER_DISPATCH") == sc.SERVED_DISPATCH
    env.pop("TT_METAL_VISIBLE_DEVICES")
    env.pop("HAMER_WARMUP_ITERS")
    assert {k: str(v) for k, v in env.items()} == sc.SERVED_ENV
    assert sc.SERVED_WORKER_L1_SIZE == 1461376 - int(sc.SERVED_ENV["HAMER_KCFG_EXTRA_KB"]) * 1024


def test_apply_served_config_keeps_caller_values(monkeypatch):
    from hamer.served_config import SERVED_ENV

    for k in SERVED_ENV:
        monkeypatch.delenv(k, raising=False)
    monkeypatch.setenv("TT_FUSED_SDPA", "default")
    differ = hamer.apply_served_config()
    assert differ == {"TT_FUSED_SDPA": ("default", "32x192")}
    assert os.environ["TT_FUSED_SDPA"] == "default" and os.environ["TT_FUSED_INPUT"] == "u8"


# --------------------------------------------------------------------------- #
# Inputs
# --------------------------------------------------------------------------- #
def test_to_pil_accepts_common_types():
    want = np.asarray(Image.open(SAMPLE).convert("RGB"))
    cands = [str(SAMPLE), SAMPLE, Image.open(SAMPLE), want, want.astype(np.float32) / 255.0,
             torch.from_numpy(want.copy()), torch.from_numpy(want.copy()).permute(2, 0, 1),
             torch.from_numpy(want.copy()).permute(2, 0, 1)[None], Image.open(SAMPLE).convert("RGBA")]
    for c in cands:
        got = np.asarray(to_pil(c))
        assert got.shape == want.shape and np.array_equal(got, want), type(c)
    with pytest.raises(ValueError):
        to_pil(np.zeros((4, 4, 3), np.float32) + 2.0)
    with pytest.raises(TypeError):
        to_pil(123)


def test_bbox_rules():
    img = Image.new("RGB", (100, 80))
    assert _check_bbox(None, img, 2.0)[1] == "center_square"
    assert _check_bbox([1, 2, 30, 40], img, 2.0) == ((1.0, 2.0, 30.0, 40.0), "request")
    for bad in ([5, 5, 5, 9], [0, 0, float("nan"), 3], [200, 200, 300, 300], [0, 0, 0.5, 0.5], "abcd"):
        with pytest.raises(ValueError):
            _check_bbox(bad, img, 2.0)


# --------------------------------------------------------------------------- #
# Outputs vs the HTTP server's handler (same reference model)
# --------------------------------------------------------------------------- #
def _server_predict(model_obj, mano_obj, **req):
    pytest.importorskip("fastapi")
    import hamer.server.app as app

    app.STATE.clear()
    app.STATE.update(model=model_obj, mano=mano_obj, ready=True)
    try:
        body = app.PredictRequest(image=base64.b64encode(SAMPLE.read_bytes()).decode(), **req)
        out = app.predict(body)
    finally:
        app.STATE.clear()
    out.pop("timing_ms")
    out.pop("trace_active")
    return out


@pytest.mark.parametrize("box,right", [(BOX_R, True), (BOX_L, False), (None, True)])
def test_matches_server_handler(ref, cpu_model, box, right):
    want = _server_predict(ref, None, bbox=box, is_right=right)
    hand = cpu_model(SAMPLE, bbox=box, is_right=right)
    assert isinstance(hand, HandPrediction)
    assert hand.to_dict() == want
    assert hand.rotmats.shape == (16, 3, 3) and hand.betas.shape == (10,) and hand.cam.shape == (3,)
    assert hand.cam_t_full.shape == (3,) and hand.vertices is None and not hand.mano_available
    assert set(hand.timing_ms) == {"preprocess", "model", "mano", "total"}


@pytest.mark.parametrize("return_vertices", [True, False])
def test_matches_server_handler_with_mano(ref, mano, cpu_model_mano, return_vertices):
    want = _server_predict(ref, mano, bbox=BOX_L, is_right=False, return_vertices=return_vertices)
    hand = cpu_model_mano(SAMPLE, bbox=BOX_L, is_right=False, return_vertices=return_vertices)
    assert hand.to_dict() == want
    assert hand.joints.shape == (21, 3) and hand.keypoints_2d.shape == (21, 2)
    assert (hand.vertices is not None) == return_vertices
    assert cpu_model_mano.faces.shape == (1538, 3)


def test_several_hands_and_input_types(cpu_model):
    frame = np.asarray(Image.open(SAMPLE).convert("RGB"))
    hands = cpu_model(frame, bbox=np.array([BOX_R, BOX_L]), is_right=[True, False])
    assert isinstance(hands, list) and len(hands) == 2
    single = [cpu_model(SAMPLE, bbox=b, is_right=r) for b, r in ((BOX_R, True), (BOX_L, False))]
    for a, b in zip(hands, single):
        assert a.to_dict(rounded=False) == b.to_dict(rounded=False)
    with pytest.raises(ValueError):
        cpu_model(frame, bbox=[BOX_R, BOX_L], is_right=[True])
    with pytest.raises(ValueError):
        cpu_model(frame, bbox=BOX_R, rescale_factor=0.2)


def test_mano_params_layout(cpu_model):
    h = cpu_model(SAMPLE, bbox=BOX_R)
    mp = h.mano_params
    assert mp["global_orient"].shape == (1, 3, 3) and mp["hand_pose"].shape == (15, 3, 3)
    assert np.array_equal(np.concatenate([mp["global_orient"], mp["hand_pose"]]), h.rotmats)
    r = h.rotmats.astype(np.float64)
    assert np.allclose(r @ r.transpose(0, 2, 1), np.eye(3), atol=1e-5)  # proper rotations


def test_forward_crops_upstream_layout(ref, cpu_model_mano):
    from hamer.preprocess import crop_bbox, normalize_raw, preprocess_pil

    img = Image.open(SAMPLE).convert("RGB")
    u8, _ = preprocess_pil(img, BOX_R, raw=True)                         # (1, 3, 256, 192)
    full, _ = crop_bbox(img, BOX_R)                                      # 256 x 256 upstream-style crop
    u256 = torch.from_numpy(np.asarray(full).copy()).permute(2, 0, 1)[None]
    res = cpu_model_mano.forward_crops(torch.cat([u8, u8]))
    assert res["pred_cam"].shape == (2, 3) and res["pred_cam_t"].shape == (2, 3)
    assert res["pred_mano_params"]["global_orient"].shape == (2, 1, 3, 3)
    assert res["pred_mano_params"]["hand_pose"].shape == (2, 15, 3, 3)
    assert res["pred_mano_params"]["betas"].shape == (2, 10)
    assert res["pred_vertices"].shape == (2, 778, 3) and res["pred_keypoints_3d"].shape == (2, 21, 3)
    res_n = cpu_model_mano.forward_crops(normalize_raw(u8))              # uint8 == its normalized form
    assert torch.equal(res["pred_cam"][0], res["pred_cam"][1]) and torch.equal(res_n["pred_cam"][0], res["pred_cam"][0])
    want = ref(normalize_raw(u8))[0]
    got = torch.cat([res["pred_mano_params"]["global_orient"][0].reshape(-1),
                     res["pred_mano_params"]["hand_pose"][0].reshape(-1),
                     res["pred_mano_params"]["betas"][0], res["pred_cam"][0]])
    assert torch.equal(got, want)
    r256 = cpu_model_mano.forward_crops(u256)                            # the [32:224] slice of a 256 crop
    assert r256["pred_cam"].shape == (1, 3)
    with pytest.raises(ValueError):
        cpu_model_mano.forward_crops(torch.zeros(1, 3, 128, 128))


def test_close_and_context_manager(ref):
    with HamerModel.from_reference(ref, backend="cpu") as m:
        assert m(SAMPLE, bbox=BOX_R).rotmats.shape == (16, 3, 3)
    with pytest.raises(RuntimeError):
        m(SAMPLE, bbox=BOX_R)
    m.close()  # idempotent
    assert "closed" in repr(m)


# --------------------------------------------------------------------------- #
# Warm-up API (from_pretrained(warmup_variants=...), model.warmup(**variant))
# --------------------------------------------------------------------------- #
def test_warmup_signature_and_defaults():
    import inspect

    from hamer.api import DEFAULT_WARMUP_VARIANTS, WARMUP_FRAME_SIZE, normalize_warmup_variant

    sig = inspect.signature(HamerModel.from_pretrained)
    assert sig.parameters["warmup_variants"].default == "default"
    assert sig.parameters["verbose"].default is False
    assert inspect.signature(HamerModel.from_reference).parameters["warmup_variants"].default == "auto"
    assert "DEFAULT_WARMUP_VARIANTS" in hamer.api.__all__
    keys = [normalize_warmup_variant(**v) for v in DEFAULT_WARMUP_VARIANTS]
    assert len(set(keys)) == len(keys)  # no duplicates
    calls = {dict(k)["call"] for k in keys}
    assert calls == {"frame", "crops", "run_crop"}
    frame = [dict(k) for k in keys if dict(k)["call"] == "frame"]
    assert {f["input"] for f in frame} == {"path", "pil", "array", "tensor"}
    assert {f.get("fmt") for f in frame if f["input"] == "path"} == {"jpeg", "png"}
    assert {f["is_right"] for f in frame} >= {True, False, "mixed"}
    assert {f["bbox"] for f in frame} == {"box", "center"}
    assert all(f["frame_size"] == WARMUP_FRAME_SIZE for f in frame if f.get("fmt") != "png")


def test_normalize_warmup_variant_rules():
    from hamer.api import normalize_warmup_variant as nv

    assert nv() == nv(call="frame", input="path", fmt="jpeg", hands=1, is_right=True, bbox="box")
    assert dict(nv(hands=2))["is_right"] == "mixed"
    assert "fmt" not in dict(nv(input="array", fmt="png"))      # fmt only matters for files
    assert dict(nv(frame_size=[512, 334]))["frame_size"] == (512, 334)
    assert dict(nv(call="crops"))["batch"] == 2
    for bad in (dict(call="nope"), dict(bogus=1), dict(input="url"), dict(hands=0), dict(is_right="mixed"),
                dict(frame_size=(8, 8)), dict(call="crops", crop_width=128), dict(call="run_crop", batch=2),
                dict(call="crops", batch=0), dict(call="run_crop", crop_dtype="fp16")):
        with pytest.raises(ValueError):
            nv(**bad)


def test_warmup_spec_parsing():
    from hamer.api import DEFAULT_WARMUP_VARIANTS, _warmup_spec

    assert _warmup_spec("default") == [dict(v) for v in DEFAULT_WARMUP_VARIANTS] == _warmup_spec("all")
    for none in (None, "none", False, []):
        assert _warmup_spec(none) == []
    assert _warmup_spec({"call": "run_crop"}) == [{"call": "run_crop"}]
    with pytest.raises(ValueError):
        _warmup_spec("fast")
    with pytest.raises(ValueError):
        _warmup_spec(["run_crop"])


def test_bad_warmup_spec_fails_before_build(ref):
    with pytest.raises(ValueError):
        HamerModel.from_reference(ref, backend="cpu", warmup_variants=[{"call": "frame", "input": "url"}])


def test_warmup_is_idempotent_and_does_not_change_outputs(ref, mano):
    small = (320, 240)  # the CPU backend is slow: small dummy frames
    variants = [{"call": "run_crop"}, {"call": "frame", "input": "array", "is_right": False, "frame_size": small}]
    with HamerModel.from_reference(ref, backend="cpu", mano=mano) as plain:
        want_l = plain(SAMPLE, bbox=BOX_L, is_right=False).to_dict(rounded=False)
        want_r = plain(SAMPLE, bbox=BOX_R).to_dict(rounded=False)
        assert plain.warmed_variants == [] and plain.info.get("warmup_variants") == []
    with HamerModel.from_reference(ref, backend="cpu", mano=mano, warmup_variants=variants) as m:
        assert len(m.warmed_variants) == 2 and len(m.info["warmup_variants"]) == 2
        assert all(r["ms"] > 0 for r in m.info["warmup_variants"])
        again = m.warmup(call="frame", input="array", is_right=False, frame_size=list(small))
        assert again["warmed"] is False and len(m.warmed_variants) == 2
        new = m.warmup(call="frame", input="path", fmt="png", hands=2, bbox="center", frame_size=small)
        assert new["warmed"] is True and new["variant"]["is_right"] == "mixed" and len(m.warmed_variants) == 3
        assert m(SAMPLE, bbox=BOX_L, is_right=False).to_dict(rounded=False) == want_l
        assert m(SAMPLE, bbox=BOX_R).to_dict(rounded=False) == want_r
    with pytest.raises(RuntimeError):
        m.warmup(call="run_crop")


def test_quiet_logs(monkeypatch):
    from hamer import _verbosity
    from hamer.api import QUIET_LOG_ENV, _quiet_logs

    for k in QUIET_LOG_ENV:
        monkeypatch.delenv(k, raising=False)
    monkeypatch.delitem(sys.modules, "ttnn", raising=False)
    monkeypatch.setenv("LOGURU_LEVEL", "DEBUG")  # a caller value stays
    try:
        assert _quiet_logs(verbose=False) == {"TT_LOGGER_LEVEL": "error"}
        assert os.environ["TT_LOGGER_LEVEL"] == "error" and os.environ["LOGURU_LEVEL"] == "DEBUG"
        assert _verbosity.quiet()
        monkeypatch.delenv("TT_LOGGER_LEVEL")
        assert _quiet_logs(verbose=True) == {} and "TT_LOGGER_LEVEL" not in os.environ
        assert not _verbosity.quiet()
    finally:
        _verbosity.set_quiet(False)