# SPDX-License-Identifier: Apache-2.0
"""Device test of the Python API: ``Sam3(...)`` output == the server / pipeline output for the demo image (bit for
bit), and the warm ``model(...)`` time is not slower than the published model-call number.

    source tools/chipenv.sh <chip> <model>/.venv
    timeout -s INT 1800 python -m pytest -q -s code/tt_sam3/tests/test_api_device.py

The published model-call number (served_bench "selected", OPT_REPORT audit integration 2026-10-04, chip 7: decoded
1920x1080 demo image -> forward + selection, 1 prompt) is 82.32 ms median / 80.82 ms min.  ``API_PERF_JSON=<path>``
writes the measured numbers.
"""
import base64
import json
import os
import statistics
import time
import types

import numpy as np
import pytest
import torch

pytestmark = pytest.mark.skipif(os.environ.get("TT_VISIBLE_DEVICES", "") == "none", reason="needs a device")

MODEL_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
DEMO = os.path.join(MODEL_DIR, "media", "source_1.png")
PUBLISHED_MODEL_CALL_MS = {"median": 82.32, "min": 80.82}
# published demo results (served_bench / server_smoke, rounds 4-6 and the audit integration: identical every time)
PUBLISHED = {
    "car": {"scores": [0.9731, 0.8171, 0.7969], "areas": [175155, 390, 3587],
            "boxes": [[545.17, 447.36, 1159.19, 854.73], [392.59, 533.37, 419.39, 549.49], [964.91, 510.56, 1077.54, 584.65]]},
    "person": {"scores": [0.5936], "areas": [386]},
    "tree": {"scores": [0.9174, 0.87, 0.7769, 0.5161, 0.5029], "areas": [33417, 38059, 4255, 5903, 1674]},
}
ITERS = int(os.environ.get("API_PERF_ITERS", "30"))


@pytest.fixture(scope="module")
def model():
    from tt_sam3 import Sam3

    t0 = time.perf_counter()
    m = Sam3.from_pretrained(device_id=0)
    print(f"\nfrom_pretrained: {time.perf_counter() - t0:.1f} s, info={m.info}", flush=True)
    yield m
    m.close()
    assert m.closed


def _server_results(model, prompts):
    """The HTTP server's request handler (``_predict_core``) on the same pipeline."""
    from tt_sam3.server import app as A

    raw = open(DEMO, "rb").read()
    saved = dict(A.STATE)
    A.STATE.clear()
    A.STATE.update(pipeline=model.pipeline, config=types.SimpleNamespace(max_prompts=8, max_side=4096),
                   rgb_input=model.info["rgb_input"])
    try:
        A._LAST_DECODED[:] = [None, None]
        return A._predict_core(raw, prompts, 0.5, 50, "png", A._decode_bytes, "enc")
    finally:
        A.STATE.clear()
        A.STATE.update(saved)


