"""NCA-003: startup fails closed — no registration without a fresh matching proof.

Two layers are covered here:

* the gate itself (`meshnet_node.admission.admit`) — which reports admit a
  selection, and which are refused as failed, stale, or about something else;
* `run_startup` — that a refused report means the tracker is never called, and
  that an admitted one travels with the registration payload.

Torch is a stub in the dev venv, so the backend is faked by duck-typing
`TorchModelShard` (see `_FakeBackend`); the *production* validator still runs a
real `doctor` forward against it, so the fail-closed path is exercised without a
bypass. Tests that cannot supply an executable backend pass the explicit
test-safe validator from `meshnet_node.testing`.
"""

import base64
import struct
import time

import pytest

from meshnet_node.admission import (
    REASON_BACKEND_MISMATCH,
    REASON_MODEL_MISMATCH,
    REASON_NO_REPORT,
    REASON_NOT_PASSED,
    REASON_RECIPE_MISMATCH,
    REASON_SHARD_MISMATCH,
    REASON_STALE,
    AdmissionRequirement,
    CapabilityAdmissionError,
    CapabilityContext,
    admit,
    probe_capability,
)
from meshnet_node.capability import STATUS_FAILED, STATUS_SKIPPED
from meshnet_node.doctor import DoctorSelection
from meshnet_node.recipe_manifest import DEFAULT_RECIPE_ID, load_recipe_manifest
from meshnet_node.startup import run_startup
from meshnet_node.testing import capability_report_for, capability_stub

MODEL = "acme/opaque-model-7b"


class _FakeDevice:
    def __init__(self, type_: str = "cpu"):
        self.type = type_


class _FakeOutput:
    def __init__(self, hidden_size: int, tokens: int = 4):
        self.body = b"\x00" * (tokens * hidden_size * 2)
        self.shape = [1, tokens, hidden_size]
        self.attention_mask_header = _int64_header([[1] * tokens])
        self.position_ids_header = _int64_header([list(range(tokens))])


def _int64_header(rows):
    flat = [int(v) for row in rows for v in row]
    raw = struct.pack(f"<{len(flat)}q", *flat)
    return f"{len(rows)},{len(rows[0])}:{base64.b64encode(raw).decode('ascii')}"


class _FakeBackend:
    """Duck-types the parts of `TorchModelShard` the doctor probe touches."""

    total_layers = 24
    hidden_size = 8

    def __init__(self, *, shard_start=0, shard_end=23, device="cpu", forward_error=None):
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.is_head = shard_start == 0
        self.is_tail = shard_end == self.total_layers - 1
        self.device = _FakeDevice(device)
        self.model_id = MODEL
        self._forward_error = forward_error

    def encode_prompt(self, _prompt):
        if self._forward_error:
            raise self._forward_error
        return _FakeOutput(self.hidden_size)

    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
        if self._forward_error:
            raise self._forward_error
        return _FakeOutput(self.hidden_size)


def _context(backend=None, *, model_id=MODEL, shard_start=0, shard_end=23, device="cpu"):
    manifest = load_recipe_manifest()
    return CapabilityContext(
        backend=backend,
        selection=DoctorSelection(
            model_id=model_id,
            shard_start=shard_start,
            shard_end=shard_end,
            quantization="bfloat16",
        ),
        recipe=manifest.require(DEFAULT_RECIPE_ID),
        manifest=manifest,
        device=device,
    )


# ---------------------------------------------------------------------------
# The gate: which reports admit a selection
# ---------------------------------------------------------------------------


def test_a_fresh_matching_passing_report_admits_the_selection():
    "The proof covers exactly what is about to be advertised, so the node may register.\n\nTags: node, admission"
    ctx = _context()
    report = capability_report_for(ctx)

    assert admit(AdmissionRequirement.for_context(ctx), report) is report


def test_a_missing_report_is_refused():
    "No proof at all is the default state, and it must not register.\n\nTags: node, admission"
    ctx = _context()

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), None)

    assert exc.value.reason == REASON_NO_REPORT


