# SPDX-License-Identifier: Apache-2.0
"""Host tests of the API's host half (``BEVFormer._build_host`` / ``_prepare``, no device, no weights): the runner
inputs a frame of the shipped sample produces equal the golden network inputs and host geometry (``make_goldens.py``,
port form), the stream history semantics (new id / reset -> ``use_prev_bev = 0``), the per-calibration cache, the
SCA-capacity refusal (HTTP 400) and the can_bus layout knob."""
from __future__ import annotations

import numpy as np
import pytest

from tt_bevformer import io as tio
from tt_bevformer.api import BEVFormer
from tt_bevformer.reference import config as C
from tt_bevformer.tests import paths

SAMPLES = [paths.PANDASET_SAMPLES / f"pandaset_019_f{f}.json" for f in paths.SAMPLE_FRAMES]


def _model(**compile_params) -> BEVFormer:
    m = BEVFormer.__new__(BEVFormer)
    m.compile_params = dict(compile_params)
    m._build_host()
    return m


def _kw(i: int, sid: str = "s", reset=None):
    paths.need_pandaset()
    from tt_bevformer.host.inputs import load_sample

    kw = load_sample(SAMPLES[i])
    st = dict(kw["stream"], id=sid)
    if reset is not None:
        st["reset"] = reset
    kw["stream"] = st
    return kw


@pytest.mark.parametrize("pre", ["numpy", "numba", "bf16", "numpy-s2d", "numba-s2d", "bf16-s2d"])
def test_frame_inputs_equal_the_golden_inputs(pre):
    """f30 (reset) then f35 as one stream: can_bus, use_prev_bev, the rotation map, the TSA shift and every
    per-calibration table equal what the golden run fed the port-form reference (both host pre-processing paths,
    BEVFORMER_HOST_PRE: the network input bit for bit)."""
    from tt_bevformer.host import fast
    from tt_bevformer.tt.model import calibration_inputs
    from tt_bevformer.tt.params import tsa_shift_rows

    from tt_bevformer.tt.params import stem_s2d_image

    if pre.startswith(("numba", "bf16")) and not fast.available():
        pytest.skip("numba is not installed")
    s2d = pre.endswith("-s2d")
    m = _model(HOST_PRE=pre.split("-")[0], STEM="s2d" if s2d else "plain", UPLOAD_PIPE=False, STEM_IN="kin")
    for i, f in ((0, 30), (1, 35)):
        p = m._prepare(**_kw(i, reset=(i == 0)))
        gp = paths.golden_dir() / "pandaset_019" / f"f{f}.npz"
        if not gp.is_file():
            pytest.skip(f"{gp} missing")
        with np.load(gp) as z:
            g = {k: z[k] for k in z.files if k != "__meta__"}
        up = p["upload"]
        np.testing.assert_array_equal(up["params"]["can_bus"].reshape(18), g["input.can_bus"])
        assert up["params"]["use"] == float(g["input.use_prev_bev"][0]) == (0.0 if i == 0 else 1.0)
        np.testing.assert_array_equal(up["inputs"]["rot_index"].reshape(-1), g["geometry.rot_index"])
        np.testing.assert_array_equal(up["params"]["tsa_shift"], tsa_shift_rows(g["geometry.shift"],
                                                                                 up["params"]["use"]))
        from tt_bevformer.tt.model import img_shape

        assert up["inputs"]["img"].shape == img_shape("s2d" if s2d else "plain", m.knobs.UPLOAD == "rows")
        if "input.image" in g:
            want = np.ascontiguousarray(g["input.image"].reshape(C.NUM_CAMS, 3, *C.IMG_HW).transpose(0, 2, 3, 1))
            if s2d:
                want = stem_s2d_image(want)
            got = up["inputs"]["img"]
            if pre == "bf16":                    # bf16 rows: the golden input cast as the upload casts it
                import torch

                got = got.reshape(want.shape)
                want = torch.from_numpy(np.ascontiguousarray(want, np.float32)).to(torch.bfloat16)
                assert got.dtype == torch.bfloat16 and torch.equal(got.view(torch.int16), want.view(torch.int16))
            else:
                np.testing.assert_array_equal(got.reshape(want.shape), want)
        geo = {k[len("geometry."):]: v for k, v in g.items() if k.startswith("geometry.")}
        want_cal = calibration_inputs(geo)
        for k, v in p["geometry"]["inputs"].items():
            np.testing.assert_array_equal(v, want_cal[k], err_msg=k)


