# SPDX-License-Identifier: Apache-2.0
"""Host tests of the bundle wiring (no device, no weights): the vendored ttaw copy is intact and every module
imports without ttnn / torch; the identity in api.py / device.py matches tt-model.yaml; the device configuration
resolves like the server's; the stdlib client and smoke test work against this bundle's app served over real HTTP
(ETH dispatch and 12x10 asserted, serve-profile pins, agreement with a stored reference output); the pip project is the
repo-root pyproject.toml; code/scripts/container_smoke.sh (run against fake tt-model / docker tools) keeps its evidence
and always stops the container.

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

import contextlib
import fnmatch
import hashlib
import json
import os
import shutil
import socket
import subprocess
import sys
import threading
from pathlib import Path
from typing import Iterator

import numpy as np
import pytest

import tt_bevformer
from tt_bevformer import device as tdev
from tt_bevformer.api import BEVFormer
from tt_bevformer.tests.stubs import ETH_P150, StubModel

PKG_DIR = Path(tt_bevformer.__file__).resolve().parent
CODE_DIR, BUNDLE_DIR, TTAW_DIR = PKG_DIR.parent, PKG_DIR.parents[1], PKG_DIR / "ttaw"
PATCH_MARKER = "single_chip_arch_1cq_no_dispatch_s"   # added to topology.cpp by patches/tt-metal-eth-dispatch.patch
VENDOR_IGNORE = ("__pycache__", "*.pyc", "*.pyo", ".pytest_cache", "*.bin", "*.log", ".DS_Store", "VENDORED.json")

BLOCKER = """
import importlib, importlib.abc, sys
BLOCKED = set(sys.argv[2].split(","))
class Block(importlib.abc.MetaPathFinder):
    def find_spec(self, name, path, target=None):
        if name.split(".")[0] in BLOCKED:
            raise ImportError("blocked import of " + name)
sys.meta_path.insert(0, Block())
sys.path.insert(0, sys.argv[1])
for target in sys.argv[3:]:
    module, _, attr = target.partition(":")
    mod = importlib.import_module(module)
    if attr:
        getattr(mod, attr)
