# SPDX-License-Identifier: Apache-2.0
"""Host tests (no device, no weights) for the Python API ``tt_sam3.Sam3``: import surface, image input conversion,
argument checks, result format, the image cache, lifetime, and equality with the HTTP server's results for the same
pipeline outputs (fake pipeline).

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sam3/tests/test_api_host.py
"""
import base64
import io
import os
import subprocess
import sys
import types

import numpy as np
import pytest
import torch
from PIL import Image

import tt_sam3
from tt_sam3 import api

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")


# ------------------------------------------------------------------------------------------------
# import surface
# ------------------------------------------------------------------------------------------------
def test_import_is_light_and_exports():
    code = ("import sys, tt_sam3; assert 'torch' not in sys.modules and 'ttnn' not in sys.modules; "
            "from tt_sam3 import Sam3, Sam3Result, load_image, device_open_kwargs, PINNED_REVISION; "
            "assert 'ttnn' not in sys.modules, 'ttnn imported at API import'; print(tt_sam3.__version__)")
    env = {**os.environ, "TT_VISIBLE_DEVICES": "none"}
    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, env=env, timeout=300)
    assert out.returncode == 0, out.stderr[-2000:]
    assert out.stdout.strip() == tt_sam3.__version__
    for name in ("from_pretrained", "__call__", "predict_many", "close", "__enter__", "__exit__"):
        assert callable(getattr(api.Sam3, name))
    assert api.Sam3.from_pretrained.__doc__ and api.Sam3.__call__.__doc__ and api.Sam3Result.__doc__


def test_no_machine_specific_paths_in_runtime_code():
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    bad = []
    for d, _, files in os.walk(root):
        if os.sep + "tests" in d or "__pycache__" in d:
            continue
        for f in files:
            if f.endswith(".py"):
                text = open(os.path.join(d, f), encoding="utf-8").read()
                if "/home/" in text or "deepgadget" in text:
                    bad.append(os.path.join(d, f))
    assert not bad, bad


def test_package_data_declared():
    text = open(os.path.join(MODEL_DIR, "code", "pyproject.toml")).read()
    assert '"assets/*.gz"' in text and '"kernels/*.cpp"' in text
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    assert os.path.isfile(os.path.join(root, "assets", "bpe_simple_vocab_16e6.txt.gz"))


# ------------------------------------------------------------------------------------------------
# image input
# ------------------------------------------------------------------------------------------------
def _img(seed, h=37, w=53, c=3):
    return np.random.default_rng(seed).integers(0, 256, size=(h, w, c), dtype=np.uint8)


def _png(arr):
    b = io.BytesIO()
    Image.fromarray(arr).save(b, format="PNG")
    return b.getvalue()


def test_load_image_kinds(tmp_path):
    a = _img(0)
    p = tmp_path / "x.png"
    p.write_bytes(_png(a))
    for src in (str(p), p, p.read_bytes(), Image.fromarray(a), Image.fromarray(np.dstack([a, a[..., :1]])), a,
                torch.from_numpy(a).permute(2, 0, 1), torch.from_numpy(a).permute(2, 0, 1)[None],
                np.moveaxis(a, -1, 0), a.astype(np.float32) / 255.0):
        rgb = api.load_image(src)
        assert rgb.dtype == np.uint8 and rgb.shape == (37, 53, 3)
        assert np.array_equal(rgb, a), type(src)
    g = a[..., 0]
    assert np.array_equal(api.load_image(g), np.dstack([g, g, g]))
    assert np.array_equal(api.load_image(Image.fromarray(g)), np.dstack([g, g, g]))
    assert np.array_equal(api.load_image(a.astype(np.float64)), a)  # 0..255 floats
    with pytest.raises(TypeError):
        api.load_image(123)
    with pytest.raises(ValueError):
        api.load_image(np.zeros((2, 2, 2, 2)))


