"""
SharedState — typed, versioned, observable shared state for multi-agent workflows.

Design notes:
  - Pydantic models for type safety
  - Every mutation is recorded as a diff (for checkpointing & time-travel)
  - Optimistic locking via version numbers (for concurrent agent access)
  - Subscribable: agents can watch for state changes
"""

from __future__ import annotations

import copy
from collections.abc import Callable
from datetime import datetime
from typing import Any, Generic, TypeVar

from pydantic import BaseModel, Field

from nexus.types import new_id

T = TypeVar("T", bound=BaseModel)


class StateDiff(BaseModel):
    """Records a single state mutation."""

    version: int
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    agent_id: str | None = None
    path: str = ""  # Dot-separated path to the changed field
    old_value: Any = None
    new_value: Any = None
    description: str = ""


class SharedState(Generic[T]):
    """
    A typed, versioned, observable shared state container.

    Usage:
        class MyState(BaseModel):
            plan: str = ""
            code: str = ""
            review_score: float = 0.0

        state = SharedState(MyState())
        state.update(plan="Write a web scraper", agent_id="planner-001")
        print(state.data.plan)
        print(state.version)  # 1
        print(state.history)  # [StateDiff(...)]
    """

    def __init__(self, initial: T) -> None:
        self._id = new_id()
        self._data: T = initial
        self._version: int = 0
        self._history: list[StateDiff] = []
        self._subscribers: list[Callable[[StateDiff], Any]] = []

    @property
    def id(self) -> str:
        return self._id

    @property
    def data(self) -> T:
        """Read-only access to the current state."""
        return self._data

    @property
    def version(self) -> int:
        return self._version

    @property
    def history(self) -> list[StateDiff]:
        return list(self._history)

    def update(
        self,
        agent_id: str | None = None,
        expected_version: int | None = None,
        description: str = "",
        **changes: Any,
    ) -> T:
        """
        Apply updates to the state.

        Args:
            agent_id: Who is making this change (for audit trail).
            expected_version: If provided, enables optimistic locking.
                              Raises ValueError if current version != expected.
            description: Human-readable description of the change.
            **changes: Key-value pairs to update on the state model.

        Returns:
            The updated state.

        Raises:
            ValueError: If expected_version doesn't match (optimistic lock failure).
            AttributeError: If a key doesn't exist on the state model.
        """
        # Optimistic locking
        if expected_version is not None and expected_version != self._version:
            raise ValueError(
                f"Optimistic lock failure: expected version {expected_version}, "
                f"but current version is {self._version}"
            )

        old_data = self._data.model_dump()

        for key, value in changes.items():
            if not hasattr(self._data, key):
                raise AttributeError(
                    f"State model {type(self._data).__name__} has no field '{key}'"
                )
            setattr(self._data, key, value)

        # Record diffs
        new_data = self._data.model_dump()
        self._version += 1

        for key, new_val in new_data.items():
            old_val = old_data.get(key)
            if old_val != new_val:
                diff = StateDiff(
                    version=self._version,
                    agent_id=agent_id,
                    path=key,
                    old_value=old_val,
                    new_value=new_val,
                    description=description,
                )
                self._history.append(diff)
                # Notify subscribers
                for sub in self._subscribers:
                    sub(diff)

        return self._data

    def subscribe(self, callback: Callable[[StateDiff], Any]) -> None:
        """Register a callback to be invoked on every state change."""
        self._subscribers.append(callback)

    def unsubscribe(self, callback: Callable[[StateDiff], Any]) -> None:
        """Remove a subscription."""
        self._subscribers.remove(callback)

    def snapshot(self) -> dict[str, Any]:
        """Create a serializable snapshot of the current state (for checkpointing)."""
        return {
            "id": self._id,
            "version": self._version,
            "data": self._data.model_dump(),
            "history_length": len(self._history),
        }

    def rollback(self, to_version: int) -> T:
        """
        Roll back the state to a previous version by replaying diffs.

        This is a simple implementation that replays from the initial state.
        For production use, consider storing periodic snapshots.
        """
        if to_version < 0 or to_version > self._version:
            raise ValueError(f"Invalid version {to_version} (current: {self._version})")

        if to_version == self._version:
            return self._data

        # Rebuild state from initial + diffs up to target version
        # We need to reconstruct — get all diffs up to target version
        target_diffs = [d for d in self._history if d.version <= to_version]

        # Reset to a blank state and apply diffs
        initial_data = type(self._data).model_validate(
            type(self._data).model_construct().model_dump()
        )

        # Apply diffs sequentially
        state_dict = initial_data.model_dump()
        for diff in target_diffs:
            if diff.path in state_dict:
                state_dict[diff.path] = diff.new_value

        self._data = type(self._data).model_validate(state_dict)
        self._version = to_version
        # Trim history
        self._history = [d for d in self._history if d.version <= to_version]

        return self._data

    def clone(self) -> SharedState[T]:
        """Create an independent deep copy of this state."""
        new_state: SharedState[T] = SharedState(self._data.model_copy(deep=True))
        new_state._version = self._version
        new_state._history = copy.deepcopy(self._history)
        return new_state
