"""
AXON Runtime — Context Manager
================================
Maintains mutable execution state that flows between steps
within a single execution unit.

The ContextManager is the working memory of a running AXON program.
It tracks:
    - Step results: named outputs from completed steps
    - Message history: the multi-turn conversation with the model
    - Flow parameters: input arguments passed to the flow
    - System prompt: the compiled persona + anchor instructions

Every mutation is traced through the Tracer for full observability.
"""

from __future__ import annotations

import copy
from dataclasses import dataclass, field
from typing import Any

from axon.runtime.tracer import Tracer


# ═══════════════════════════════════════════════════════════════════
#  MESSAGE ROLE — typed conversation roles
# ═══════════════════════════════════════════════════════════════════

ROLE_SYSTEM = "system"
ROLE_USER = "user"
ROLE_ASSISTANT = "assistant"
VALID_ROLES = frozenset({ROLE_SYSTEM, ROLE_USER, ROLE_ASSISTANT})


# ═══════════════════════════════════════════════════════════════════
#  CONTEXT SNAPSHOT — immutable state capture
# ═══════════════════════════════════════════════════════════════════


@dataclass(frozen=True)
class ContextSnapshot:
    """An immutable point-in-time capture of execution state.

    Used by the Tracer and debugging tools to record the exact
    state of the context at any moment during execution.

    Attributes:
        step_results:     Copy of all step name → result mappings.
        message_count:    Number of messages in the conversation.
        variables:        Copy of all flow parameter bindings.
        current_step:     Name of the step being executed (if any).
    """

    step_results: dict[str, Any] = field(default_factory=dict)
    message_count: int = 0
    variables: dict[str, Any] = field(default_factory=dict)
    current_step: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-compatible dictionary."""
        result: dict[str, Any] = {
            "step_results": {
                k: repr(v) for k, v in self.step_results.items()
            },
            "message_count": self.message_count,
        }
        if self.variables:
            result["variables"] = {
                k: repr(v) for k, v in self.variables.items()
            }
        if self.current_step:
            result["current_step"] = self.current_step
        return result


# ═══════════════════════════════════════════════════════════════════
#  CONTEXT MANAGER — mutable execution state
# ═══════════════════════════════════════════════════════════════════


class ContextManager:
    """Maintains execution state between steps in a flow.

    The ContextManager is scoped to a single ``CompiledExecutionUnit``
    (one ``run`` statement). Each run gets its own context.

    Usage::

        ctx = ContextManager(system_prompt="You are LegalExpert.", tracer=tracer)
        ctx.set_variable("document", contract_text)
        ctx.append_message("user", "Analyze this contract.")
        ctx.append_message("assistant", "The contract contains...")
        ctx.set_step_result("analyze", {"clauses": [...]})
        result = ctx.get_step_result("analyze")
        snapshot = ctx.snapshot()
    """

    def __init__(
        self,
        system_prompt: str = "",
        tracer: Tracer | None = None,
    ) -> None:
        self._system_prompt = system_prompt
        self._tracer = tracer
        self._step_results: dict[str, Any] = {}
        self._variables: dict[str, Any] = {}
        self._messages: list[dict[str, str]] = []
        self._current_step: str = ""

    # — System prompt —

    @property
    def system_prompt(self) -> str:
        """The compiled system prompt for this execution unit."""
        return self._system_prompt

    # — Step state tracking —

    @property
    def current_step(self) -> str:
        """The name of the step currently being executed."""
        return self._current_step

    @current_step.setter
    def current_step(self, name: str) -> None:
        """Set the current step being executed."""
        self._current_step = name

    def set_step_result(self, step_name: str, result: Any) -> None:
        """Record the output of a completed step.

        Args:
            step_name:  The name of the step that produced the result.
            result:     The step's output value.

        Raises:
            ValueError: If ``step_name`` is empty.
        """
        if not step_name:
            raise ValueError("step_name must not be empty")

        self._step_results[step_name] = result

    def get_step_result(self, step_name: str) -> Any:
        """Retrieve the output of a previously completed step.

        Args:
            step_name:  The name of the step whose result is needed.

        Returns:
            The step's output value.

        Raises:
            KeyError: If the step has not completed yet.
        """
        if step_name not in self._step_results:
            raise KeyError(
                f"Step '{step_name}' has no result. "
                f"Available: {list(self._step_results.keys())}"
            )
        return self._step_results[step_name]

    def has_step_result(self, step_name: str) -> bool:
        """Check whether a step has a recorded result."""
        return step_name in self._step_results

    @property
    def completed_steps(self) -> list[str]:
        """Names of all steps that have recorded results, in insertion order."""
        return list(self._step_results.keys())

    # — Variable bindings (flow parameters & intermediate values) —

    def set_variable(self, name: str, value: Any) -> None:
        """Bind a named variable in the execution context.

        Args:
            name:   The variable name.
            value:  The variable's value.

        Raises:
            ValueError: If ``name`` is empty.
        """
        if not name:
            raise ValueError("Variable name must not be empty")

        self._variables[name] = value

    def get_variable(self, name: str) -> Any:
        """Retrieve a named variable from the execution context.

        Args:
            name:  The variable name.

        Returns:
            The variable's value.

        Raises:
            KeyError: If the variable has not been set.
        """
        if name not in self._variables:
            raise KeyError(
                f"Variable '{name}' is not defined. "
                f"Available: {list(self._variables.keys())}"
            )
        return self._variables[name]

    def has_variable(self, name: str) -> bool:
        """Check whether a named variable exists."""
        return name in self._variables

    def get_variables(self) -> dict[str, Any]:
        """Return a shallow copy of all variable bindings."""
        return dict(self._variables)

    # — Message history (multi-turn conversation) —

    def append_message(self, role: str, content: str) -> None:
        """Add a message to the conversation history.

        Args:
            role:     One of ``"system"``, ``"user"``, ``"assistant"``.
            content:  The message content.

        Raises:
            ValueError: If ``role`` is not a valid conversation role.
            ValueError: If ``content`` is empty.
        """
        if role not in VALID_ROLES:
            raise ValueError(
                f"Invalid role '{role}'. Must be one of: {sorted(VALID_ROLES)}"
            )
        if not content:
            raise ValueError("Message content must not be empty")

        self._messages.append({"role": role, "content": content})

    def get_message_history(self) -> list[dict[str, str]]:
        """Return a copy of the full message history."""
        return list(self._messages)

    @property
    def message_count(self) -> int:
        """The number of messages in the conversation history."""
        return len(self._messages)

    def clear_messages(self) -> None:
        """Clear the entire message history."""
        self._messages.clear()

    # — Snapshot (immutable state capture) —

    def snapshot(self) -> ContextSnapshot:
        """Capture an immutable snapshot of the current execution state.

        The snapshot deep-copies step results and variables to prevent
        mutations from affecting the captured state.

        Returns:
            An immutable ``ContextSnapshot`` representing the current state.
        """
        return ContextSnapshot(
            step_results=copy.deepcopy(self._step_results),
            message_count=self.message_count,
            variables=copy.deepcopy(self._variables),
            current_step=self._current_step,
        )

    # — Reset —

    def reset(self) -> None:
        """Clear all state, returning the context to its initial condition.

        The system prompt is preserved; everything else is cleared.
        """
        self._step_results.clear()
        self._variables.clear()
        self._messages.clear()
        self._current_step = ""
