Blame for src/dashboard/state.py

"""Thread-safe dashboard state publication and management.

This module provides a thread-safe state store for the dashboard that
supports atomic updates, partial state changes, and event notifications
for state subscribers.
"""

from __future__ import annotations

import datetime
import logging
from collections.abc import Callable
from dataclasses import replace
from threading import RLock
from typing import Any, Protocol

from .models import (
    ComponentStatus,
    DashboardSnapshot,
    ExperimentSession,
    ExperimentState,
    HealthSnapshot,
    ParameterChangeRecord,
    ParameterSchema,
    PendingParameterChange,
)

logger = logging.getLogger(__name__)


# ============================================================================
# Type Definitions
# ============================================================================

StateChangeCallback = Callable[[DashboardSnapshot, DashboardSnapshot], None]
"""Callback function type for state change notifications.

Args:
    old_snapshot: The previous dashboard state.
    new_snapshot: The new dashboard state.
"""


# ============================================================================
# Protocols
# ============================================================================


class StateSource(Protocol):
    """Protocol for objects that can provide a dashboard snapshot."""

    def snapshot(self) -> DashboardSnapshot:
        """Return the current dashboard snapshot."""
        ...


# ============================================================================
# State Store
# ============================================================================


class DashboardStateStore:
    """Publish and retrieve immutable dashboard snapshots safely.

    This class provides thread-safe access to the current dashboard state
    with support for atomic updates, partial updates, and event notifications.

    Features:
        - Thread-safe with RLock
        - Atomic full-state updates
        - Partial updates using a builder pattern
        - Event callbacks for state changes
        - Optional state history for debugging
        - Compatible with the simple interface expected by server.py

    Example:
        >>> store = DashboardStateStore()
        >>> store.update(system=SystemMetrics(neurons=1000))
        >>> current = store.snapshot()
        >>> store.on_change(lambda old, new: print(f"State changed: {new.status}"))
    """

    def __init__(
        self,
        initial: DashboardSnapshot | None = None,
        *,
        max_history: int = 0,
        notify_on_update: bool = True,
    ) -> None:
        """Initialize the state store.

        Args:
            initial: Optional initial snapshot. If None, a default is used.
            max_history: Maximum number of historical states to keep.
                0 disables history tracking.
            notify_on_update: Whether to notify callbacks on updates.
        """
        self._lock = RLock()
        self._snapshot = initial or DashboardSnapshot()
        self._version = 0
        self._callbacks: list[StateChangeCallback] = []
        self._history: list[DashboardSnapshot] = []
        self._max_history = max_history
        self._notify_on_update = notify_on_update
        self._update_count = 0

        # Store the initial snapshot in history if enabled
        if max_history > 0:
            self._history.append(self._snapshot)

    # =========================================================================
    # Core Operations
    # =========================================================================

    def publish(self, snapshot: DashboardSnapshot) -> None:
        """Replace the currently visible snapshot atomically.

        Args:
            snapshot: The new dashboard snapshot.
        """
        with self._lock:
            old = self._snapshot
            self._snapshot = snapshot
            self._version += 1
            self._update_count += 1

            # Track history if enabled
            if self._max_history > 0:
                self._history.append(snapshot)
                if len(self._history) > self._max_history:
                    self._history = self._history[-self._max_history :]

            # Notify callbacks
            if self._notify_on_update:
                self._notify_callbacks(old, snapshot)

    def snapshot(self) -> DashboardSnapshot:
        """Return the latest immutable snapshot.

        Returns:
            The current dashboard snapshot.
        """
        with self._lock:
            return self._snapshot

    def update(self, **kwargs: Any) -> DashboardSnapshot:
        """Update specific fields of the dashboard state.

        This is a convenience method that creates a new snapshot with
        updated fields. It uses the dataclass replace() pattern.

        Args:
            **kwargs: Fields to update. Supported keys correspond to
                the fields of DashboardSnapshot.

        Returns:
            The new snapshot that was published.

        Example:
            >>> store.update(status="running", system=SystemMetrics(neurons=1000))
        """
        with self._lock:
            old = self._snapshot

            # Create a new snapshot with updated fields
            new_snapshot = self._update_snapshot(old, **kwargs)

            # Publish the new snapshot
            self._snapshot = new_snapshot
            self._version += 1
            self._update_count += 1

            # Track history if enabled
            if self._max_history > 0:
                self._history.append(new_snapshot)
                if len(self._history) > self._max_history:
                    self._history = self._history[-self._max_history :]

            # Notify callbacks
            if self._notify_on_update:
                self._notify_callbacks(old, new_snapshot)

            return new_snapshot

    def _update_snapshot(
        self, snapshot: DashboardSnapshot, **kwargs: Any
    ) -> DashboardSnapshot:
        """Create a new snapshot with updated fields."""
        # Get all fields of DashboardSnapshot
        field_names = [
            "system",
            "learning",
            "storage",
            "self_organization",
            "homeostasis",
            "structural",
            "spikes",
            "network",
            "language_organ",
            "knowledge_intake",
            "signal_metrics",
            "experiment",
            "embodiment",
            "runtime",
            "components",
            "parameters",
            "pending_changes",
            "change_history",
            "experiment_state",
            "health",
            "status",
            "version",
        ]

        # Build replacement dict
        replacements: dict[str, Any] = {}
        for name in field_names:
            if name in kwargs:
                replacements[name] = kwargs[name]

        # Update the snapshot using dataclass replace
        return replace(snapshot, **replacements)

    def update_component(self, status: ComponentStatus) -> DashboardSnapshot:
        """Update or insert a single component status.

        Args:
            status: The new component status.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            components = dict(current.components or {})
            components[status.component] = status
            return self.update(components=components)

    def update_parameter(self, parameter: ParameterSchema) -> DashboardSnapshot:
        """Update or insert a single parameter schema entry.

        Args:
            parameter: The new parameter schema entry.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            parameters = dict(current.parameters or {})
            parameters[parameter.name] = parameter
            return self.update(parameters=parameters)

    def set_health(self, health: HealthSnapshot) -> DashboardSnapshot:
        """Set the aggregated health snapshot.

        Args:
            health: The new health snapshot.

        Returns:
            The new snapshot.
        """
        return self.update(health=health)

    def set_pending_change(
        self,
        change: PendingParameterChange,
    ) -> DashboardSnapshot:
        """Record or update a pending parameter change.

        Args:
            change: The pending change to store.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            pending = dict(current.pending_changes or {})
            pending[change.name] = change
            return self.update(pending_changes=pending)

    def remove_pending_change(self, name: str) -> DashboardSnapshot:
        """Remove a pending parameter change.

        Args:
            name: Parameter name whose pending change should be removed.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            pending = dict(current.pending_changes or {})
            pending.pop(name, None)
            return self.update(pending_changes=pending)

    def clear_pending_changes(self) -> DashboardSnapshot:
        """Clear all pending parameter changes.

        Returns:
            The new snapshot.
        """
        return self.update(pending_changes={})

    def append_change_history(
        self,
        record: ParameterChangeRecord,
    ) -> DashboardSnapshot:
        """Append a change record to the history.

        Args:
            record: The record to append.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            history = list(current.change_history)
            history.append(record)
            return self.update(change_history=tuple(history))

    def set_experiment_mode(
        self,
        mode: str,
    ) -> DashboardSnapshot:
        """Set the current experiment mode.

        Args:
            mode: One of 'operator', 'experiment', 'debug'.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            new_state = ExperimentState(
                current_mode=mode,
                active_session=current.experiment_state.active_session,
                sessions=current.experiment_state.sessions,
            )
            return self.update(experiment_state=new_state)

    def start_experiment_session(
        self,
        session: ExperimentSession,
    ) -> DashboardSnapshot:
        """Start a new experiment/debug session.

        Any currently active session is closed automatically.

        Args:
            session: The session to start.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            old_state = current.experiment_state
            sessions = list(old_state.sessions)

            # Close any active session first
            if old_state.active_session is not None:
                now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
                closed = ExperimentSession(
                    session_id=old_state.active_session.session_id,
                    mode=old_state.active_session.mode,
                    hypothesis=old_state.active_session.hypothesis,
                    notes=old_state.active_session.notes,
                    start_tick=old_state.active_session.start_tick,
                    end_tick=session.start_tick,
                    start_time=old_state.active_session.start_time,
                    end_time=now,
                    config_snapshot=old_state.active_session.config_snapshot,
                    active=False,
                )
                sessions.append(closed)

            sessions.append(session)
            new_state = ExperimentState(
                current_mode=session.mode,
                active_session=session,
                sessions=tuple(sessions),
            )
            return self.update(experiment_state=new_state)

    def stop_experiment_session(
        self,
        end_tick: int = 0,
    ) -> DashboardSnapshot:
        """Stop the active experiment/debug session.

        Args:
            end_tick: Simulation tick at session stop.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            old_state = current.experiment_state
            active = old_state.active_session

            if active is None:
                return current

            now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
            sessions = [s for s in old_state.sessions if not s.active]
            closed = ExperimentSession(
                session_id=active.session_id,
                mode=active.mode,
                hypothesis=active.hypothesis,
                notes=active.notes,
                start_tick=active.start_tick,
                end_tick=end_tick,
                start_time=active.start_time,
                end_time=now,
                config_snapshot=active.config_snapshot,
                active=False,
            )
            sessions.append(closed)

            new_state = ExperimentState(
                current_mode=old_state.current_mode,
                active_session=None,
                sessions=tuple(sessions),
            )
            return self.update(experiment_state=new_state)

    def add_experiment_note(
        self,
        note: str,
    ) -> DashboardSnapshot:
        """Append a note to the active experiment session.

        Args:
            note: Note text to append.

        Returns:
            The new snapshot.
        """
        with self._lock:
            current = self._snapshot
            old_state = current.experiment_state
            active = old_state.active_session

            if active is None:
                return current

            now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat()
            updated = ExperimentSession(
                session_id=active.session_id,
                mode=active.mode,
                hypothesis=active.hypothesis,
                notes=(*active.notes, f"[{now}] {note}"),
                start_tick=active.start_tick,
                end_tick=active.end_tick,
                start_time=active.start_time,
                end_time=active.end_time,
                config_snapshot=active.config_snapshot,
                active=active.active,
            )

            sessions = [
                s if s.session_id != active.session_id else updated
                for s in old_state.sessions
            ]
            # Ensure active session is present in sessions list
            if not any(s.session_id == active.session_id for s in sessions):
                sessions.append(updated)

            new_state = ExperimentState(
                current_mode=old_state.current_mode,
                active_session=updated,
                sessions=tuple(sessions),
            )
            return self.update(experiment_state=new_state)

    # =========================================================================
    # Event System
    # =========================================================================

    def on_change(self, callback: StateChangeCallback) -> None:
        """Register a callback for state changes.

        Args:
            callback: Function to call when the state changes.
                Receives old and new snapshots as arguments.

        Raises:
            ValueError: If the callback is already registered.
        """
        with self._lock:
            if callback in self._callbacks:
                raise ValueError("Callback already registered")
            self._callbacks.append(callback)

    def remove_callback(self, callback: StateChangeCallback) -> bool:
        """Remove a previously registered callback.

        Args:
            callback: The callback to remove.

        Returns:
            True if the callback was removed, False if not found.
        """
        with self._lock:
            try:
                self._callbacks.remove(callback)
                return True
            except ValueError:
                return False

    def clear_callbacks(self) -> None:
        """Remove all registered callbacks."""
        with self._lock:
            self._callbacks.clear()

    def _notify_callbacks(self, old: DashboardSnapshot, new: DashboardSnapshot) -> None:
        """Notify all registered callbacks of a state change."""
        # Make a copy of the callback list to avoid issues if callbacks
        # modify the list during iteration.
        with self._lock:
            callbacks = list(self._callbacks)

        for callback in callbacks:
            try:
                callback(old, new)
            except Exception as e:
                logger.warning(f"State change callback failed: {e}")

    # =========================================================================
    # Query Methods
    # =========================================================================

    def version(self) -> int:
        """Return the current state version number.

        The version increments on every state change.
        """
        with self._lock:
            return self._version

    def update_count(self) -> int:
        """Return the total number of state updates."""
        with self._lock:
            return self._update_count

    def get_history(self, limit: int | None = None) -> list[DashboardSnapshot]:
        """Return the state history.

        Args:
            limit: Maximum number of historical states to return.
                If None, returns all history.

        Returns:
            List of historical snapshots (newest first).
        """
        with self._lock:
            if limit is None or limit <= 0:
                return list(reversed(self._history))
            return list(reversed(self._history[-limit:]))

    def get_version_at(self, version: int) -> DashboardSnapshot | None:
        """Get the state at a specific version.

        Args:
            version: Version number to retrieve.

        Returns:
            The snapshot at that version, or None if not found.
        """
        with self._lock:
            if version < 0 or version >= len(self._history):
                return None
            return self._history[version]

    # =========================================================================
    # Convenience Methods
    # =========================================================================

    def set_status(self, status: str) -> DashboardSnapshot:
        """Update the dashboard status.

        Args:
            status: New status string (e.g., "idle", "running", "paused").

        Returns:
            The new snapshot.
        """
        return self.update(status=status)

    def set_version(self, version: str) -> DashboardSnapshot:
        """Update the system version.

        Args:
            version: New version string.

        Returns:
            The new snapshot.
        """
        return self.update(version=version)

    def update_metrics(
        self,
        system: Any = None,
        learning: Any = None,
        storage: Any = None,
        homeostasis: Any = None,
        structural: Any = None,
        spikes: Any = None,
        network: Any = None,
        language_organ: Any = None,
        knowledge_intake: Any = None,
        experiment: Any = None,
        embodiment: Any = None,
    ) -> DashboardSnapshot:
        """Update multiple metric fields at once.

        This is a convenience method for updating all metric-related
        fields in a single call.

        Args:
            system: SystemMetrics instance.
            learning: LearningMetrics instance.
            storage: StorageMetrics instance.
            homeostasis: HomeostasisMetrics instance.
            structural: StructuralMetrics instance.
            spikes: SpikeMetrics instance.
            network: NetworkMetrics instance.
            language_organ: LanguageOrganMetrics instance.
            knowledge_intake: KnowledgeIntakeMetrics instance.
            experiment: ExperimentMetrics instance.
            embodiment: EmbodimentMetrics instance.

        Returns:
            The new snapshot.
        """
        updates: dict[str, Any] = {}
        if system is not None:
            updates["system"] = system
        if learning is not None:
            updates["learning"] = learning
        if storage is not None:
            updates["storage"] = storage
        if homeostasis is not None:
            updates["homeostasis"] = homeostasis
        if structural is not None:
            updates["structural"] = structural
        if spikes is not None:
            updates["spikes"] = spikes
        if network is not None:
            updates["network"] = network
        if language_organ is not None:
            updates["language_organ"] = language_organ
        if knowledge_intake is not None:
            updates["knowledge_intake"] = knowledge_intake
        if experiment is not None:
            updates["experiment"] = experiment
        if embodiment is not None:
            updates["embodiment"] = embodiment

        return self.update(**updates)

    # =========================================================================
    # Testing Helpers
    # =========================================================================

    def reset(self, initial: DashboardSnapshot | None = None) -> None:
        """Reset the state store to a clean state.

        This is primarily useful for testing.

        Args:
            initial: Optional initial snapshot. If None, uses a default.
        """
        with self._lock:
            self._snapshot = initial or DashboardSnapshot()
            self._version = 0
            self._update_count = 0
            self._history = []
            self._callbacks.clear()
            if self._max_history > 0:
                self._history.append(self._snapshot)


