"""Deterministic fake ``ShardEngine`` fixture (DGR-032).

``FakeShardEngine`` is a pure-Python, allocation-cheap subclass of
:class:`~meshnet_node.shard_engine.ShardEngine`: no llama.cpp, no native
buffers, no GPU, no filesystem or network I/O. Every prefill/decode output is
a deterministic pure function of ``(loaded range, request inputs,
idempotency_step)`` — hashed with SHA-256 — so replaying identical inputs on
a fresh session always yields byte-identical output. It exists so worker
wiring, gRPC harnesses (DGR-033), and lifecycle/session logic can be
exercised end-to-end before a real llama.cpp-backed engine (DGR-037) exists.

This is FIXTURE evidence only. ``EVIDENCE_CLASS`` is set to ``"fixture"`` (as
opposed to ``"real"``) precisely so a later story comparing engines
programmatically — DGR-036's fixture-vs-real-model parity check — can assert
it is actually comparing a fixture against a real engine rather than two
fixtures. This module proves lifecycle/session/epoch/fault-injection
semantics; it says nothing about numerical parity with a real model. Real-
model certification is DGR-036 onward (DGR-053/DGR-054 for V4 alpha).

Fault injection (delay, memory pressure, malformed output, crash) is
deterministic and opt-in via :class:`FakeShardEngineConfig`. Every knob
defaults to off, so a bare ``FakeShardEngine()`` reproduces plain
deterministic fixture behavior and passes
:func:`tests.shard_engine_contract.assert_shard_engine_contract` unmodified.
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass
from typing import Callable

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

__all__ = ["FakeShardEngineConfig", "FakeShardEngine", "TOKEN_ID_VOCAB_SIZE", "MALFORMED_TOKEN_ID_FLOOR"]

TOKEN_ID_VOCAB_SIZE = 50_000
# A malformed tail output is deterministically pushed past the fixture's own
# advertised vocabulary range, so a downstream consumer checking "is this
# token_id within the vocab this fixture promises" can detect it without any
# extra signalling from the engine.
MALFORMED_TOKEN_ID_FLOOR = 100_000_000


def _default_crash_exception() -> BaseException:
    return RuntimeError(
        "FakeShardEngine: injected crash (simulated process failure, not a StructuredStatus)"
    )


@dataclass(frozen=True)
class FakeShardEngineConfig:
    """Deterministic fault-injection knobs.

    Every knob is off (``0``/``None``/``False``) by default. ``sleep`` is
    injectable so tests can assert a delay was requested without an actual
    process sleep; ``crash_exception_factory`` is injectable so tests can
    assert on a specific exception type/instance.
    """

    step_delay_seconds: float = 0.0
    sleep: Callable[[float], None] = time.sleep
    memory_budget_bytes: int | None = None
    malformed_output: bool = False
    crash_after_calls: int | None = None
    crash_exception_factory: Callable[[], BaseException] = _default_crash_exception

    def __post_init__(self) -> None:
        if self.step_delay_seconds < 0:
            raise ValueError("step_delay_seconds must be non-negative")
        if self.memory_budget_bytes is not None and self.memory_budget_bytes < 0:
            raise ValueError("memory_budget_bytes must be non-negative")
        if self.crash_after_calls is not None and self.crash_after_calls <= 0:
            raise ValueError("crash_after_calls must be positive when set")


@dataclass
class _SessionState:
    epoch: int
    cancelled: bool = False


class FakeShardEngine(ShardEngine):
    """Deterministic fixture ``ShardEngine``. See module docstring."""

    EVIDENCE_CLASS = "fixture"

    def __init__(self, config: FakeShardEngineConfig | None = None) -> None:
        self._config = config or FakeShardEngineConfig()
        self._loaded: LoadRequest | None = None
        self._sessions: dict[str, _SessionState] = {}
        self._cancelled_total = 0
        self._generated_tokens = 0
        self._call_count = 0
        self._bytes_used = 0

    # -- lifecycle -----------------------------------------------------

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

    def capabilities(self) -> EngineCapabilities:
        if self._loaded is None:
            return EngineCapabilities(
                status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine 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=str(request.recipe.get("architecture", "fake")),
            max_concurrent_sessions=64,
            max_context_tokens=131072,
            supports_mtp=False,
        )

    def prefill(self, request: PrefillRequest) -> StepResult:
        return self._step(
            session_id=request.session_id,
            route_epoch=request.route_epoch,
            idempotency_step=request.idempotency_step,
            token_ids=request.token_ids,
            input_bundle=request.input,
            cache_result_on_success=CacheResult.STORED,
            opens_session=True,
        )

    def decode(self, request: DecodeRequest) -> StepResult:
        token_ids = (request.token_id,) if request.token_id is not None else None
        return self._step(
            session_id=request.session_id,
            route_epoch=request.route_epoch,
            idempotency_step=request.idempotency_step,
            token_ids=token_ids,
            input_bundle=request.input,
            cache_result_on_success=CacheResult.HIT,
            opens_session=False,
        )

    def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus:
        session = self._sessions.get(session_id)
        if session is None:
            session = _SessionState(epoch=0)
            self._sessions[session_id] = session
        if not session.cancelled:
            self._cancelled_total += 1
        session.cancelled = True
        return StructuredStatus(StatusCode.CANCELLED, reason or "fake engine: session cancelled")

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

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

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

    # -- shared step machinery ------------------------------------------

    def _step(
        self,
        *,
        session_id: str,
        route_epoch: int,
        idempotency_step: int,
        token_ids: tuple[int, ...] | None,
        input_bundle: BoundaryBundle | None,
        cache_result_on_success: CacheResult,
        opens_session: bool,
    ) -> StepResult:
        if self._loaded is None:
            return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine not loaded"))

        self._call_count += 1
        if self._config.crash_after_calls is not None and self._call_count == self._config.crash_after_calls:
            raise self._config.crash_exception_factory()

        session = self._sessions.get(session_id)
        if session is None:
            if not opens_session:
                return StepResult(
                    status=StructuredStatus(StatusCode.NOT_FOUND, "no cached session state for decode"),
                    cache_result=CacheResult.MISS,
                )
            session = _SessionState(epoch=route_epoch)
            self._sessions[session_id] = session

        if session.cancelled:
            return StepResult(status=StructuredStatus(StatusCode.CANCELLED, "session cancelled"))

        if route_epoch < session.epoch:
            return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "stale route epoch"))
        session.epoch = route_epoch

        if self._config.step_delay_seconds:
            self._config.sleep(self._config.step_delay_seconds)

        seed = self._seed_bytes(token_ids, input_bundle)
        self._bytes_used += len(seed)
        budget = self._config.memory_budget_bytes
        if budget is not None and self._bytes_used > budget:
            return StepResult(
                status=StructuredStatus(
                    StatusCode.RESOURCE_EXHAUSTED,
                    "fake engine memory pressure budget exceeded",
                    retryable=True,
                    details={"memory_budget_bytes": str(budget), "bytes_used": str(self._bytes_used)},
                )
            )

        output = self._transform(seed, idempotency_step, input_bundle)
        if isinstance(output, TokenOutput):
            self._generated_tokens += 1
        return StepResult(status=StructuredStatus(StatusCode.OK, "ok"), cache_result=cache_result_on_success, output=output)

    @staticmethod
    def _seed_bytes(token_ids: tuple[int, ...] | None, bundle: BoundaryBundle | None) -> bytes:
        if token_ids:
            seed = b"".join(int(t).to_bytes(8, "big") for t in token_ids)
        elif bundle is not None:
            seed = b"".join(tensor.data for tensor in bundle.tensors)
            if bundle.token_id_sideband:
                seed += b"".join(int(t).to_bytes(8, "big") for t in bundle.token_id_sideband)
        else:
            seed = b""
        return seed

    def _transform(
        self, seed: bytes, idempotency_step: int, input_bundle: BoundaryBundle | None
    ) -> BoundaryBundle | TokenOutput:
        assert self._loaded is not None
        digest = hashlib.sha256(seed + idempotency_step.to_bytes(8, "big")).digest()
        loaded = self._loaded
        is_tail = loaded.shard_end >= loaded.total_layers - 1
        is_head = loaded.shard_start == 0

        if is_tail:
            token_id = int.from_bytes(digest[:4], "big") % TOKEN_ID_VOCAB_SIZE
            if self._config.malformed_output:
                token_id = MALFORMED_TOKEN_ID_FLOOR + token_id
            return TokenOutput(token_id=token_id)

        boundary_point = "post_head_residual" if is_head else "post_middle_residual"
        architecture = (
            input_bundle.architecture if input_bundle is not None else str(loaded.recipe.get("architecture", "fake"))
        )
        data = digest
        if self._config.malformed_output:
            architecture = f"malformed:{architecture}"
            data = digest[:1]
        tensor = EngineTensor(
            name="hidden_states",
            shape=(1, max(len(seed) // 8, 1)),
            dtype="bfloat16",
            data=data,
        )
        token_id_sideband = input_bundle.token_id_sideband if input_bundle is not None else None
        return BoundaryBundle(
            tensors=(tensor,),
            architecture=architecture,
            boundary_point=boundary_point,
            token_id_sideband=token_id_sideband,
        )