@pytest.mark.parametrize("status", [STATUS_FAILED, STATUS_SKIPPED])
def test_a_report_that_did_not_pass_is_refused(status):
    "A failed or skipped validation is evidence against admission, not for it.\n\nTags: node, admission"
    ctx = _context()
    report = capability_report_for(
        ctx, status=status, diagnostics=["the shard forward returned no output"]
    )

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_NOT_PASSED
    assert "the shard forward returned no output" in str(exc.value)


def test_a_passing_report_for_another_model_is_refused():
    "A proof about one model says nothing about another — no reuse across models.\n\nTags: node, admission"
    ctx = _context()
    report = capability_report_for(ctx, model_id="other/model-1b")

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_MODEL_MISMATCH
    assert "other/model-1b" in str(exc.value)


def test_a_passing_report_for_another_shard_range_is_refused():
    "Layers 0–11 running is no proof that layers 12–23 fit.\n\nTags: node, admission"
    ctx = _context(shard_start=12, shard_end=23)
    report = capability_report_for(ctx, shard_start=0, shard_end=11)

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_SHARD_MISMATCH


def test_a_passing_report_for_another_recipe_version_is_refused():
    "A recipe's execution params changed, so its old proof no longer applies.\n\nTags: node, admission"
    ctx = _context()
    report = capability_report_for(ctx, recipe_version="0")

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_RECIPE_MISMATCH


def test_a_passing_report_from_another_backend_or_device_is_refused():
    "A CUDA proof does not admit a node that would serve the shard on CPU.\n\nTags: node, admission"
    ctx = _context(device="cpu")
    report = capability_report_for(ctx, device="cuda")

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_BACKEND_MISMATCH

    other_backend = capability_report_for(ctx, backend_id="some-other-runtime")
    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), other_backend)
    assert exc.value.reason == REASON_BACKEND_MISMATCH


def test_a_report_older_than_the_freshness_window_is_refused():
    "Hardware, drivers and weights move; an old proof is not a current one.\n\nTags: node, admission"
    ctx = _context()
    requirement = AdmissionRequirement.for_context(ctx, max_age_seconds=900)
    report = capability_report_for(ctx, age_seconds=901)

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(requirement, report)

    assert exc.value.reason == REASON_STALE

    still_fresh = capability_report_for(ctx, age_seconds=899)
    assert admit(requirement, still_fresh) is still_fresh


def test_a_future_dated_report_is_refused():
    "A report from the future is a broken clock, not a fresh proof.\n\nTags: node, admission"
    ctx = _context()
    report = capability_report_for(ctx, validated_at=time.time() + 3600)

    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)

    assert exc.value.reason == REASON_STALE


# ---------------------------------------------------------------------------
# The production validator: a real forward through the loaded backend
# ---------------------------------------------------------------------------


def test_the_production_validator_proves_a_working_backend_with_a_real_forward():
    "probe_capability runs the doctor's bounded forward on the backend that would serve.\n\nTags: node, admission"
    ctx = _context(_FakeBackend())

    report = probe_capability(ctx)

    assert report.passed
    assert admit(AdmissionRequirement.for_context(ctx), report) is report


def test_the_production_validator_fails_a_backend_that_cannot_execute():
    "A shard that loads but cannot run a forward yields a failed report, not a pass.\n\nTags: node, admission"
    ctx = _context(_FakeBackend(forward_error=RuntimeError("CUDA out of memory")))

    report = probe_capability(ctx)

    assert not report.passed
    with pytest.raises(CapabilityAdmissionError) as exc:
        admit(AdmissionRequirement.for_context(ctx), report)
    assert exc.value.reason == REASON_NOT_PASSED


def test_the_production_validator_fails_a_node_with_no_backend_at_all():
    "A server with no model backend cannot prove anything, so it is not admitted.\n\nTags: node, admission"
    ctx = _context(None)

    assert not probe_capability(ctx).passed