# ============================================================================
# Factory Functions
# ============================================================================


def create_state_store(
    initial: DashboardSnapshot | None = None,
    *,
    with_history: bool = True,
    history_limit: int = 100,
    notify: bool = True,
) -> DashboardStateStore:
    """Create a configured DashboardStateStore.

    Args:
        initial: Optional initial snapshot.
        with_history: Whether to enable history tracking.
        history_limit: Maximum number of historical states to keep.
        notify: Whether to notify callbacks on updates.

    Returns:
        Configured DashboardStateStore instance.

    Example:
        >>> store = create_state_store(with_history=True, history_limit=50)
    """
    return DashboardStateStore(
        initial=initial,
        max_history=history_limit if with_history else 0,
        notify_on_update=notify,
    )


# ============================================================================
# Global State Manager (Optional)
# ============================================================================


class StateManager:
    """Global state manager for the dashboard.

    This is a singleton-like manager that provides access to the
    dashboard state store. It can be used as a central point for
    state management across different parts of the application.
    """

    _instance: StateManager | None = None
    _store: DashboardStateStore | None = None
    _config_dict: dict[str, Any] | None = None

    def __new__(cls) -> StateManager:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def initialize(
        self,
        initial: DashboardSnapshot | None = None,
        *,
        with_history: bool = True,
        history_limit: int = 100,
    ) -> None:
        """Initialize the global state store.

        Args:
            initial: Optional initial snapshot.
            with_history: Whether to enable history tracking.
            history_limit: Maximum number of historical states to keep.
        """
        if self._store is None:
            self._store = create_state_store(
                initial=initial,
                with_history=with_history,
                history_limit=history_limit,
            )

    def set_config(self, config_dict: dict[str, Any] | None) -> None:
        """Store the runtime configuration for health derivation."""
        self._config_dict = config_dict

    @property
    def store(self) -> DashboardStateStore:
        """Get the global state store."""
        if self._store is None:
            # Auto-initialize with default settings
            self._store = create_state_store()
        return self._store

    def publish(self, snapshot: DashboardSnapshot) -> None:
        """Publish a new snapshot."""
        self.store.publish(snapshot)

    def snapshot(self) -> DashboardSnapshot:
        """Get the current snapshot."""
        return self.store.snapshot()

    def update(self, **kwargs: Any) -> DashboardSnapshot:
        """Update the current state."""
        return self.store.update(**kwargs)