def test_api_equals_server_and_pipeline(model):
    from PIL import Image

    from tt_sam3.api import load_image

    prompts = ["car", "person", "tree"]
    res = model(DEMO, prompts, return_mask_logits=True)
    # 1) the server's handler, same pipeline, same image
    j = _server_results(model, prompts)
    assert j["width"] == 1920 and j["height"] == 1080
    for r, jr in zip(res, j["results"]):
        assert r.prompt == jr["prompt"] and len(r) == jr["num_instances"] and round(r.presence, 4) == jr["presence"]
        for i, inst in enumerate(jr["instances"]):
            assert round(float(r.scores[i]), 4) == inst["score"]
            assert [round(float(v), 2) for v in r.boxes[i]] == inst["box_xyxy"]
            png = np.asarray(Image.open(__import__("io").BytesIO(base64.b64decode(inst["mask_png"])))) > 0
            assert np.array_equal(r.masks[i], png) and int(r.masks[i].sum()) == inst["mask_area"]
    # 2) the pipeline's selection tuples, unrounded: bit-identical scores / boxes / mask logits / presence
    rgb = load_image(DEMO)
    model._invalidate_cache()
    with torch.inference_mode():
        sel = model.pipeline.forward_rgb_selected_multi(rgb, prompts, 0.5, 50)
    for r, (scores, boxes, logits, presence) in zip(res, sel):
        assert np.array_equal(r.scores, scores.float().numpy()) and r.presence == float(presence)
        assert np.array_equal(r.mask_logits, logits.float().numpy())
        from tt_sam3.api import boxes_cxcywh_to_xyxy_pixels

        assert np.array_equal(r.boxes, boxes_cxcywh_to_xyxy_pixels(boxes, 1920, 1080))
    # 3) the published demo results
    for r in res:
        pub = PUBLISHED[r.prompt]
        assert [round(float(s), 4) for s in r.scores] == pub["scores"], r
        assert [int(m.sum()) for m in r.masks] == pub["areas"], r
        if "boxes" in pub:
            assert [[round(float(v), 2) for v in b] for b in r.boxes] == pub["boxes"]
    # 4) every input kind and the image cache give the same bits
    pil = Image.open(DEMO)
    for src in (rgb, np.ascontiguousarray(rgb), pil, torch.from_numpy(np.ascontiguousarray(rgb)).permute(2, 0, 1), DEMO):
        r2 = model(src, "car")
        assert np.array_equal(r2.masks, res[0].masks) and np.array_equal(r2.scores, res[0].scores)
        assert np.array_equal(r2.boxes, res[0].boxes)
    model(rgb, "car")  # the loop ended with the file (an encoded-bytes key): this call makes rgb resident
    hits = model.pipeline._full.stats.get("image_cache_hits", 0)
    model(rgb, "person")  # same pixels: the image part is skipped
    assert model.pipeline._full.stats.get("image_cache_hits", 0) == hits + 1
    many = model.predict_many([DEMO, rgb], "car")
    assert all(np.array_equal(m.masks, res[0].masks) for m in many)


def _times(fn, n):
    t = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        t.append((time.perf_counter() - t0) * 1e3)
    return {"median": round(statistics.median(t), 2), "min": round(min(t), 2)}


def test_api_warm_call_time(model):
    from tt_sam3.api import load_image

    rgb = np.ascontiguousarray(load_image(DEMO))
    pipe = model.pipeline

    def server_path():  # what the server's handler runs for a decoded image (no cache key, 1 prompt)
        with torch.inference_mode():
            pipe.forward_rgb_selected_multi(rgb, ["car"], 0.5, 50)

    out = {"published_model_call_ms": PUBLISHED_MODEL_CALL_MS, "iters": ITERS, "info": {k: str(v) for k, v in model.info.items()}}
    model.cache_image = False  # every call runs the whole image graph (as for a new image)
    for f in (server_path, lambda: model(rgb, "car", return_masks=False), lambda: model(rgb, "car")):
        _times(f, 5)  # warm
    rounds = {"pipeline_direct": [], "api_no_masks": [], "api_masks": []}
    for _ in range(3):  # interleaved A/B rounds
        rounds["pipeline_direct"].append(_times(server_path, ITERS // 3))
        rounds["api_no_masks"].append(_times(lambda: model(rgb, "car", return_masks=False), ITERS // 3))
        rounds["api_masks"].append(_times(lambda: model(rgb, "car"), ITERS // 3))
    for k, v in rounds.items():
        out[k] = {"median": round(statistics.median(x["median"] for x in v), 2), "min": min(x["min"] for x in v)}
    out["api_path_with_decode"] = _times(lambda: model(DEMO, "car"), ITERS)  # PNG read + decode every call (no cache)
    model.cache_image = True
    model(rgb, "car")
    out["api_repeated_image_cached"] = _times(lambda: model(rgb, "car"), ITERS)
    out["api_3_prompts_new_image"] = None
    model.cache_image = False
    out["api_3_prompts_new_image"] = _times(lambda: model(rgb, ["car", "person", "tree"]), ITERS // 2)
    model.cache_image = True
    print("\n" + json.dumps(out, indent=1), flush=True)
    if os.environ.get("API_PERF_JSON"):
        with open(os.environ["API_PERF_JSON"], "w") as f:
            json.dump(out, f, indent=1)
    # forward + selection through the API (no mask upsample) vs the published number and vs the direct call
    tol = float(os.environ.get("API_PERF_TOL", "1.05"))
    assert out["api_no_masks"]["median"] <= PUBLISHED_MODEL_CALL_MS["median"] * tol, out
    assert out["api_no_masks"]["median"] <= out["pipeline_direct"]["median"] * tol, out