# ---------------------------------------------------------------------------
# run_startup: a refused report means the tracker is never called
# ---------------------------------------------------------------------------


class _FakeTorchNodeServer:
    started = False

    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.backend = kwargs.pop("_backend", None) or _FakeBackend()

    def start(self):
        type(self).started = True
        return 7099

    def stop(self):
        pass


@pytest.fixture
def startup_env(monkeypatch):
    """Fake hardware, wallet and tracker; records registrations *for this model*.

    Heartbeat threads that other tests leave running are daemon threads that
    outlive their test and re-register through this same patched `_post_json`, so
    a plain call log would be polluted by whatever ran before. Only posts naming
    this test's model can have come from this test's node.
    """
    import meshnet_node.startup as startup_mod

    posted: list[tuple[str, dict]] = []
    _FakeTorchNodeServer.started = False

    def _record(url, payload, timeout=10.0):
        if url.endswith("/v1/nodes/register") and payload.get("hf_repo") == MODEL:
            posted.append((url, payload))
        return {"node_id": "node-nca"}

    monkeypatch.setattr(
        startup_mod,
        "detect_hardware",
        lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
    )
    monkeypatch.setattr(
        startup_mod, "benchmark_throughput_checked", lambda _device: (10.0, True, None)
    )
    monkeypatch.setattr(startup_mod, "TorchNodeServer", _FakeTorchNodeServer)
    monkeypatch.setattr(
        startup_mod, "load_or_create_wallet", lambda **_kw: (b"", b"", "wallet-nca")
    )
    monkeypatch.setattr(
        startup_mod, "_get_json", lambda _url, timeout=10.0: {"relay_url": None, "nodes": []}
    )
    monkeypatch.setattr(startup_mod, "_post_json", _record)
    monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)
    return posted


def _start(**kwargs):
    return run_startup(
        tracker_url="http://127.0.0.1:8080",
        model_id=MODEL,
        shard_start=0,
        shard_end=23,
        **kwargs,
    )


def test_backend_validation_failure_registers_nothing(startup_env, monkeypatch):
    "A shard that cannot run a forward must never reach /v1/nodes/register.\n\nTags: node, admission, startup"
    import meshnet_node.startup as startup_mod

    broken = _FakeBackend(forward_error=RuntimeError("CUDA out of memory"))
    monkeypatch.setattr(
        startup_mod,
        "TorchNodeServer",
        lambda **kw: _FakeTorchNodeServer(_backend=broken, **kw),
    )

    with pytest.raises(CapabilityAdmissionError) as exc:
        _start()  # no validator: the production real-forward path

    assert exc.value.reason == REASON_NOT_PASSED
    assert startup_env == [], "the tracker was called despite a failed validation"
    assert not _FakeTorchNodeServer.started, "a failed node still opened an endpoint"


def test_a_report_for_a_different_model_cannot_be_reused_to_register(startup_env):
    "A success for one model must not admit another — the shape of a replayed proof.\n\nTags: node, admission, startup"
    with pytest.raises(CapabilityAdmissionError) as exc:
        _start(capability_validator=capability_stub(model_id="other/model-1b"))

    assert exc.value.reason == REASON_MODEL_MISMATCH
    assert startup_env == []


def test_a_stale_report_cannot_be_reused_to_register(startup_env):
    "An aged-out proof is refused before the node advertises itself.\n\nTags: node, admission, startup"
    with pytest.raises(CapabilityAdmissionError) as exc:
        _start(capability_validator=capability_stub(age_seconds=86_400))

    assert exc.value.reason == REASON_STALE
    assert startup_env == []


# ---------------------------------------------------------------------------
# Re-registration: the proof presented is fresh, never the one captured at boot
# ---------------------------------------------------------------------------