leaked = sorted(k for k in sys.modules if k.split(".")[0] in BLOCKED)
print("ok" if not leaked else "leaked: " + " ".join(leaked))
"""


# ------------------------------------------------------------------------------------- shipped data and docs

def test_only_redistributable_samples_ship():
    """code/ (what the image and the Hub get) holds the synthetic sample generated by code/scripts/make_sample.py
    only: no PandaSet-derived file until the user approves PandaSet samples (research/DATASETS.md N5; they wait in
    the git-ignored staging_samples_pandaset/), no nuScenes / Argoverse 2 / METEOR / rosbag data ever, and no
    ``*.bin`` (tt-model's staging would drop it)."""
    shipped = [p.relative_to(CODE_DIR).as_posix() for p in CODE_DIR.rglob("*") if p.is_file()
               and "__pycache__" not in p.parts and ".pytest_cache" not in p.parts]
    banned = ("pandaset", "nuscenes", "av2_", "argoverse", "meteor", "rosbag", "sample0")
    data = [f for f in shipped if not f.endswith((".py", ".md", ".cpp", ".hpp", ".h", ".sh"))
            and not f.startswith("tt_bevformer/ttaw/")]
    assert not [f for f in data if any(b in Path(f).name.lower() for b in banned)]
    assert not [f for f in shipped if f.endswith(".bin")]
    samples = sorted(p.relative_to(PKG_DIR / "samples").as_posix() for p in (PKG_DIR / "samples").rglob("*")
                     if p.is_file())
    cams = sorted(f"synthetic_6cam/{c}.png" for c in BEVFormer.CAMERA_ORDER)
    assert samples == sorted(["README.md", "synthetic_6cam.json", "synthetic_6cam.reference.json"] + cams)
    assert sorted(p.name for p in (PKG_DIR / "calib").iterdir()) == ["README.md", "synthetic_6cam.json"]


def test_python_md_top_level_copy():
    """The repo's top-level PYTHON.md (published next to README.md) is a byte copy of code/PYTHON.md."""
    top = BUNDLE_DIR / "PYTHON.md"
    if not (BUNDLE_DIR / "tt-model.yaml").is_file():
        pytest.skip("not a bundle source checkout")
    assert top.is_file() and top.read_bytes() == (CODE_DIR / "PYTHON.md").read_bytes()


# ------------------------------------------------------------------------------------------ vendored ttaw

def test_vendored_ttaw_matches_its_manifest():
    """The copy under code/tt_bevformer/ttaw is exactly what common/tools/vendor.py wrote (never edit it here)."""
    manifest = json.loads((TTAW_DIR / "VENDORED.json").read_text())
    recorded = manifest["files"]
    actual = {}
    for p in sorted(TTAW_DIR.rglob("*")):
        rel = p.relative_to(TTAW_DIR)
        if p.is_file() and not any(fnmatch.fnmatch(part, pat) for part in rel.parts for pat in VENDOR_IGNORE):
            actual[rel.as_posix()] = hashlib.sha256(p.read_bytes()).hexdigest()
    problems = ([f"missing: {k}" for k in sorted(set(recorded) - set(actual))]
                + [f"extra: {k}" for k in sorted(set(actual) - set(recorded))]
                + [f"modified: {k}" for k in sorted(set(recorded) & set(actual)) if recorded[k] != actual[k]])
    assert problems == [], "re-vendor with common/tools/vendor.py instead of editing the copy"
    from tt_bevformer import ttaw

    assert manifest["package"] == "ttaw" and manifest["version"] == ttaw.__version__


@pytest.mark.parametrize("blocked,targets", [
    ("ttnn,torch", ["tt_bevformer:BEVFormer", "tt_bevformer:Output", "tt_bevformer:open_device",
                    "tt_bevformer:load_points", "tt_bevformer:load_image", "tt_bevformer:load_sample",
                    "tt_bevformer.device", "tt_bevformer.io", "tt_bevformer.server", "tt_bevformer.server.smoke_test",
                    "tt_bevformer.host.inputs", "tt_bevformer.host.calib", "tt_bevformer.host.geometry",
                    "tt_bevformer.host.temporal", "tt_bevformer.host.postprocess", "tt_bevformer.host.preprocess",
                    "tt_bevformer.reference", "tt_bevformer.reference.config", "tt_bevformer.reference.weights",
                    "tt_bevformer.reference.ort"]),
    # the torch half of the CPU reference also runs where ttnn is absent (research venv)
    ("ttnn", ["tt_bevformer.reference.ops", "tt_bevformer.reference.model", "tt_bevformer.reference.pipeline"]),
])
def test_imports_have_no_side_effects(blocked, targets):
    if blocked == "ttnn,torch":
        pytest.importorskip("fastapi")
        targets = targets + ["tt_bevformer.server.app:app"]
    out = subprocess.run([sys.executable, "-c", BLOCKER, str(CODE_DIR), blocked] + targets, capture_output=True,
                         text=True)
    assert out.returncode == 0 and out.stdout.strip() == "ok", out.stderr[-2000:]


# ----------------------------------------------------------------------------------- identity and config

@pytest.fixture(scope="module")
def manifest():
    yaml = pytest.importorskip("yaml")
    path = BUNDLE_DIR / "tt-model.yaml"
    if not path.is_file():
        pytest.skip("no tt-model.yaml next to code/ (installed package or container image)")
    return yaml.safe_load(path.read_text())


def test_identity_matches_manifest(manifest):
    cls = BEVFormer
    assert manifest["name"] == cls.MODEL_NAME and manifest["repo"].endswith("/" + cls.MODEL_NAME)
    w = manifest["weights"]
    assert (w["repo"], w["revision"]) == (cls.DEFAULT_REPO, cls.DEFAULT_REVISION)
    assert list(w.get("allow_patterns") or []) == list(cls.ALLOW_PATTERNS or [])
    assert manifest["runtime"]["app"] == "tt_bevformer.server.app:app"
    env = manifest["serve"]["env"]
    assert env["TT_WEIGHTS_REVISION"] == cls.DEFAULT_REVISION
    assert int(env["BEVFORMER_NUM_CQS"]) == tdev.DEVICE_DEFAULTS["num_command_queues"]
    for profile in [{"name": "serve", "env": {}}] + list(manifest.get("serve_profiles") or []):
        merged = {**env, **(profile.get("env") or {})}
        assert merged.get("BEVFORMER_DISPATCH") == "eth", f"{profile['name']}: the p150 target is ETH dispatch"
        assert merged.get("BEVFORMER_VARIANT") in cls.VARIANTS, profile["name"]
    verify = "\n".join(manifest["verify"])
    assert PATCH_MARKER in verify and cls.DEFAULT_REVISION in verify


def test_device_config_resolution(monkeypatch):
    for name in ("DISPATCH", "NUM_CQS", "L1_SMALL", "TRACE_REGION", "WORKER_L1_SIZE"):
        monkeypatch.delenv(f"{tdev.ENV_PREFIX}_{name}", raising=False)
    monkeypatch.delenv("TT_DEVICE_ID", raising=False)
    cfg = tdev.device_config()
    assert (cfg.dispatch, cfg.device_id, cfg.allow_fallback) == ("eth", 0, True)
    assert {k: getattr(cfg, k) for k in tdev.DEVICE_DEFAULTS} == tdev.DEVICE_DEFAULTS
    assert cfg == BEVFormer.device_config()
    monkeypatch.setenv("BEVFORMER_DISPATCH", "worker")
    monkeypatch.setenv("BEVFORMER_TRACE_REGION", str(1 << 20))
    monkeypatch.setenv("TT_DEVICE_ID", "1")
    cfg = tdev.device_config(num_command_queues=2, worker_l1_size=None)
    assert (cfg.dispatch, cfg.trace_region_size, cfg.device_id, cfg.num_command_queues) == ("worker", 1 << 20, 1, 2)
    assert cfg == BEVFormer.device_config(num_command_queues=2)
    assert tdev.device_config(dispatch="eth").dispatch == "eth"     # an explicit argument beats the environment


# ------------------------------------------------------------------------------- client and smoke test

def test_client_script_runs_standalone(tmp_path):
    """server/client.py runs with any python3 in isolated mode: no package install, no numpy needed."""
    images = []
    for cam in BEVFormer.CAMERA_ORDER:
        f = tmp_path / f"{cam}.npy"
        np.save(f, np.zeros((9, 16, 3), np.uint8))
        images += ["--image", f"{cam}={f}"]
    subprocess.run([sys.executable, "-I", str(PKG_DIR / "server" / "client.py"), *images, "--calib-preset",
                    "synthetic_6cam", "--param", "score_threshold=0.5", "--out", str(tmp_path / "req.json")],
                   check=True, capture_output=True, text=True)
    req = json.loads((tmp_path / "req.json").read_text())
    assert [i["camera"] for i in req["images"]] == list(BEVFormer.CAMERA_ORDER)
    assert req["calibration"] == {"preset": "synthetic_6cam"} and req["params"] == {"score_threshold": 0.5}


class WorkerStub(StubModel):
    """A server whose ETH open fell back to WORKER dispatch."""

    device_info = dict(ETH_P150, dispatch="worker", fallback="RuntimeError: no ETH", grid="11x10", grid_x=11,
                       cores=110)


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


@contextlib.contextmanager
def _serve(stub_cls) -> Iterator[str]:
    """This bundle's real app under uvicorn on a free port with ``stub_cls`` as the model; yields the base URL."""
    uvicorn = pytest.importorskip("uvicorn")
    pytest.importorskip("fastapi")
    from tt_bevformer.server import app as server
    from tt_bevformer.server import smoke_test

    state = server.app.state.ttaw
    saved, state.model_factory = state.model_factory, stub_cls
    srv = uvicorn.Server(uvicorn.Config(server.app, host="127.0.0.1", port=_free_port(), log_level="warning",
                                        lifespan="on"))
    thread = threading.Thread(target=srv.run, daemon=True)
    thread.start()
    base = f"http://127.0.0.1:{srv.config.port}"
    try:
        assert smoke_test.ttaw_client.wait_ready(base, wait_s=30, poll_s=0.1)["status"] == "ok"
        yield base
    finally:
        srv.should_exit = True
        thread.join(timeout=30)
        state.model_factory = saved


def _staged_manifest(tmp_path: Path) -> Path:
    """A staged tt_kernel_manifest.json (wire format) with a default profile and one whose pins differ."""
    env = {"TT_WEIGHTS_REVISION": BEVFormer.DEFAULT_REVISION, "BEVFORMER_DISPATCH": "eth",
           "BEVFORMER_NUM_CQS": str(tdev.DEVICE_DEFAULTS["num_command_queues"]),
           "BEVFORMER_VARIANT": BEVFormer.DEFAULT_VARIANT}
    wire = {"name": BEVFormer.MODEL_NAME,
            "weights": {"repo_id": BEVFormer.DEFAULT_REPO, "revision": BEVFormer.DEFAULT_REVISION},
            "container": {"serve": {"env": env}, "default_profile": None,
                          "serve_profiles": [{"name": "default", "env": {}},
                                             {"name": "other", "env": {"BEVFORMER_VARIANT": "not-a-variant"}}]}}
    path = tmp_path / "tt_kernel_manifest.json"
    path.write_text(json.dumps(wire))
    return path


def test_smoke_test_against_live_server(tmp_path, monkeypatch, capsys):
    from tt_bevformer.server import smoke_test

    monkeypatch.setenv("TT_MESH_SHAPE", "1x1")
    monkeypatch.setenv("TT_MODEL_WEIGHTS_REVISION", BEVFormer.DEFAULT_REVISION)
    for name in ("HF_MODEL", "BEVFORMER_VARIANT", "BEVFORMER_DISPATCH", "BEVFORMER_NUM_CQS"):
        monkeypatch.delenv(name, raising=False)
    shipped = PKG_DIR / "samples" / "synthetic_6cam.json"           # six camera files + the synthetic_6cam preset
    spec = json.loads(shipped.read_text())
    spec["images"] = {cam: str((shipped.parent / rel).resolve()) for cam, rel in spec["images"].items()}
    sample = tmp_path / "sample.json"
    sample.write_text(json.dumps(spec))
    from tt_bevformer.host.inputs import load_sample

    reference = StubModel()(**load_sample(sample)).to_dict()          # the stub plays the CPU reference
    (tmp_path / "sample.reference.json").write_text(json.dumps(reference))
    shifted = json.loads(json.dumps(reference))
    for d in shifted["detections"]:
        d["center"][0] += 2.0
    (tmp_path / "shifted.json").write_text(json.dumps(shifted))
    args = ["--input", str(sample), "--expect", "", "--out", str(tmp_path / "out.json")]
    pinned = args + ["--manifest", str(_staged_manifest(tmp_path))]

    with _serve(StubModel) as url:
        assert smoke_test.main(["--url", url] + pinned) == 0
        out = capsys.readouterr().out
        assert out.startswith("PASS bevformer-p150: profile=default") and "dispatch=eth grid=12x10" in out
        assert "reference=sample.reference.json (detections 3/3 recall=1.000" in out
        assert json.loads((tmp_path / "out.json").read_text())["num_detections"] == 3
        assert smoke_test.main(["--url", url, "--profile", "other"] + pinned) == 1
        assert "profile 'other' pins BEVFORMER_VARIANT=not-a-variant" in capsys.readouterr().out
        assert smoke_test.main(["--url", url, "--reference", str(tmp_path / "shifted.json")] + args) == 1
        assert "detections recall 0.000 < 0.95" in capsys.readouterr().out
        assert smoke_test.main(["--url", url, "--reference", str(tmp_path / "missing.json")] + args) == 1
        assert "reference " in capsys.readouterr().out
        assert smoke_test.main(["--url", url, "--input", str(tmp_path / "missing.npy"), "--expect", ""]) == 1
        assert "/predict: FileNotFoundError" in capsys.readouterr().out    # a FAIL line, not a traceback
    with _serve(WorkerStub) as url:
        assert smoke_test.main(["--url", url] + args) == 1
        out = capsys.readouterr().out
        assert out.startswith("FAIL bevformer-p150:")
        assert "dispatch is 'worker', expected 'eth'" in out and "grid is '11x10', expected '12x10'" in out


# ------------------------------------------------------------------------- pip project and container smoke

def test_pip_project_at_repo_root():
    """The pip project is the repo-root pyproject.toml over code/ (never code/pyproject.toml: tt-model copies code/
    over the tt-metal tree before building ttnn, PACKAGING_PILOT.md problem 1)."""
    tomllib = pytest.importorskip("tomllib")
    path = BUNDLE_DIR / "pyproject.toml"
    if not (BUNDLE_DIR / "tt-model.yaml").is_file():
        pytest.skip("not a bundle source checkout (installed package or container image)")
    assert not (CODE_DIR / "pyproject.toml").exists(), "the pip project lives at the repo root, never in code/"
    project = tomllib.loads(path.read_text())
    find = project["tool"]["setuptools"]["packages"]["find"]
    assert find["where"] == ["code"] and any(fnmatch.fnmatch("tt_bevformer", pat) for pat in find["include"])
    data = project["tool"]["setuptools"]["package-data"]
    assert {"API.md", "VENDORED.json"} <= set(data["tt_bevformer.ttaw"])
    assert {"kernels/*.cpp", "kernels/*.hpp", "kernels/*.h"} <= set(data["tt_bevformer.ttaw.ops"])
    assert (BUNDLE_DIR / project["project"]["readme"]).is_file()
    assert project["tool"]["setuptools"]["dynamic"]["version"]["attr"] == "tt_bevformer.__version__"
    yaml = pytest.importorskip("yaml")
    shipped = [p for e in yaml.safe_load((BUNDLE_DIR / "tt-model.yaml").read_text())["source"]["extra_code"]
               for p in e["paths"]]
    assert "pyproject.toml" not in shipped and "setup.py" not in shipped
    setuptools = pytest.importorskip("setuptools")
    found = set(setuptools.find_packages(where=str(CODE_DIR), include=find["include"], exclude=find.get("exclude", [])))
    assert {"tt_bevformer", "tt_bevformer.ttaw", "tt_bevformer.ttaw.server", "tt_bevformer.ttaw.ops"} <= found
    assert not any(p.startswith("tt_bevformer.tests") for p in found)


FAKE_TT_MODEL = r'''
import json, os, signal, subprocess, sys, time, urllib.request
from pathlib import Path
state = Path(os.environ["FAKE_STATE"])
args = sys.argv[1:]
with open(state / "calls.txt", "a") as f:
    f.write(" ".join(args) + "\n")
opt = lambda name: args[args.index(name) + 1] if name in args else None
if args[0] == "serve":
    port, manifest = int(opt("--port")), json.load(open(args[-1]))
    (state / "container").write_text("tt-model-" + manifest["name"] + "-" + (opt("--profile") or "default"))
    (state / "container.log").write_text("boot: Opening device 0 (dispatch eth, 1 CQ)\n")
    if os.environ.get("FAKE_SERVE_FAIL"):
        sys.exit("fake serve: the boot failed")
    srv = subprocess.Popen([sys.executable, str(state / "server.py"), str(port)], start_new_session=True)
    (state / "server.pid").write_text(str(srv.pid))
    for _ in range(200):
        try:
            urllib.request.urlopen("http://127.0.0.1:%d/health" % port, timeout=1)
            break
        except OSError:
            time.sleep(0.05)
elif args[0] == "stop":
    if (state / "server.pid").exists():
        os.kill(int((state / "server.pid").read_text()), signal.SIGTERM)
    with open(state / "container.log", "a") as f:
        f.write("shutdown: traces released, device closed\n")
    (state / "stopped").write_text("1")
elif args[0] == "logs":
    print("fallback: tt-model logs")
else:
    sys.exit(2)
'''

FAKE_DOCKER = r'''
import os, sys, time
from pathlib import Path
state = Path(os.environ["FAKE_STATE"])
args = sys.argv[1:]
with open(state / "docker_calls.txt", "a") as f:
    f.write(" ".join(args) + "\n")
known = (state / "container").read_text() if (state / "container").exists() else None
if args[0] == "inspect":
    sys.exit(0 if args[-1] == known else 1)
if args[0] == "logs" and args[-1] == known:
    deadline = time.time() + 20
    while "--follow" in args and not (state / "stopped").exists() and time.time() < deadline:
        time.sleep(0.05)
    sys.stdout.write((state / "container.log").read_text())
    sys.exit(0)
sys.exit(1)
'''

FAKE_SERVER = r'''
import json, sys
from http.server import BaseHTTPRequestHandler, HTTPServer
INFO = {"model": "fake", "device": {"dispatch": "eth", "grid": "12x10", "cores": 120}}
class Handler(BaseHTTPRequestHandler):
    def _send(self, code, body):
        data = json.dumps(body).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)
    def do_GET(self):
        self._send(200, {"status": "ok"} if self.path.startswith("/health") else INFO if self.path == "/info" else {})
    def do_POST(self):
        self._send(500, {"detail": "fake server"})
    def log_message(self, *args):
        pass
HTTPServer(("127.0.0.1", int(sys.argv[1])), Handler).serve_forever()
'''


def _fake_tools(tmp_path: Path) -> dict:
    """PATH with fake `tt-model` / `docker` / `python3` (the test interpreter) and a stub server: no docker, no chip."""
    state, bin_dir = tmp_path / "state", tmp_path / "bin"
    state.mkdir()
    bin_dir.mkdir()
    (state / "server.py").write_text(FAKE_SERVER)
    for name, body in (("tt-model", FAKE_TT_MODEL), ("docker", FAKE_DOCKER)):
        (bin_dir / name).write_text(f"#!{sys.executable}\n{body}")
        (bin_dir / name).chmod(0o755)
    (bin_dir / "python3").write_text(f'#!/bin/sh\nexec "{sys.executable}" "$@"\n')
    (bin_dir / "python3").chmod(0o755)
    env = {k: v for k, v in os.environ.items() if k not in ("SMOKE_OUT", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN")}
    env.update(PATH=f"{bin_dir}:{env.get('PATH', '/usr/bin:/bin')}", FAKE_STATE=str(state))
    for tool in ("tt-model", "docker", "python3"):   # never the real tools: a real serve would claim the chip
        assert shutil.which(tool, path=env["PATH"]) == str(bin_dir / tool)
    return env


@pytest.mark.parametrize("serve_fails", [False, True])
def test_container_smoke_keeps_evidence_and_stops(tmp_path, serve_fails):
    """code/scripts/container_smoke.sh passes --profile through, saves /info once READY and the container log through
    the shutdown (whatever the outcome), and always stops the container."""
    staged = tmp_path / "staged"
    staged.mkdir()
    manifest = _staged_manifest(staged)
    env, logs, port = _fake_tools(tmp_path), tmp_path / "evidence", _free_port()
    profile = [] if serve_fails else ["--profile", "other"]
    if serve_fails:
        env["FAKE_SERVE_FAIL"] = "1"
    r = subprocess.run(["bash", str(CODE_DIR / "scripts" / "container_smoke.sh"), str(staged), str(port), *profile,
                        f"--log-dir={logs}"], env=env, capture_output=True, text=True, timeout=180)
    assert r.returncode == 1, r.stdout + r.stderr          # the stub serves no valid /predict: the smoke FAILS
    calls = (tmp_path / "state" / "calls.txt").read_text().splitlines()
    assert calls[0] == " ".join(["serve", "--port", str(port), *profile, str(manifest)])
    assert calls[-1] == " ".join(["stop", *profile, str(manifest)]) and len(calls) == 2
    stem = "bevformer-p150" + ("" if serve_fails else "-other") + "-"
    files = {p.name.split(".", 1)[1]: p for p in logs.iterdir()}
    assert all(p.name.startswith(stem) for p in files.values()), sorted(files)
    log = files["container.log"].read_text()
    assert "boot: Opening device 0" in log and "shutdown: traces released" in log   # followed through the stop
    result = json.loads(files["result.json"].read_text())
    assert (result["rc"], result["result"], result["profile"]) == (1, "FAIL", "default" if serve_fails else "other")
    if serve_fails:
        assert "info.json" not in files
    else:
        assert json.loads(files["info.json"].read_text())["device"] == {"dispatch": "eth", "grid": "12x10",
                                                                        "cores": 120}
        assert "FAIL bevformer-p150:" in r.stdout