def test_load_image_float_matches_upstream_todtype():
    from torchvision.transforms import v2

    x = torch.rand(3, 19, 23)
    ref = v2.ToDtype(torch.uint8, scale=True)(x).permute(1, 2, 0).numpy()
    assert np.array_equal(api.load_image(x), ref)


@pytest.mark.skipif(not os.path.isfile(DEMO), reason="demo image missing")
def test_demo_decode_equals_server_decode():
    from tt_sam3.server import app as A

    raw = open(DEMO, "rb").read()
    assert np.array_equal(api.load_image(DEMO), A._decode_bytes(raw).rgb)
    assert np.array_equal(api.load_image(Image.open(DEMO)), A._decode_bytes(raw).rgb)


def test_prompt_checks():
    assert api._check_prompts("car") == (["car"], True)
    assert api._check_prompts([" car ", "tree"]) == (["car", "tree"], False)
    for bad in ([], [""], ["  "], [1]):
        with pytest.raises(ValueError):
            api._check_prompts(bad)


def test_boxes_match_server_rounding():
    from tt_sam3.server import app as A

    b = torch.tensor([[0.5, 0.5, 0.2, 0.3], [0.01, 0.99, 0.5, 0.5], [0.3, 0.2, 0.1, 0.05]])
    ours = api.boxes_cxcywh_to_xyxy_pixels(b, 1920, 1080)
    assert ours.dtype == np.float32 and ours.shape == (3, 4)
    assert [[round(float(v), 2) for v in r] for r in ours] == A._boxes_to_xyxy_pixels(b, 1920, 1080)
    assert api.boxes_cxcywh_to_xyxy_pixels(torch.zeros(0, 4), 10, 10).shape == (0, 4)


# ------------------------------------------------------------------------------------------------
# fake pipeline: the API against the server on the same pipeline outputs
# ------------------------------------------------------------------------------------------------
class _FakeRunner:
    def __init__(self):
        self.stats = {}
        self.img_key = None


class _FakePipeline:
    """Deterministic outputs from the pixels it receives; emulates the device runner's image-cache contract (a wrong
    hit would return the resident image's answer)."""

    def __init__(self):
        self._full = _FakeRunner()
        self.calls = []
        self.resident = None
        self.fail = False
        self.released = 0

    def prefetch_text(self, prompts):
        pass

    def release(self):
        self.released += 1

    def forward_rgb_selected_multi(self, rgb, prompts, threshold, max_instances, image_key=None):
        if self.fail:
            raise RuntimeError("device error")
        hit = image_key is not None and image_key == self._full.img_key
        if not hit:
            self.resident = np.ascontiguousarray(rgb).astype(np.float64)
            self._full.img_key = image_key
        self.calls.append((image_key, hit))
        out = []
        for k, _p in enumerate(prompts):
            m = float(self.resident.mean()) / 255.0
            logits = torch.randn(3, api.MASK_RES, api.MASK_RES, generator=torch.Generator().manual_seed(k + int(m * 1e6)))
            logits[0, 10:40 + int(50 * m), 20:60] += 3.0
            scores = torch.tensor([0.5 + 0.4 * m, 0.45, 0.3])
            boxes = torch.tensor([[0.5, 0.5, 0.2 + 0.1 * m, 0.3], [0.1, 0.1, 0.3, 0.3], [0.9, 0.9, 0.5, 0.5]])
            keep = scores > threshold
            out.append((scores[keep][:max_instances], boxes[keep][:max_instances], logits[keep][:max_instances], 0.9))
        return out


def _model(cache=True, fast_mask=True):
    pytest.importorskip("xxhash")
    pipe = _FakePipeline()
    m = api.Sam3(pipe, device=None, owns_device=False, info={"rgb_input": True, "fast_mask": fast_mask}, cache_image=cache)
    return m, pipe