def test_stream_history_semantics():
    m = _model()
    assert m._prepare(**_kw(0, "a"))["fresh"]                         # new id
    p = m._prepare(**_kw(1, "a", reset=False))
    assert not p["fresh"] and p["use_prev_bev"][0] == 1.0
    assert m._prepare(**_kw(2, "a", reset=True))["fresh"]             # explicit reset
    assert m._prepare(**_kw(1, "b", reset=False))["fresh"]            # another id starts fresh
    with pytest.raises(tio.InputError):
        m._prepare(**dict(_kw(1, "a"), stream={"id": "a", "reset": "yes"}))


def test_calibration_cache():
    m = _model()
    a = m._prepare(**_kw(0))["geometry"]
    b = m._prepare(**_kw(0))["geometry"]
    assert a is b                                                     # same lidar2img -> same tables
    c = m._prepare(**_kw(1))["geometry"]                              # per-frame compensated calibration
    assert c["key"] != a["key"]


def test_sca_overflow_is_a_client_error(monkeypatch):
    from tt_bevformer.host import geometry as hgeo

    def overflow(*_a, **_k):
        raise hgeo.SCAOverflow("a camera sees 7000 BEV queries")

    monkeypatch.setattr(hgeo, "sca_tables", overflow)
    with pytest.raises(tio.InputError, match="calibration not supported"):
        _model()._prepare(**_kw(0))


def test_can_bus_layout_knob():
    t = _model()._prepare(**_kw(0))["can_bus"]
    a = _model(can_bus_layout="autoware")._prepare(**_kw(0))["can_bus"]
    np.testing.assert_array_equal(t[7:10], a[13:16])                  # accel <-> velocity slots
    np.testing.assert_array_equal(t[13:16], a[7:10])
    np.testing.assert_array_equal(t[10:13], a[10:13])


def test_knobs_defaults():
    from tt_bevformer.api import KNOBS

    k = KNOBS.read(env={})
    assert k.PRECISION == "HiFi2+fp32+l1acc:w=bf16" and k.MAX_STREAMS == 4 and k.SHIFT_MODE == "surrogate"
    assert k.CAN_BUS_LAYOUT == "training" and k.MSDA == "kernel" and k.DCN == "kernel" and k.DCN_OFFSET_L1
    assert k.MSDA_PX and k.MSDA_VALUE == "tile" and k.HOST_PRE == "bf16" and k.BACKBONE_MM == "exact2"
    assert k.UPLOAD == "rows" and k.STEM == "plain" and not k.MSDA_SKIP and k.STEM_IN == "kpre"
    assert k.MSDA_LOC == "compact" and k.MSDA_SOFTMAX == "kernel" and k.SCA_GATHER == "tile" and k.BEV_MM == "fast" and k.DEC_ATTN == "sdpa"
    assert k.PRECISION_RULES.endswith("img_backbone.layer3.*=HiFi2+fp32+l1acc:w=bfp8") and k.RES3_NEXT and k.ACT8 and k.ACT8_L3
    assert k.DCN_IM2COL == "l1" and k.DCN_SCALE == "full" and k.BACKBONE_S2 == "prod" and k.S2_SKIP and k.MSDA_FRAC16 and k.DCN_SAMPLE == "bilinear" and k.SCA_COPY == "kernel" and k.DEC_VALUE == "fused" and k.BACKBONE_C3 == "l1_2"
    # D20 round 2 (all bit-identical data movement)
    assert k.MMRES_MCAST and k.MSDA_MK_MCAST and k.DCN_GLUE == "kernel" and k.RES3_RM == "on" and k.DCN_OUT == "l1"
    assert k.PAIRMAP == "v2" and k.ENC_L1 == "all" and k.NECK_L1 and k.DCN_SUM == "mul" and not k.DCN_PAIRS
    assert k.HOST_MALLOC == "keep" and k.RES_L1 and k.OUTPUTS == "tile" and k.MSDA_PAIRS == "all"
    assert k.VALUES_DIRECT and k.SCA_VALUE == "l1" and k.HOST_STAGING and k.PIPELINE == 1


