"""Lab 01 — The Agent Kernel.

Build a runtime that hosts agents the way an OS hosts processes: an explicit lifecycle,
enforced budgets, externalized state, tiered memory, and a recorded execution chain.

Everything ambient is injected — the model (``Policy``), the clock (``now``), the tools,
and the summarizer. No wall clock, no randomness, no ``uuid4``.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

from __future__ import annotations

import hashlib
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

# ======================================================================================
# 1. Lifecycle
# ======================================================================================


class RunState(str, Enum):
    CREATED = "created"
    PLANNING = "planning"
    ACTING = "acting"
    WAITING_INPUT = "waiting_input"
    SUSPENDED = "suspended"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class Event(str, Enum):
    START = "start"
    PROPOSE = "propose"
    OBSERVE = "observe"
    NEED_INPUT = "need_input"
    RESUME_INPUT = "resume_input"
    SUSPEND = "suspend"
    RESUME = "resume"
    FINISH = "finish"
    FAIL = "fail"
    CANCEL = "cancel"


TERMINAL_STATES: frozenset = frozenset(
    {RunState.COMPLETED, RunState.FAILED, RunState.CANCELLED}
)

# TODO: complete the transition table.
#
# Required edges (anything not listed must be ILLEGAL):
#   CREATED       --start--------> PLANNING        CREATED       --cancel--> CANCELLED
#   PLANNING      --propose------> ACTING          PLANNING      --need_input--> WAITING_INPUT
#   PLANNING      --suspend------> SUSPENDED       PLANNING      --finish--> COMPLETED
#   PLANNING      --fail---------> FAILED          PLANNING      --cancel--> CANCELLED
#   ACTING        --observe------> PLANNING        ACTING        --need_input--> WAITING_INPUT
#   ACTING        --suspend------> SUSPENDED       ACTING        --fail--> FAILED
#   ACTING        --cancel-------> CANCELLED
#   WAITING_INPUT --resume_input-> PLANNING        WAITING_INPUT --suspend--> SUSPENDED
#   WAITING_INPUT --cancel-------> CANCELLED       WAITING_INPUT --fail--> FAILED
#   SUSPENDED     --resume-------> PLANNING        SUSPENDED     --cancel--> CANCELLED
#   SUSPENDED     --fail---------> FAILED
TRANSITIONS: Mapping[Tuple[RunState, Event], RunState] = {}


class IllegalTransition(RuntimeError):
    def __init__(self, state: RunState, event: Event) -> None:
        super().__init__(f"illegal transition: {state.value} --{event.value}-->")
        self.state = state
        self.event = event


def transition(state: RunState, event: Event) -> RunState:
    """Apply one lifecycle event.

    A terminal state accepts NOTHING. An undeclared pair raises IllegalTransition.
    """
    # TODO
    raise NotImplementedError


# ======================================================================================
# 2. Memory
# ======================================================================================


def estimate_tokens(text: str) -> int:
    """Deterministic tokenizer stand-in: ceil(len / 4). Empty string -> 0."""
    # TODO
    raise NotImplementedError


@dataclass(frozen=True)
class Step:
    index: int
    thought: str
    tool: Optional[str] = None
    arguments: Tuple[Tuple[str, str], ...] = ()
    observation: Optional[str] = None
    error: Optional[str] = None
    tokens_in: int = 0
    tokens_out: int = 0
    started_at: float = 0.0
    ended_at: float = 0.0

    def render(self) -> str:
        parts = [f"[{self.index}] thought: {self.thought}"]
        if self.tool is not None:
            args = ", ".join(f"{k}={v}" for k, v in self.arguments)
            parts.append(f"    action: {self.tool}({args})")
        if self.observation is not None:
            parts.append(f"    observation: {self.observation}")
        if self.error is not None:
            parts.append(f"    error: {self.error}")
        return "\n".join(parts)


@dataclass
class Scratchpad:
    """Short-term working memory. Owns compaction, because it owns the quadratic term."""

    max_tokens: int
    summarize: Callable[[Sequence["Step"]], str]
    keep_recent: int = 2
    steps: List[Step] = field(default_factory=list)
    summary: str = ""
    compactions: int = 0

    def append(self, step: Step) -> None:
        self.steps.append(step)

    def token_count(self) -> int:
        """Tokens in the summary plus every rendered step."""
        # TODO
        raise NotImplementedError

    def render(self) -> str:
        """``[summary of earlier steps]\\n<summary>`` (if any) then each step, in order,
        newline-joined."""
        # TODO
        raise NotImplementedError

    def compact_if_needed(self) -> bool:
        """Fold all but the last ``keep_recent`` steps into the summary.

        Returns True only if a compaction happened. Do NOT compact when already under
        budget, and NEVER compact away the recent window even if still over budget —
        the model needs those steps to decide what to do next. Append the new summary to
        any existing one (space-separated) and increment ``compactions``.
        """
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class Fact:
    key: str
    value: str
    scope: str          # "user" | "tenant" | "app"
    owner: str
    tags: Tuple[str, ...] = ()


class SemanticMemory:
    """Durable facts partitioned by (scope, owner). The partition is not optional."""

    def __init__(self) -> None:
        self._facts: Dict[Tuple[str, str, str], Fact] = {}

    def put(self, fact: Fact) -> None:
        """Key on ``(scope, owner, key)``. Unknown scope -> ValueError."""
        # TODO
        raise NotImplementedError

    def get(self, scope: str, owner: str, key: str) -> Optional[Fact]:
        return self._facts.get((scope, owner, key))

    def search(self, *, scopes: Mapping[str, str], tags: Sequence[str], limit: int = 5) -> List[Fact]:
        """Facts VISIBLE to the caller, ranked by tag overlap desc then key asc.

        ``scopes`` maps a scope name to the caller's owner id. A fact is visible only if
        ``scopes.get(fact.scope) == fact.owner`` — filter BEFORE ranking, never after.
        With a non-empty ``tags``, drop facts with zero overlap.
        """
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class Episode:
    episode_id: str
    session_id: str
    goal: str
    outcome: str
    steps: int
    tags: Tuple[str, ...] = ()
    lesson: str = ""


class EpisodicMemory:
    """Append-only memory of what happened, recalled by tag overlap then recency."""

    def __init__(self) -> None:
        self._episodes: List[Episode] = []

    def record(self, episode: Episode) -> None:
        self._episodes.append(episode)

    def recall(self, tags: Sequence[str], limit: int = 3) -> List[Episode]:
        """Highest tag overlap first; ties broken by MOST RECENT first. Zero-overlap
        episodes are excluded entirely."""
        # TODO
        raise NotImplementedError

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


# ======================================================================================
# 3. Session state, checkpoints and optimistic concurrency
# ======================================================================================


class ConcurrentModification(RuntimeError):
    """Two workers tried to advance the same session. Exactly one may win."""


@dataclass(frozen=True)
class SessionSnapshot:
    session_id: str
    tenant: str
    user_id: str
    goal: str
    state: RunState
    version: int
    steps: Tuple[Step, ...]
    summary: str
    tokens_used: int
    cost_micros: int
    pending_question: Optional[str] = None
    created_at: float = 0.0
    updated_at: float = 0.0


class SessionStore:
    """Externalized state with optimistic concurrency control."""

    def __init__(self) -> None:
        self._sessions: Dict[str, SessionSnapshot] = {}

    def create(self, snapshot: SessionSnapshot) -> SessionSnapshot:
        """Store at version 1. Duplicate session_id -> KeyError."""
        # TODO (use dataclasses.replace to set the version)
        raise NotImplementedError

    def load(self, session_id: str) -> SessionSnapshot:
        """Missing -> KeyError."""
        # TODO
        raise NotImplementedError

    def save(self, snapshot: SessionSnapshot, *, expected_version: int) -> SessionSnapshot:
        """Compare-and-swap on ``version``.

        If the stored version != expected_version -> ConcurrentModification.
        On success, store at ``expected_version + 1`` and return the stored snapshot.
        """
        # TODO
        raise NotImplementedError

    def exists(self, session_id: str) -> bool:
        return session_id in self._sessions


# ======================================================================================
# 4. Session affinity
# ======================================================================================


def _hash_to_int(value: str) -> int:
    """Deterministic 64-bit hash from blake2b.

    NOT Python's ``hash()``: that is salted per process, so a ring built with it
    reshuffles every session on every pod restart.
    """
    digest = hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest()
    return int.from_bytes(digest, "big")


class AffinityRouter:
    """Consistent-hash routing of sessions to replicas, with draining."""

    def __init__(self, replicas: Iterable[str], *, virtual_nodes: int = 64) -> None:
        if virtual_nodes <= 0:
            raise ValueError("virtual_nodes must be > 0")
        self.virtual_nodes = virtual_nodes
        self._replicas: List[str] = []
        self._draining: set = set()
        self._ring: List[Tuple[int, str]] = []
        for replica in replicas:
            self.add(replica)

    def _rebuild(self) -> None:
        """Build a SORTED list of ``(hash, replica)`` with ``virtual_nodes`` points per
        replica, hashing ``f"{replica}#{i}"``. Sorting is what makes the walk work."""
        # TODO
        raise NotImplementedError

    def add(self, replica: str) -> None:
        """Idempotent. Keep ``_replicas`` sorted so construction order cannot change
        routing, then rebuild the ring."""
        # TODO
        raise NotImplementedError

    def remove(self, replica: str) -> None:
        """Unknown replica -> KeyError."""
        # TODO
        raise NotImplementedError

    def drain(self, replica: str) -> None:
        """Stop routing NEW sessions here without removing it from the ring — this is
        what a rolling deploy needs. Unknown replica -> KeyError."""
        # TODO
        raise NotImplementedError

    def route(self, session_id: str) -> str:
        """Walk the ring clockwise from ``_hash_to_int(session_id)``, skipping draining
        replicas, and wrap around. Empty ring or all-draining -> RuntimeError."""
        # TODO
        raise NotImplementedError

    def distribution(self, session_ids: Sequence[str]) -> Dict[str, int]:
        """Count routed sessions per non-draining replica (zeros included)."""
        # TODO
        raise NotImplementedError


# ======================================================================================
# 5. Tools and budgets
# ======================================================================================


@dataclass(frozen=True)
class ToolResult:
    ok: bool
    output: str
    tokens: int = 0
    cost_micros: int = 0


Tool = Callable[[Mapping[str, str]], ToolResult]


@dataclass(frozen=True)
class Budgets:
    """Kernel-enforced quotas. Not the agent author's responsibility."""

    max_steps: int = 12
    max_tokens: int = 20_000
    max_cost_micros: int = 500_000
    deadline_seconds: float = 30.0

    def __post_init__(self) -> None:
        # TODO: every field must be > 0, else ValueError
        raise NotImplementedError