@pytest.mark.parametrize("fast_mask", [True, False])
def test_results_equal_server(monkeypatch, fast_mask):
    pytest.importorskip("fastapi")
    from fastapi.testclient import TestClient

    from tt_sam3.server import app as A

    model, pipe = _model(fast_mask=fast_mask)
    arr = _img(3, 97, 131)
    png = _png(arr)
    res = model(png, ["car", "tree"], threshold=0.4)
    assert isinstance(res, list) and [r.prompt for r in res] == ["car", "tree"]
    monkeypatch.setattr(A, "STATE", {"pipeline": pipe, "config": types.SimpleNamespace(max_prompts=4, max_side=4096),
                                     "rgb_input": True})
    j = TestClient(A.app).post("/predict", json={"image": base64.b64encode(png).decode(), "prompts": ["car", "tree"],
                                                 "threshold": 0.4}).json()
    for r, jr in zip(res, j["results"]):
        assert r.prompt == jr["prompt"] and round(r.presence, 4) == jr["presence"] and len(r) == jr["num_instances"]
        assert r.masks.shape == (len(r), 97, 131) and r.masks.dtype == bool
        assert r.boxes.shape == (len(r), 4) and r.scores.shape == (len(r),)
        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_mask = np.asarray(Image.open(io.BytesIO(base64.b64decode(inst["mask_png"])))) > 0
            assert np.array_equal(r.masks[i], png_mask) and int(r.masks[i].sum()) == inst["mask_area"]


def test_single_prompt_and_result_access():
    model, _ = _model()
    r = model(_img(1), "car")
    assert isinstance(r, api.Sam3Result)
    assert r["masks"] is r.masks and set(r.keys()) >= {"masks", "boxes", "scores"}
    assert (r.width, r.height) == (53, 37) and r.mask_logits is None
    d = r.to_dict()
    assert d["num_instances"] == len(r) and len(d["mask_areas"]) == len(r)
    assert "Sam3Result(prompt='car'" in repr(r)
    r2 = model(_img(1), "car", return_masks=False, return_mask_logits=True, max_instances=1, threshold=0.0)
    assert r2.masks is None and r2.mask_logits.shape == (1, 288, 288) and len(r2) == 1
    with pytest.raises(ValueError):
        model(_img(1), "car", threshold=1.5)
    with pytest.raises(ValueError):
        model(_img(1), "car", max_instances=0)
    with pytest.raises(ValueError):
        model(_img(1), [])


def test_image_cache_keys(tmp_path):
    model, pipe = _model()
    a, b = _img(1), _img(2)
    r1 = model(a, "car")
    r2 = model(a.copy(), "tree")  # same pixels, other array object: hit
    assert pipe.calls[0][1] is False and pipe.calls[1][1] is True
    a2 = a.copy()
    a2[0, 0, 0] ^= 1  # one pixel changed: miss
    model(a2, "car")
    assert pipe.calls[2][1] is False
    model(b, "car")
    assert pipe.calls[3][1] is False
    p = tmp_path / "a.png"
    p.write_bytes(_png(a))
    model(str(p), "car")
    model(str(p), "car")
    assert pipe.calls[-1][1] is True and pipe.calls[-1][0][0] == "enc"
    assert np.array_equal(r1.masks, model(a, "car").masks) and len(r2) >= 0


def test_image_cache_off():
    model, pipe = _model(cache=False)
    a = _img(1)
    model(a, "car")
    model(a, "car")
    assert [k for k, _ in pipe.calls] == [None, None]


def test_failure_invalidates_cache():
    model, pipe = _model()
    a = _img(1)
    model(a, "car")
    assert pipe._full.img_key is not None
    pipe.fail = True
    with pytest.raises(RuntimeError):
        model(a, "car")
    assert pipe._full.img_key is None
    pipe.fail = False
    model(a, "car")
    assert pipe.calls[-1][1] is False