def test_run_startup_hands_the_heartbeat_a_refresher_for_the_current_shard(startup_env, monkeypatch):
    "The tracker refuses aged proofs, so the heartbeat must be able to re-prove what the node serves now.\n\nTags: node, admission, startup"
    import meshnet_node.startup as startup_mod

    captured: dict = {}
    monkeypatch.setattr(
        startup_mod, "_start_heartbeat", lambda *a, **kw: captured.update(kw)
    )

    _start(capability_validator=capability_stub())

    refresh = captured.get("refresh_capability")
    assert callable(refresh), "run_startup no longer wires a capability refresher"
    fresh = refresh({"hf_repo": MODEL, "model": MODEL.split("/")[-1]})
    assert fresh is not None
    assert fresh["model"]["model_id"] == MODEL
    assert (fresh["shard"]["start"], fresh["shard"]["end"]) == (0, 23)
    assert fresh["validated_at"] > time.time() - 60


def test_a_reregistration_presents_a_refreshed_proof(monkeypatch):
    "Replaying the boot-time report after an outage re-registers the node unroutable; the re-register path must present a fresh proof.\n\nTags: node, admission, startup"
    import json

    import meshnet_node.startup as startup_mod

    model = "acme/refresh-model-7b"
    boot_report = {"validated_at": 1.0, "marker": "boot"}
    fresh_report = {"validated_at": 2.0, "marker": "fresh"}
    posted: list[dict] = []

    def _record(url, payload, timeout=10.0):
        if url.endswith("/v1/nodes/register") and payload.get("hf_repo") == model:
            posted.append(json.loads(json.dumps(payload)))
            return {"node_id": "node-refresh"}
        raise SystemExit  # first heartbeat POST ends the daemon loop

    monkeypatch.setattr(startup_mod, "_post_json", _record)

    payload = {
        "hf_repo": model,
        "model": model.split("/")[-1],
        "capability_report": dict(boot_report),
    }
    thread = startup_mod._start_heartbeat(
        "http://tracker.invalid",
        startup_mod._PENDING_NODE_ID,  # forces a re-registration on the first tick
        payload,
        interval=0.02,
        refresh_capability=lambda _payload: dict(fresh_report),
    )

    # The loop must be dead before this test returns: once monkeypatch restores
    # `_post_json`, a surviving thread would re-register through whatever the
    # *next* test patches in and corrupt its call counts.
    thread.join(timeout=5.0)

    assert not thread.is_alive(), "the heartbeat loop outlived the test"
    assert posted, "the heartbeat never re-registered"
    assert posted[0]["capability_report"] == fresh_report


def test_a_matching_passing_report_registers_and_travels_with_the_payload(startup_env):
    "Registration carries the proof for exactly the model/shard/recipe it advertises.\n\nTags: node, admission, startup"
    node = _start()  # production validator against a working fake backend
    node.stop()

    assert len(startup_env) == 1
    url, payload = startup_env[0]
    assert url.endswith("/v1/nodes/register")

    report = payload["capability_report"]
    assert report["status"] == "passed"
    assert report["model"]["model_id"] == MODEL
    assert (report["shard"]["start"], report["shard"]["end"]) == (0, 23)
    assert report["recipe"]["recipe_id"] == DEFAULT_RECIPE_ID
    assert report["backend"]["device"] == "cpu"


def test_the_served_backend_is_loaded_with_the_recipe_that_was_validated(startup_env):
    "The recipe named in the report is the one the serving backend actually ran.\n\nTags: node, admission, startup"
    node = _start(recipe_id="eager-attention")
    node.stop()

    assert node.kwargs["recipe_params"] == {"attn_implementation": "eager"}
    report = startup_env[0][1]["capability_report"]
    assert report["recipe"]["recipe_id"] == "eager-attention"


def test_an_unknown_recipe_fails_before_any_weights_are_loaded(startup_env):
    "A typo'd recipe id is caught at resolution, not after a multi-minute load.\n\nTags: node, admission, startup"
    from meshnet_node.recipe_manifest import RecipeManifestError

    with pytest.raises(RecipeManifestError):
        _start(recipe_id="does-not-exist")

    assert startup_env == []
    assert not _FakeTorchNodeServer.started
