"""DGR-031 ``ShardEngine`` contract tests.

``_ReferenceEngine`` below is a minimal, in-memory ``ShardEngine`` that exists
only to prove :func:`assert_shard_engine_contract` is non-vacuous and to pin
the abstract contract's own validation rules. It is deliberately not the
DGR-032 deterministic fixture (delay/memory-pressure/malformed/crash
injection, full session/epoch modeling for the fake worker) — that is a
separate, larger story. DGR-032 and DGR-037 are expected to import
``assert_shard_engine_contract`` from ``tests/shard_engine_contract.py``
against their own engines.
"""

from __future__ import annotations

import hashlib

import pytest

from meshnet_node.shard_engine import (
    ArchitectureAuxStateHook,
    BoundaryBundle,
    DecodeRequest,
    EngineCapabilities,
    EngineTensor,
    HealthResult,
    LoadRequest,
    LoadResult,
    MetricsResult,
    MtpHook,
    PrefillRequest,
    ShardEngine,
    StepResult,
    TokenOutput,
)
from meshnet_node.shard_lifecycle import CacheResult, StatusCode, StructuredStatus

from shard_engine_contract import assert_shard_engine_contract


class _ReferenceEngine(ShardEngine):
    """Minimal in-memory engine used only to exercise the shared contract."""

    def __init__(self) -> None:
        self._loaded: LoadRequest | None = None
        self._sessions: dict[str, dict] = {}
        self._cancelled_total = 0

    def load(self, request: LoadRequest) -> LoadResult:
        self._loaded = request
        return LoadResult(
            status=StructuredStatus(StatusCode.OK, "loaded"),
            effective_start=request.shard_start,
            architecture="dense",
        )

    def capabilities(self) -> EngineCapabilities:
        if self._loaded is None:
            return EngineCapabilities(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "not loaded"))
        request = self._loaded
        return EngineCapabilities(
            status=StructuredStatus(StatusCode.OK, "ready"),
            shard_start=request.shard_start,
            shard_end=request.shard_end,
            effective_start=request.shard_start,
            total_layers=request.total_layers,
            architecture="dense",
            max_concurrent_sessions=8,
            max_context_tokens=131072,
            supports_mtp=False,
        )

    def prefill(self, request: PrefillRequest) -> StepResult:
        if self._loaded is None:
            return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine not loaded"))
        self._sessions[request.session_id] = {"epoch": request.route_epoch, "cancelled": False}
        output = self._transform(self._seed_bytes(request.token_ids, request.input), request.idempotency_step)
        return StepResult(status=StructuredStatus(StatusCode.OK, "prefilled"), cache_result=CacheResult.STORED, output=output)

    def decode(self, request: DecodeRequest) -> StepResult:
        session = self._sessions.get(request.session_id)
        if session is None:
            return StepResult(
                status=StructuredStatus(StatusCode.NOT_FOUND, "no cached session state"),
                cache_result=CacheResult.MISS,
            )
        if session["cancelled"]:
            return StepResult(status=StructuredStatus(StatusCode.CANCELLED, "session cancelled"))
        if request.route_epoch < session["epoch"]:
            return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "stale route epoch"))
        session["epoch"] = request.route_epoch
        token_ids = (request.token_id,) if request.token_id is not None else None
        output = self._transform(self._seed_bytes(token_ids, request.input), request.idempotency_step)
        return StepResult(status=StructuredStatus(StatusCode.OK, "decoded"), cache_result=CacheResult.HIT, output=output)

    def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus:
        session = self._sessions.setdefault(session_id, {"epoch": 0, "cancelled": False})
        if not session["cancelled"]:
            self._cancelled_total += 1
        session["cancelled"] = True
        return StructuredStatus(StatusCode.CANCELLED, reason or "cancelled")

    def release(self, session_id: str) -> StructuredStatus:
        self._sessions.pop(session_id, None)
        return StructuredStatus(StatusCode.OK, "released")

    def health(self) -> HealthResult:
        return HealthResult(
            status=StructuredStatus(StatusCode.OK, "ok"),
            serving=self._loaded is not None,
            state="SERVING" if self._loaded is not None else "NOT_LOADED",
            active_sessions=len(self._sessions),
        )

    def metrics(self) -> MetricsResult:
        return MetricsResult(
            status=StructuredStatus(StatusCode.OK, "ok"),
            active_sessions=len(self._sessions),
            cancelled_sessions=self._cancelled_total,
        )

    @staticmethod
    def _seed_bytes(token_ids, bundle: BoundaryBundle | None) -> bytes:
        if token_ids:
            return b"".join(int(t).to_bytes(4, "big") for t in token_ids)
        if bundle is not None:
            return b"".join(tensor.data for tensor in bundle.tensors)
        return b""

    def _transform(self, seed: bytes, idempotency_step: int) -> BoundaryBundle | TokenOutput:
        digest = hashlib.sha256(seed + idempotency_step.to_bytes(4, "big")).digest()
        assert self._loaded is not None
        if self._loaded.shard_end >= self._loaded.total_layers - 1:
            token_id = int.from_bytes(digest[:4], "big") % 50_000
            return TokenOutput(token_id=token_id)
        tensor = EngineTensor(name="hidden_states", shape=(1, max(len(seed) // 4, 1)), dtype="bfloat16", data=digest)
        return BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")


def test_reference_engine_obeys_the_shared_shard_engine_contract():
    assert_shard_engine_contract(_ReferenceEngine)


def test_shard_engine_is_abstract_and_cannot_be_instantiated_directly():
    with pytest.raises(TypeError):
        ShardEngine()  # type: ignore[abstract]


def test_engine_tensor_rejects_empty_name_shape_or_dtype():
    with pytest.raises(ValueError):
        EngineTensor(name="", shape=(1,), dtype="bfloat16", data=b"x")
    with pytest.raises(ValueError):
        EngineTensor(name="t", shape=(), dtype="bfloat16", data=b"x")
    with pytest.raises(ValueError):
        EngineTensor(name="t", shape=(0,), dtype="bfloat16", data=b"x")
    with pytest.raises(ValueError):
        EngineTensor(name="t", shape=(1,), dtype="", data=b"x")


def test_boundary_bundle_requires_at_least_one_tensor():
    with pytest.raises(ValueError):
        BoundaryBundle(tensors=(), architecture="dense", boundary_point="pre_tail_residual")


def test_boundary_bundle_tensor_lookup_by_name():
    tensor = EngineTensor(name="hidden_states", shape=(1, 1), dtype="bfloat16", data=b"\x00\x00")
    bundle = BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")
    assert bundle.tensor("hidden_states") is tensor
    with pytest.raises(KeyError):
        bundle.tensor("router_logits")


def test_token_output_rejects_negative_token_id():
    with pytest.raises(ValueError):
        TokenOutput(token_id=-1)


def test_mtp_hook_is_reserved_and_refuses_to_enable():
    MtpHook()  # disabled is fine
    with pytest.raises(ValueError):
        MtpHook(enabled=True)
    with pytest.raises(ValueError):
        MtpHook(draft_token_count=-1)


def test_architecture_aux_state_hook_carries_opaque_shard_local_state():
    hook = ArchitectureAuxStateHook(kind="csa", state={"window": 128})
    assert hook.kind == "csa"
    assert hook.state == {"window": 128}


def test_prefill_and_decode_requests_require_exactly_one_input_kind():
    with pytest.raises(ValueError):
        PrefillRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)
    with pytest.raises(ValueError):
        PrefillRequest(
            session_id="s",
            route_epoch=0,
            position=0,
            idempotency_step=0,
            token_ids=(1,),
            input=BoundaryBundle(
                tensors=(EngineTensor(name="hidden_states", shape=(1,), dtype="bfloat16", data=b"x"),),
                architecture="dense",
                boundary_point="pre_tail_residual",
            ),
        )
    with pytest.raises(ValueError):
        DecodeRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)


def test_load_request_validates_shard_range_against_total_layers():
    LoadRequest(artifact_path="a", shard_start=0, shard_end=3, total_layers=4)
    with pytest.raises(ValueError):
        LoadRequest(artifact_path="a", shard_start=0, shard_end=4, total_layers=4)
    with pytest.raises(ValueError):
        LoadRequest(artifact_path="", shard_start=0, shard_end=0, total_layers=1)
    with pytest.raises(ValueError):
        LoadRequest(artifact_path="a", shard_start=3, shard_end=1, total_layers=4)


def test_step_result_requires_an_output_when_status_is_ok():
    with pytest.raises(ValueError):
        StepResult(status=StructuredStatus(StatusCode.OK, "ok"), output=None)
    # A non-OK status is allowed to carry no output.
    StepResult(status=StructuredStatus(StatusCode.NOT_FOUND, "missing"), output=None)


def test_shard_engine_module_imports_no_native_or_grpc_or_wire_abi_types():
    import meshnet_node.shard_engine as shard_engine_module

    # The boundary module must not *import* anything that would let a
    # ggml_tensor, llama context/scheduler handle, ctypes native handle, or a
    # generated-protobuf (ABI) message leak into a project-owned dataclass
    # field. Checking bound globals (not docstring prose) proves this
    # structurally rather than by convention.
    forbidden_modules = {"ctypes", "grpc", "meshnet_node.native_protocol"}
    for name, value in vars(shard_engine_module).items():
        module_name = getattr(value, "__name__", None)
        assert module_name not in forbidden_modules, (
            f"shard_engine.{name} binds forbidden module {module_name!r}"
        )