def test_upload_pipe_uploads_every_camera_bit_exact():
    """BEVFORMER_UPLOAD_PIPE: _prepare leaves the images pending (no img input), _forward's camera pipeline uploads
    img0 .. img5 = the bf16 rows of HOST_PRE=bf16, each exactly once, bit for bit."""
    import torch

    from tt_bevformer.host import fast
    from tt_bevformer.tt.model import CAM_ROWS_SHAPE

    if not fast.available():
        pytest.skip("numba is not installed")
    m = _model(HOST_PRE="bf16", UPLOAD_PIPE=True, STEM_IN="kin")
    p = m._prepare(**_kw(0, reset=True))
    assert p["pending_images"] is not None and not any(k.startswith("img") for k in p["upload"]["inputs"])
    want = fast.preprocess_rows_bf16(p["pending_images"], rgb=True, threads=2).copy()

    class _Runner:
        def __init__(self):
            self.got = {}

        def upload(self, inputs=None, params=None):
            for k, v in (inputs or {}).items():
                assert k not in self.got
                self.got[k] = v.clone()

    m.runner = _Runner()
    m._upload_cameras(p["pending_images"])
    assert sorted(m.runner.got) == [f"img{i}" for i in range(C.NUM_CAMS)]
    for i in range(C.NUM_CAMS):
        t = m.runner.got[f"img{i}"]
        assert t.dtype == torch.bfloat16 and tuple(t.shape) == CAM_ROWS_SHAPE
        assert np.array_equal(t.view(torch.int16).numpy().view(np.uint16).reshape(want[i].shape), want[i])
    assert not _model(HOST_PRE="numba", UPLOAD_PIPE=True, STEM_IN="kin")._upload_pipe()
    assert not _model(UPLOAD_PIPE=True)._upload_pipe()                 # the default STEM_IN=kpre uploads uint8


def test_kpre_uploads_the_node_images():
    """BEVFORMER_STEM_IN=kpre: the frame's inputs are the six uint8 node images (img0 .. img5, [1, 1, 900, 4800]),
    the node resize of the 1920x1080 sample equal to the numpy node transport; no float rows on the host."""
    import torch

    from tt_bevformer.host import fast
    from tt_bevformer.host import preprocess as hp
    from tt_bevformer.tt.model import NODE_U8_SHAPE

    if not fast.available():
        pytest.skip("numba is not installed")
    m = _model(STEM_IN="kpre")
    assert m._kpre() and not m._upload_pipe()
    kw = _kw(0, reset=True)
    p = m._prepare(**kw)
    ins = p["upload"]["inputs"]
    assert "img" not in ins and sorted(k for k in ins if k.startswith("img")) == [f"img{i}" for i in range(6)]
    from tt_bevformer.host.inputs import prepare_request

    frame = prepare_request(kw["images"], kw.get("calibration"), kw.get("stream"), kw.get("ego"), bgr=True)
    for i in range(C.NUM_CAMS):
        t = ins[f"img{i}"]
        assert t.dtype == torch.uint8 and tuple(t.shape) == NODE_U8_SHAPE
        want = hp.node_transport(frame.images_bgr[i])[:, :, ::-1]          # BGR -> the RGB the kernel reads
        assert np.array_equal(t.numpy().reshape(want.shape), want)


