# SPDX-License-Identifier: Apache-2.0
"""Host tests (no device) for the serving layer: the image-result cache (SAM3_IMG_CACHE) key / hit / invalidation
logic in the device runner and in server/app.py, and the HTTP routes with a fake pipeline (FastAPI TestClient).

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

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

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))  # <repo>/code

from tt_sam3.server import app as A  # noqa: E402
from tt_sam3.tt import ttnn_sam3_device_graph as dg  # noqa: E402

xxhash = pytest.importorskip("xxhash")


# ------------------------------------------------------------------------------------------------
# device runner: _run_image key logic (fake ttnn, fake upload)
# ------------------------------------------------------------------------------------------------
def _fake_runner(monkeypatch):
    log = []
    fake_ttnn = types.SimpleNamespace(
        execute_trace=lambda dev, tid, cq_id=0, blocking=False: log.append(("replay", tid)),
        release_trace=lambda dev, tid: log.append(("release", tid)),
    )
    monkeypatch.setattr(dg, "ttnn", fake_ttnn)
    monkeypatch.setattr(dg, "_BANDED_INPUT", True)
    r = object.__new__(dg.TtSam3FullDevice)
    r.device = None
    r.stats = {}
    r.img_key = None
    r.trace_img, r.trace_id = "IMG", "PROMPT"
    r.hx_bands = ["band"]
    r.g = types.SimpleNamespace(_dec0=None)
    r._upload_bands = lambda px: log.append(("upload", px))
    return r, log


def test_run_image_hit_skips_upload_and_trace(monkeypatch):
    r, log = _fake_runner(monkeypatch)
    r._run_image("px1", ("enc", 1, b"a"))
    assert log == [("upload", "px1"), ("replay", "IMG")]
    r._run_image("px1", ("enc", 1, b"a"))  # same key: nothing on the device
    assert len(log) == 2 and r.stats["image_cache_hits"] == 1
    r._run_image("px2", ("enc", 1, b"b"))  # different image: full path, the key moves
    assert log[2:] == [("upload", "px2"), ("replay", "IMG")]
    r._run_image("px1", ("enc", 1, b"a"))  # one entry only: the first image is no longer resident
    assert log[4:] == [("upload", "px1"), ("replay", "IMG")]
    assert r.stats["image_cache_hits"] == 1 and r.stats["image_replays"] == 3


def test_run_image_no_key_always_runs_and_clears(monkeypatch):
    r, log = _fake_runner(monkeypatch)
    r._run_image("px1", ("enc", 1, b"a"))
    r._run_image("px1")  # forward_image / __call__ path: no key -> runs, and the resident image is now unkeyed
    assert r.img_key is None
    r._run_image("px1", ("enc", 1, b"a"))
    assert [e for e in log if e[0] == "replay"] == [("replay", "IMG")] * 3
    assert "image_cache_hits" not in r.stats


def test_run_image_failed_upload_leaves_no_key(monkeypatch):
    r, log = _fake_runner(monkeypatch)
    r._run_image("px1", ("enc", 1, b"a"))

    def boom(px):
        raise RuntimeError("upload failed")

    r._upload_bands = boom
    with pytest.raises(RuntimeError):
        r._run_image("px2", ("enc", 1, b"b"))
    assert r.img_key is None  # the resident outputs may be half rewritten: no hit for either image
    r._upload_bands = lambda px: log.append(("upload", px))
    r._run_image("px1", ("enc", 1, b"a"))
    assert log[-1] == ("replay", "IMG") and "image_cache_hits" not in r.stats


def test_release_clears_key(monkeypatch):
    r, log = _fake_runner(monkeypatch)
    r._run_image("px1", ("enc", 1, b"a"))
    r.release()
    assert r.img_key is None and r.trace_img is None


# ------------------------------------------------------------------------------------------------
# app: routes with a fake pipeline (deterministic outputs from the pixels it receives)
# ------------------------------------------------------------------------------------------------
class _FakeRunner:
    def __init__(self):
        self.stats = {}
        self.img_key = None


class _FakePipeline:
    """forward_rgb_selected_multi emulates the device runner's cache contract: a hit (same key) must return what the
    resident image gives, so a wrong hit (stale pixels) shows up as a wrong answer."""

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

    def prefetch_text(self, prompts):
        pass

    def forward_rgb_selected_multi(self, rgb, prompts, threshold, max_instances, image_key=None):
        if self.fail:
            raise RuntimeError("device error")
        if image_key is None or image_key != self._full.img_key:
            self.resident = np.ascontiguousarray(rgb).astype(np.float64)
            self._full.img_key = image_key
            hit = False
        else:
            self._full.stats["image_cache_hits"] = self._full.stats.get("image_cache_hits", 0) + 1
            hit = True
        self.calls.append((image_key, hit))
        out = []
        for k, p in enumerate(prompts):
            m = float(self.resident.mean()) / 255.0
            logits = torch.full((1, A.MASK_RES, A.MASK_RES), -1.0)
            logits[0, 10:40 + int(50 * m), 20:60] = 1.0 + k
            out.append((torch.tensor([0.5 + 0.4 * m]), torch.tensor([[0.5, 0.5, 0.2 + 0.1 * m, 0.3]]), logits, 0.9))
        return out


@pytest.fixture
def client(monkeypatch):
    from fastapi.testclient import TestClient

    pipe = _FakePipeline()
    cfg = types.SimpleNamespace(max_prompts=4, max_side=4096)
    monkeypatch.setattr(A, "STATE", {"pipeline": pipe, "config": cfg, "rgb_input": True})
    monkeypatch.setattr(A, "_IMG_CACHE", True)
    A._LAST_DECODED[:] = [None, None]
    yield TestClient(A.app), pipe
    A._LAST_DECODED[:] = [None, None]


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


def _img(seed, h=97, w=131):
    return np.random.default_rng(seed).integers(0, 256, size=(h, w, 3), dtype=np.uint8)


def _post_json(c, png, prompts=("car",)):
    r = c.post("/predict", json={"image": base64.b64encode(png).decode(), "prompts": list(prompts)})
    assert r.status_code == 200, r.text
    return r.json()


def test_predict_cache_hit_same_results(client):
    c, pipe = client
    p1, p2 = _png(_img(1)), _png(_img(2))
    a = _post_json(c, p1)
    b = _post_json(c, p1, ("tree",))  # same image, other prompt: hit
    assert pipe.calls[0][1] is False and pipe.calls[1][1] is True and pipe.calls[0][0] == pipe.calls[1][0]
    assert a["results"][0]["instances"] == b["results"][0]["instances"]
    _post_json(c, p2)  # different image: miss
    assert pipe.calls[2][1] is False and pipe.calls[2][0] != pipe.calls[0][0]
    a2 = _post_json(c, p1)  # back to the first image: miss again (one entry), same answer
    assert pipe.calls[3][1] is False and a2["results"] == a["results"]


def test_predict_cache_off(client, monkeypatch):
    c, pipe = client
    monkeypatch.setattr(A, "_IMG_CACHE", False)
    p1 = _png(_img(1))
    a, b = _post_json(c, p1), _post_json(c, p1)
    assert [k for k, _ in pipe.calls] == [None, None] and a["results"] == b["results"]


def test_predict_failure_invalidates(client):
    c, pipe = client
    p1 = _png(_img(1))
    _post_json(c, p1)
    assert pipe._full.img_key is not None and A._LAST_DECODED[0] is not None
    pipe.fail = True
    r = c.post("/predict", json={"image": base64.b64encode(p1).decode(), "prompts": ["car"]})
    assert r.status_code == 500
    assert pipe._full.img_key is None and A._LAST_DECODED == [None, None]
    pipe.fail = False
    _post_json(c, p1)
    assert pipe.calls[-1][1] is False


def test_predict_response_schema_unchanged(client):
    c, _ = client
    j = _post_json(c, _png(_img(3)), ("car", "tree"))
    assert set(j) == {"width", "height", "input_size", "threshold", "mask_format", "results", "timing_ms"}
    assert set(j["timing_ms"]) == {"decode", "preprocess", "forward", "forward_per_prompt", "postprocess", "total"}
    assert [r["prompt"] for r in j["results"]] == ["car", "tree"]
    assert set(j["results"][0]["instances"][0]) == {"score", "box_xyxy", "mask_area", "mask_png"}
    r = c.post("/predict", json={"image": "!!!", "prompts": ["car"]})
    assert r.status_code == 400
    r = c.post("/predict", json={"image": base64.b64encode(_png(_img(3))).decode(), "prompts": [" "]})
    assert r.status_code == 400


# ------------------------------------------------------------------------------------------------
# /predict_raw (binary body)
# ------------------------------------------------------------------------------------------------
def test_predict_raw_png_and_rgb_match_predict(client):
    c, pipe = client
    arr = _img(5)
    png = _png(arr)
    j = _post_json(c, png, ("car", "tree"))
    r = c.post("/predict_raw", params=[("prompts", "car"), ("prompts", "tree")], content=png,
               headers={"Content-Type": "image/png"})
    assert r.status_code == 200, r.text
    jr = r.json()
    assert jr["results"] == j["results"] and set(jr) == set(j)
    assert pipe.calls[-1][1] is True  # same PNG bytes as the JSON request: same key, cache hit
    h, w = arr.shape[:2]
    r = c.post("/predict_raw", params={"prompts": ["car", "tree"]}, content=arr.tobytes(),
               headers={"Content-Type": "application/x-rgb24", "X-Width": str(w), "X-Height": str(h)})
    assert r.status_code == 200, r.text
    assert r.json()["results"] == j["results"] and (r.json()["width"], r.json()["height"]) == (w, h)
    assert pipe.calls[-1][1] is False and pipe.calls[-1][0][0] == f"rgb{w}x{h}"  # raw pixels never alias encoded bytes


def test_predict_raw_rgb_key_includes_shape(client):
    c, pipe = client
    arr = _img(6, 12, 12)
    body = arr.tobytes()
    for w, h in ((12, 12), (24, 6)):  # same bytes, different shape -> different image, no hit
        r = c.post("/predict_raw", params={"prompts": "car"}, content=body,
                   headers={"Content-Type": "application/x-rgb24", "X-Width": str(w), "X-Height": str(h)})
        assert r.status_code == 200, r.text
    assert [hit for _, hit in pipe.calls] == [False, False]


def test_predict_raw_errors(client):
    c, _ = client
    arr = _img(7)
    h, w = arr.shape[:2]
    hdr = {"Content-Type": "application/x-rgb24", "X-Width": str(w), "X-Height": str(h)}
    assert c.post("/predict_raw", params={"prompts": "car"}, content=arr.tobytes()[:-1], headers=hdr).status_code == 400
    assert c.post("/predict_raw", params={"prompts": "car"}, content=arr.tobytes(),
                  headers={"Content-Type": "application/x-rgb24"}).status_code == 400
    assert c.post("/predict_raw", params={"prompts": "car"}, content=b"x" * 30,
                  headers={**hdr, "X-Width": "5000", "X-Height": "2"}).status_code == 400
    assert c.post("/predict_raw", params={"prompts": "car"}, content=b"not an image",
                  headers={"Content-Type": "image/png"}).status_code == 400
    assert c.post("/predict_raw", params={"prompts": "car"}, content=b"", headers={"Content-Type": "image/png"}).status_code == 400
    assert c.post("/predict_raw", content=_png(arr), headers={"Content-Type": "image/png"}).status_code == 422
    assert c.post("/predict_raw", params={"prompts": " "}, content=_png(arr), headers={"Content-Type": "image/png"}).status_code == 400
    assert c.post("/predict_raw", params={"prompts": ["a"] * 5}, content=_png(arr),
                  headers={"Content-Type": "image/png"}).status_code == 400  # > max_prompts (4)
    assert c.post("/predict_raw", params={"prompts": "car", "threshold": 2}, content=_png(arr),
                  headers={"Content-Type": "image/png"}).status_code == 422


def test_predict_raw_not_ready(monkeypatch):
    from fastapi.testclient import TestClient

    monkeypatch.setattr(A, "STATE", {})
    c = TestClient(A.app)
    assert c.post("/predict_raw", params={"prompts": "car"}, content=b"x", headers={"Content-Type": "image/png"}).status_code == 503
    assert c.post("/predict", json={"image": "eA==", "prompts": ["car"]}).status_code == 503