# 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, six
cameras + calibration (inline or the shipped preset), the stream's ego pose and the ``ego`` CAN dynamics.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_bevformer/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_bevformer.api import BEVFormer  # noqa: E402
from tt_bevformer.server import app as server  # noqa: E402
from tt_bevformer.tests.stubs import StubModel  # noqa: E402
from tt_bevformer.ttaw.server.app import decode_request  # noqa: E402

CAMS = list(BEVFormer.CAMERA_ORDER)


@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_b64(seed: int, hw=(90, 160)) -> str:
    from PIL import Image

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


def _images(cams=CAMS):
    return [{"camera": c, "data": _png_b64(i)} for i, c in enumerate(cams)]


def _inline_calib(cams=CAMS):
    K = [[100.0, 0.0, 80.0], [0.0, 100.0, 45.0], [0.0, 0.0, 1.0]]
    return {"cameras": {c: {"intrinsics": K, "T_ref_from_camera": {"translation": [1.0, 0.1 * i, 1.5],
                                                                    "rotation_wxyz": [0.5, -0.5, 0.5, -0.5]}}
                        for i, c in enumerate(cams)}}


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"] == BEVFormer.MODEL_NAME and info["weights"]["revision"] == "0" * 40
    assert info["device"]["dispatch"] == "eth" and info["device"]["grid"] == "12x10"
    assert info["autoware"]["package"] == "autoware_tensorrt_bevformer"
    assert info["output"]["labels"] == list(BEVFormer.LABELS)
    assert info["input"]["camera_order"] == CAMS
    assert "synthetic_6cam" in info["input"]["calibration_presets"]
    assert info["variant"] == BEVFormer.DEFAULT_VARIANT
    assert client.get("/v1/models").json()["data"][0]["id"] == "changh95/bevformer-small-weights"


@pytest.mark.parametrize("calibration", ["inline", "preset"])
def test_predict_ok_and_equals_api(client, calibration):
    cal = _inline_calib() if calibration == "inline" else {"preset": "synthetic_6cam"}
    pose = np.eye(4)
    pose[:3, 3] = [100.0, -20.0, 0.5]
    body = {"images": _images(), "calibration": cal, "params": {"score_threshold": 0.4},
            "stream": {"id": "cam-rig-1", "reset": False, "T_world_from_ego": pose.tolist()},
            "ego": {"velocity": [5.0, 0.0, 0.0]}}
    r = client.post("/predict", json=body)
    assert r.status_code == 200, r.text
    out = r.json()
    assert out["num_detections"] == 3 and out["frame_id"] == "LIDAR_TOP" and out["model"] == BEVFormer.MODEL_NAME
    assert set(out["detections"][0]) >= {"label", "label_id", "score", "center", "size", "yaw", "velocity"}
    assert {"decode", "model_call", "total", "device"} <= set(out["timing_ms"])
    state = server.app.state.ttaw
    req = server.BEVFormerRequest(**{k: v for k, v in body.items() if k != "params"})
    api = state.model(**decode_request(req, state.spec, 1 << 30)).to_dict()["detections"]
    assert state.predict(req)["detections"] == api


@pytest.mark.parametrize("payload,code", [
    ({"images": [{"camera": "CAM_FRONT", "data": "not base64!"}]}, 400),                  # undecodable
    ({"images": _images(CAMS[:5]), "calibration": _inline_calib()}, 400),                 # a camera missing
    ({"images": _images(), "calibration": _inline_calib(CAMS[:5])}, 400),                 # calibration missing
    ({"images": _images(), "calibration": {"preset": "../etc"}}, 400),                    # bad preset
    ({"images": _images(), "calibration": {"preset": "synthetic_6cam"}, "params": {"score_threshold": 2.0}}, 400),
    ({"images": _images(), "calibration": {"preset": "synthetic_6cam"}, "params": {"no_such_knob": 1}}, 400),
    ({"images": _images(), "calibration": {"preset": "synthetic_6cam"}, "params": {"max_num": 9001}}, 400),
    ({"images": _images(), "calibration": {"preset": "synthetic_6cam"}, "ego": {"jerk": [0, 0, 0]}}, 400),
    ({"images": _images(), "calibration": {"preset": "synthetic_6cam"}, "ego": {"velocity": [1, 2]}}, 400),
    ({}, 400),                                                                            # no images
    ({"imagez": []}, 422),                                                                # schema violation
])
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")