# ============================================================================
# Global Instance
# ============================================================================

# Convenience global instance for simple use cases
_state_manager = StateManager()


def get_state_store() -> DashboardStateStore:
    """Get the global state store instance.

    Returns:
        The global DashboardStateStore instance.
    """
    return _state_manager.store


def publish_state(snapshot: DashboardSnapshot) -> None:
    """Publish a state update to the global store.

    The snapshot is automatically enriched with derived components,
    parameters and health so that every consumer sees a complete view.
    """
    from .health_builder import enrich_snapshot

    _state_manager.publish(
        enrich_snapshot(
            snapshot,
            _state_manager._config_dict,  # pyright: ignore[reportPrivateUsage]
        )
    )


def set_dashboard_config(config_dict: dict[str, Any] | None) -> None:
    """Store the runtime configuration used for health derivation."""
    _state_manager.set_config(config_dict)


def get_current_state() -> DashboardSnapshot:
    """Get the current state from the global store."""
    return _state_manager.snapshot()


# ============================================================================
# Example Usage
# ============================================================================

if __name__ == "__main__":
    # Basic usage
    store = DashboardStateStore()

    # Register a callback
    def on_state_change(old: DashboardSnapshot, new: DashboardSnapshot) -> None:
        print(f"State changed: {old.status} -> {new.status}")

    store.on_change(on_state_change)

    # Update state
    store.set_status("running")
    store.set_version("0.6.0")

    # Get current state
    current = store.snapshot()
    print(f"Current status: {current.status}")
    print(f"Version: {current.version}")