from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Protocol, Sequence

import numpy as np


@dataclass(frozen=True)
class BrainSignals:
    """Signals emitted by one agent-context brain for one transcript update.

    `readiness` is the controller-facing score for a speculative short reply:

        readiness = 1 / (1 + mean_token_entropy)

    A low-entropy draft means this agent has a confident next move. `p_end` is a
    turn-completion heuristic: in the live loop it should combine trailing
    silence, sentence-final punctuation, and/or high EOS probability.
    """

    surprise: float
    hidden: np.ndarray
    readiness: float
    p_end: float

    def __post_init__(self) -> None:
        hidden = np.asarray(self.hidden, dtype=np.float32)
        if hidden.ndim != 1:
            raise ValueError("hidden must be a 1-D float32 vector")
        if not 0.0 <= float(self.readiness) <= 1.0:
            raise ValueError("readiness must be in [0, 1]")
        if not 0.0 <= float(self.p_end) <= 1.0:
            raise ValueError("p_end must be in [0, 1]")
        object.__setattr__(self, "hidden", hidden)
        object.__setattr__(self, "surprise", float(self.surprise))
        object.__setattr__(self, "readiness", float(self.readiness))
        object.__setattr__(self, "p_end", float(self.p_end))


class Brain(Protocol):
    """Interface the live instrumented LLM must satisfy for one agent context."""

    def next_signals(self) -> BrainSignals:
        """Return signals for the newest incremental transcript update."""


class ReplayBrain:
    """Deterministic `Brain` backed by precomputed signal samples."""

    def __init__(self, samples: Sequence[BrainSignals | dict[str, object]], *, name: str = "replay") -> None:
        self.name = name
        self._samples = [coerce_signals(sample) for sample in samples]
        self._index = 0

    def __len__(self) -> int:
        return len(self._samples)

    def __iter__(self) -> Iterable[BrainSignals]:
        return iter(self._samples)

    def reset(self) -> None:
        self._index = 0

    def next_signals(self) -> BrainSignals:
        if self._index >= len(self._samples):
            raise StopIteration(f"ReplayBrain {self.name!r} is exhausted")
        sample = self._samples[self._index]
        self._index += 1
        return sample

    @classmethod
    def from_npz(
        cls,
        path: str | Path,
        *,
        readiness: Sequence[float] | None = None,
        p_end: Sequence[float] | None = None,
        name: str | None = None,
    ) -> "ReplayBrain":
        dump = np.load(path, allow_pickle=True)
        surprises = np.asarray(dump["nll_series"], dtype=np.float32)
        hidden = np.asarray(dump["hidden_states"], dtype=np.float32)
        if hidden.ndim != 2:
            raise ValueError("hidden_states in npz must be a 2-D matrix")

        n_steps = int(surprises.shape[0])
        readiness_values = (
            np.asarray(readiness, dtype=np.float32)
            if readiness is not None
            else np.linspace(0.35, 0.75, n_steps, dtype=np.float32)
        )
        p_end_values = np.asarray(p_end, dtype=np.float32) if p_end is not None else np.zeros(n_steps, dtype=np.float32)
        if n_steps:
            p_end_values[-1] = max(float(p_end_values[-1]), 0.95)

        samples = [
            BrainSignals(
                surprise=float(surprises[index]),
                hidden=hidden[index],
                readiness=float(readiness_values[index]),
                p_end=float(p_end_values[index]),
            )
            for index in range(n_steps)
        ]
        return cls(samples, name=name or Path(path).stem)


def coerce_signals(sample: BrainSignals | dict[str, object]) -> BrainSignals:
    if isinstance(sample, BrainSignals):
        return sample
    return BrainSignals(
        surprise=float(sample["surprise"]),
        hidden=np.asarray(sample["hidden"], dtype=np.float32),
        readiness=float(sample["readiness"]),
        p_end=float(sample["p_end"]),
    )