@pytest.mark.parametrize("mode", [1, 2])
def test_pipeline_orders_frames_and_survives_failures(mode, monkeypatch):
    """BEVFORMER_PIPELINE: concurrent calls take the device stage in the order their host half finished, the CQ1
    camera upload of a frame (mode 2) comes after the previous frame's image replay was enqueued, the stream half
    runs in the device stage, and a frame that fails in any stage gives up its turn (no deadlock)."""
    import random
    import threading
    import time
    import types

    m = _model(PIPELINE=mode)
    assert m.pipelined and m.pipeline == mode
    m._lock = threading.RLock()
    m._closed = False
    log, pre_order = [], []
    log_lock = threading.Lock()

    def note(*ev):
        with log_lock:
            log.append(ev)

    def prepare_host(points=None, **kw):
        time.sleep(random.uniform(0, 0.004))
        i = kw["stream"]["i"]
        pre_order.append(i)
        if i == 7:
            raise tio.InputError("bad frame 7")
        return {"i": i, "node_u8": None}

    monkeypatch.setattr(m, "_prepare_host", prepare_host)
    monkeypatch.setattr(m, "_host_images", lambda host: {"img0": host["i"]})

    def upload(imgs):
        note("upload", imgs["img0"])
        if imgs["img0"] == 11:
            raise RuntimeError("upload 11")
        return ("ev", imgs["img0"])

    monkeypatch.setattr(m, "_upload_images_cq1", upload)

    busy = []

    def device_stage(host, imgs, ev, ticket):
        i = host["i"]
        assert (ev == ("ev", i)) if mode == 2 else ev is None
        busy.append(i)
        overlap.extend(busy if len(busy) > 1 else [])
        note("image", i)
        with m._cv:                                  # what _device_stage's enqueue section implies
            m._img_turn = max(m._img_turn, m._ticket_of[i])
            m._cv.notify_all()
        time.sleep(random.uniform(0, 0.003))
        if i == 10:
            time.sleep(0.03)                         # frame 11 fails its CQ1 upload while 10 is still on the device
        busy.remove(i)
        if i == 13:
            raise RuntimeError("device 13")
        note("done", i)
        return {"i": i}, {"timing_ms": {}}

    overlap = []
    monkeypatch.setattr(m, "_device_stage", device_stage)
    m._ticket_of = {}
    orig_host = m._prepare_host

    def tracked(points=None, **kw):
        h = orig_host(points, **kw)
        m._ticket_of[h["i"]] = m._ticket
        return h

    monkeypatch.setattr(m, "_prepare_host", tracked)
    monkeypatch.setattr(m, "_postprocess", lambda raw, prep, p: types.SimpleNamespace(timing_ms={}))
    errors = {}

    def call(i):
        try:
            m(images=None, stream={"i": i})
        except Exception as e:  # noqa: BLE001
            errors[i] = type(e).__name__

    threads = [threading.Thread(target=call, args=(i,)) for i in range(24)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(10)
    assert not any(t.is_alive() for t in threads), "pipeline deadlocked"
    assert not overlap, f"device stages overlapped: {overlap}"   # a failed frame must not skip the turn order
    assert errors.get(7) == "InputError" and errors.get(13) == "RuntimeError"
    if mode == 2:
        assert errors.get(11) == "RuntimeError"
    ok = [i for i in pre_order if i not in errors]
    assert [e[1] for e in log if e[0] == "done"] == [i for i in ok]      # device order == host-half order
    if mode == 2:                                    # frame k's upload after frame k - 1's image replay
        pos = {e: n for n, e in enumerate(log)}
        seq = [i for i in pre_order if i != 7]
        for a, b in zip(seq, seq[1:]):
            if ("image", a) in pos:
                assert pos[("upload", b)] > pos[("image", a)]
    assert m._turn == m._ticket and m._img_turn == m._ticket - 1


def test_tt_model_kwargs_cover_every_device_option():
    """``api.tt_model_kwargs`` (shared by the API and the device PCC gate fixture) sets every TtBEVFormer option
    except the test / runtime-only ones, so a new device knob cannot silently fall back to the constructor default in
    the gates (verification 2026-10-13 FAIL-2)."""
    import inspect

    from tt_bevformer.api import KNOBS, tt_model_kwargs
    from tt_bevformer.tt.model import TtBEVFormer

    sig = inspect.signature(TtBEVFormer.__init__)
    options = {n for n, p in sig.parameters.items() if p.kind == p.KEYWORD_ONLY}
    kw = tt_model_kwargs(KNOBS.read())
    assert set(kw) <= options
    assert options - set(kw) == {"num_command_queues", "debug", "image_taps"}
    assert kw["stem_in"] == ("kin" if str(KNOBS.read().STEM_IN) == "kpre" else str(KNOBS.read().STEM_IN))
    assert tt_model_kwargs(KNOBS.read(), stem_in="kpre")["stem_in"] == "kpre"


def test_precision_rules_resolve(monkeypatch):
    """BEVFORMER_PRECISION_RULES (PLAN D20): first matching glob wins, unmatched modules keep the default."""
    from tt_bevformer.tt import kconf

    monkeypatch.setattr(kconf, "PREC_RULES", "*.conv_offset=HiFi4+fp32+l1acc:w=fp32;"
                                             "img_backbone.layer3.*=LoFi+fp32+l1acc:w=bfp8")
    assert kconf.prec("img_backbone.layer3.0.conv2.conv_offset", "x") == "HiFi4+fp32+l1acc:w=fp32"
    assert kconf.prec("img_backbone.layer3.4.conv1", "x") == "LoFi+fp32+l1acc:w=bfp8"
    assert kconf.prec("encoder.layers.0.ffn.0", "x") == "x"
    for off in ("", "none"):
        monkeypatch.setattr(kconf, "PREC_RULES", off)
        assert kconf.prec("img_backbone.layer3.4.conv1", "x") == "x"
    monkeypatch.setattr(kconf, "PREC_RULES", "broken")
    with pytest.raises(ValueError):
        kconf.prec("a", "x")