@dataclass(frozen=True)
class BudgetBreach:
    kind: str      # "steps" | "tokens" | "cost" | "deadline"
    detail: str


# ======================================================================================
# 6. The model policy
# ======================================================================================


@dataclass(frozen=True)
class Decision:
    kind: str                     # "act" | "finish" | "ask"
    thought: str
    tool: Optional[str] = None
    arguments: Tuple[Tuple[str, str], ...] = ()
    answer: Optional[str] = None
    question: Optional[str] = None
    tokens_in: int = 0
    tokens_out: int = 0

    def __post_init__(self) -> None:
        # TODO: unknown kind -> ValueError; "act" without a tool -> ValueError;
        #       "ask" without a question -> ValueError
        raise NotImplementedError


Policy = Callable[[str], Decision]


# ======================================================================================
# 7. The kernel
# ======================================================================================


@dataclass
class RunResult:
    session_id: str
    state: RunState
    answer: Optional[str]
    steps: Tuple[Step, ...]
    breach: Optional[BudgetBreach]
    tokens_used: int
    cost_micros: int
    compactions: int
    question: Optional[str] = None

    @property
    def succeeded(self) -> bool:
        return self.state is RunState.COMPLETED


class AgentKernel:
    def __init__(
        self,
        *,
        store: SessionStore,
        tools: Mapping[str, Tool],
        policy: Policy,
        now: Callable[[], float],
        budgets: Budgets = None,
        summarize: Optional[Callable[[Sequence[Step]], str]] = None,
        scratchpad_max_tokens: int = 400,
        episodic: Optional[EpisodicMemory] = None,
        semantic: Optional[SemanticMemory] = None,
    ) -> None:
        self.store = store
        self.tools = dict(tools)
        self.policy = policy
        self.now = now
        self.budgets = budgets if budgets is not None else Budgets()
        self.summarize = summarize or default_summarizer
        self.scratchpad_max_tokens = scratchpad_max_tokens
        self.episodic = episodic if episodic is not None else EpisodicMemory()
        self.semantic = semantic if semantic is not None else SemanticMemory()

    def create_session(self, *, session_id: str, tenant: str, user_id: str, goal: str) -> SessionSnapshot:
        """Create a CREATED snapshot with empty steps and zeroed counters."""
        # TODO
        raise NotImplementedError

    def run(self, session_id: str, *, resume_answer: Optional[str] = None) -> RunResult:
        """The loop. In order:

        1. Load the snapshot. Enter the loop by transitioning:
           CREATED --start-->, SUSPENDED --resume-->, WAITING_INPUT --resume_input-->
           (and require ``resume_answer`` for the last, else ValueError).
           A terminal session -> IllegalTransition.
        2. Rebuild the scratchpad from the snapshot's steps and summary.
        3. If resuming with an answer, append a step recording the human input.
        4. Loop:
           a. Check budgets BEFORE calling the policy — steps, tokens, cost, deadline.
              On breach: record a BudgetBreach, transition FAIL, stop.
           b. Call ``self.policy(pad.render())``; add its tokens to the total.
           c. "finish" -> append a step, transition FINISH, stop.
              "ask"    -> append a step, transition NEED_INPUT, stop, carry the question.
              "act"    -> transition PROPOSE, dispatch the tool, append a step,
                          transition OBSERVE.
              An UNKNOWN tool is a RECOVERABLE error: record ``error`` on the step and
              keep looping so the model can correct itself.
           d. ``pad.compact_if_needed()``, then checkpoint.
        5. Checkpoint the final state. If terminal, record an Episode tagged with the
           tools used, carrying the breach detail as its lesson.
        """
        # TODO
        raise NotImplementedError

    def _checkpoint(self, snapshot, state, pad, tokens_used, cost_micros, version, question):
        """Build the next SessionSnapshot and CAS it into the store.

        Return ``(stored_snapshot, stored_snapshot.version)``.
        """
        # TODO
        raise NotImplementedError

    def execution_chain(self, session_id: str) -> List[Mapping[str, object]]:
        """The audit artifact, read FROM THE STORE (compaction must not lose it).

        One dict per step with: session_id, tenant, user_id, step, tool, arguments
        (as a dict), ok (error is None), error, tokens (in + out), duration (ended -
        started, rounded to 6dp).
        """
        # TODO
        raise NotImplementedError


def _next_index(snapshot: SessionSnapshot) -> int:
    return (snapshot.steps[-1].index + 1) if snapshot.steps else 1


def default_summarizer(steps: Sequence[Step]) -> str:
    """Deterministic compaction. A real kernel calls a model here.

    Return: ``ran N step(s) using <sorted unique tools or 'no tools'>; last result: X``
    where X is the last non-None observation (empty string if none).
    """
    # TODO
    raise NotImplementedError


def main() -> None:
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()