def test_predict_many_equals_loop(tmp_path):
    model, _ = _model(cache=False)
    imgs = [_img(i, 40 + i, 50) for i in range(4)]
    paths = []
    for i, a in enumerate(imgs):
        p = tmp_path / f"{i}.png"
        p.write_bytes(_png(a))
        paths.append(str(p))
    many = model.predict_many(paths, ["car", "tree"])
    loop = [model(a, ["car", "tree"]) for a in imgs]
    assert len(many) == 4
    for rs, ls in zip(many, loop):
        for r, l in zip(rs, ls):
            assert np.array_equal(r.masks, l.masks) and np.array_equal(r.scores, l.scores) and np.array_equal(r.boxes, l.boxes)
    single = model.predict_many(imgs[:2], "car")
    assert all(isinstance(r, api.Sam3Result) for r in single)
    assert model.predict_many([], "car") == []


def test_close_and_context_manager():
    model, pipe = _model()
    with model as m:
        assert m is model and not m.closed
        m(_img(1), "car")
    assert model.closed and pipe.released == 1 and model.pipeline is None
    model.close()  # idempotent
    assert pipe.released == 1
    with pytest.raises(RuntimeError):
        model(_img(1), "car")


def test_dispatch_resolution(monkeypatch, tmp_path):
    with pytest.raises(ValueError):
        api._resolve_dispatch("tensix")
    assert api._resolve_dispatch("worker") == "worker" and api._resolve_dispatch("eth") == "eth"
    d = tmp_path / "tt_metal" / "impl" / "dispatch"
    d.mkdir(parents=True)
    (d / "topology.cpp").write_text("static const std::vector<X> single_chip_arch_1cq_no_dispatch_s = {};")
    monkeypatch.setenv("TT_METAL_HOME", str(tmp_path))
    assert api.eth_dispatch_available() and api._resolve_dispatch("auto") == "eth"
    (d / "topology.cpp").write_text("stock tree")
    with pytest.warns(RuntimeWarning, match="ETH-dispatch patch"):
        assert api._resolve_dispatch("auto") == "worker"
    assert "dispatch_core_config" not in api.device_open_kwargs(0, "worker")
    kw = api.device_open_kwargs(3, "worker")
    assert kw["device_id"] == 3 and kw.get("trace_region_size", 0) > 0 and kw.get("l1_small_size", 0) >= 32768


def test_checkpoint_resolution(tmp_path, monkeypatch):
    monkeypatch.delenv("SAM3_WEIGHTS_DIR", raising=False)
    (tmp_path / "sam3.pt").write_bytes(b"x")
    assert api._resolve_checkpoint("facebook/sam3", None, str(tmp_path)) == str(tmp_path / "sam3.pt")
    assert api._resolve_checkpoint("facebook/sam3", None, str(tmp_path / "sam3.pt")) == str(tmp_path / "sam3.pt")
    monkeypatch.setenv("SAM3_WEIGHTS_DIR", str(tmp_path))
    assert api._resolve_checkpoint("facebook/sam3", "abc", None) == str(tmp_path / "sam3.pt")
    with pytest.raises(FileNotFoundError):
        api._resolve_checkpoint("facebook/sam3", None, str(tmp_path / "missing.pt"))


def test_server_and_api_share_the_eth_default(monkeypatch, tmp_path):
    """The server (SAM3_DISPATCH, default auto) and the Python API resolve the dispatch through one module:
    with the patched tree, auto = ETH dispatch + 1 CQ."""
    from tt_sam3.server import app as appmod
    from tt_sam3.tt import device_open

    d = tmp_path / "tt_metal" / "impl" / "dispatch"
    d.mkdir(parents=True)
    (d / "topology.cpp").write_text("single_chip_arch_1cq_no_dispatch_s")
    monkeypatch.setenv("TT_METAL_HOME", str(tmp_path))
    monkeypatch.delenv("SAM3_DISPATCH", raising=False)
    assert appmod.read_config().dispatch == "auto"
    assert device_open.resolve_dispatch(appmod.read_config().dispatch) == "eth"
    monkeypatch.setenv("SAM3_DISPATCH", "worker")
    assert device_open.resolve_dispatch(appmod.read_config().dispatch) == "worker"
    assert device_open.add_dispatch_kwargs({"device_id": 0}, "worker") == {"device_id": 0}