# SPDX-License-Identifier: Apache-2.0
"""Host tests of the HTTP contract (no device, no weights): the real app of this bundle with a stub model, one
camera image per request.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sceneseg/tests/test_server_host.py
"""
from __future__ import annotations

import base64
import io

import numpy as np
import pytest

fastapi = pytest.importorskip("fastapi")
from fastapi.testclient import TestClient  # noqa: E402

from tt_sceneseg import io as tio  # noqa: E402
from tt_sceneseg.api import SceneSeg  # noqa: E402
from tt_sceneseg.server import app as server  # noqa: E402
from tt_sceneseg.tests.stubs import StubModel  # noqa: E402


@pytest.fixture
def client(monkeypatch):
    monkeypatch.setenv("TT_MESH_SHAPE", "1x1")
    monkeypatch.setenv("TT_MODEL_WEIGHTS_REVISION", "0" * 40)
    monkeypatch.setattr(server.app.state.ttaw, "model_factory", StubModel)
    with TestClient(server.app) as c:
        yield c


def _png(seed: int = 0, hw=(90, 160)) -> bytes:
    from PIL import Image

    buf = io.BytesIO()
    Image.fromarray(np.random.default_rng(seed).integers(0, 256, size=hw + (3,), dtype=np.uint8)).save(buf, "PNG")
    return buf.getvalue()


def _image(seed: int = 0, name: str = tio.CAMERA_NAME, **kw) -> dict:
    return {"camera": name, "data": base64.b64encode(_png(seed, **kw)).decode()}


def test_health_info_models(client):
    assert client.get("/health").json()["status"] == "ok"
    assert client.get("/v1/health").json()["status"] == "ok"
    info = client.get("/info").json()
    assert info["model"] == SceneSeg.MODEL_NAME and info["weights"]["revision"] == "0" * 40
    assert info["device"]["dispatch"] == "eth" and info["device"]["grid"] == "12x10"
    assert info["autoware"]["package"] == "models (VisionPilot middleware_recipes, node scene_seg_model)"
    assert info["output"]["labels"] == list(SceneSeg.LABELS)
    assert info["input"]["camera_order"] == list(SceneSeg.CAMERA_ORDER)
    assert info["variant"] == SceneSeg.DEFAULT_VARIANT
    assert client.get("/v1/models").json()["data"][0]["id"] == "changh95/sceneseg-p150"


def _decode(spec):
    from PIL import Image

    assert spec["format"] == "png"
    return np.asarray(Image.open(io.BytesIO(base64.b64decode(spec["data"]))))


def test_predict_ok_and_equals_api(client, tmp_path):
    r = client.post("/predict", json={"images": [_image(0)]})
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["frame_id"] == "camera" and body["model"] == SceneSeg.MODEL_NAME
    assert body["mask"]["format"] == "png" and body["mask"]["shape"] == [90, 160]      # the input's size
    assert body["class_map"]["shape"] == [90, 160] and body["class_names"] == list(SceneSeg.LABELS)
    mask, cmap = _decode(body["mask"]), _decode(body["class_map"])
    assert set(np.unique(mask)) <= {0, 255} and set(np.unique(cmap)) <= {0, 1, 2}
    np.testing.assert_array_equal(mask == 255, cmap == 1)
    assert sum(body["class_counts"].values()) == 90 * 160
    assert body["meta"]["channel_order"] == "bgr" and body["meta"]["network_hw"] == [320, 640]
    assert {"decode", "model_call", "total", "device"} <= set(body["timing_ms"])
    state = server.app.state.ttaw
    req = server.PredictRequest(images=[_image(1, name="any-name")])
    (tmp_path / "f.png").write_bytes(_png(1))
    api = state.model(str(tmp_path / "f.png")).to_dict()                  # the image as a path: the same pixels
    served = state.predict(req)
    assert served["mask"] == api["mask"] and served["class_map"] == api["class_map"]
    assert state.model(images=[str(tmp_path / "f.png")]).to_dict()["mask"] == api["mask"]
    no_map = client.post("/predict", json={"images": [_image(1)], "params": {"class_map": False}}).json()
    assert "class_map" not in no_map and no_map["mask"] == api["mask"] and no_map["class_counts"] == api["class_counts"]


def test_predict_npz_output(client):
    r = client.post("/predict", json={"images": [_image(2, hw=(48, 64))], "output_format": "npz"})
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["mask"]["format"] == "npz" and body["mask"]["shape"] == [48, 64]
    assert body["class_map"]["format"] == "npz" and body["class_map"]["shape"] == [48, 64]


@pytest.mark.parametrize("payload,code", [
    ({"images": [{"camera": "CAM_FRONT", "data": "not base64!"}]}, 400),                  # undecodable
    ({"images": [{"camera": "CAM_FRONT", "data": base64.b64encode(b"jpeg?").decode()}]}, 400),   # not an image
    ({"images": [_image(0, name="A"), _image(1, name="B")]}, 400),                         # one image per request
    ({"params": {"score_threshold": 0.5}, "images": [_image(0)]}, 400),                   # the node has no threshold
    ({"params": {"class_map": 1}, "images": [_image(0)]}, 400),                           # booleans must be booleans
    ({"params": {"no_such_knob": 1}, "images": [_image(0)]}, 400),
    ({"calibration": {"preset": "../etc"}, "images": [_image(0)]}, 400),
    ({"points": {"format": "list", "values": [[0, 0, 0, 0]]}}, 400),                     # a cloud is not an image
    ({}, 400),                                                                            # no image
    ({"imagez": []}, 422),                                                                # schema violation (typo)
])
def test_predict_errors(client, payload, code):
    r = client.post("/predict", json=payload)
    assert r.status_code == code, r.text


def test_503_while_starting(client, monkeypatch):
    monkeypatch.setattr(server.app.state.ttaw, "ready", False)
    assert client.post("/predict", json={}).status_code == 503
    assert client.get("/health").json()["status"] == "starting"


def test_mesh_shape_parser():
    p = server.parse_mesh_shape
    assert p("1x1") == p("(1, 1)") == p("1,1") == p(None) == (1, 1)
    with pytest.raises(RuntimeError):
        p("two by two")