Atomizer/optimization_engine/context/session_state.py

"""
Atomizer Session State - Context Isolation Management

Part of the ACE (Agentic Context Engineering) implementation for Atomizer.

Implements the "Write-Select-Compress-Isolate" pattern:
- Exposed fields are sent to LLM at every turn
- Isolated fields are accessed selectively when needed
- Automatic compression of old data

This ensures efficient context usage while maintaining
access to full historical data when needed.
"""

from typing import Dict, List, Optional, Any
from datetime import datetime
from enum import Enum
from dataclasses import dataclass, field
import json
from pathlib import Path


class TaskType(Enum):
    """Types of tasks Claude can perform in Atomizer."""
    CREATE_STUDY = "create_study"
    RUN_OPTIMIZATION = "run_optimization"
    MONITOR_PROGRESS = "monitor_progress"
    ANALYZE_RESULTS = "analyze_results"
    DEBUG_ERROR = "debug_error"
    CONFIGURE_SETTINGS = "configure_settings"
    EXPORT_DATA = "export_data"
    NEURAL_ACCELERATION = "neural_acceleration"


@dataclass
class ExposedState:
    """
    State exposed to LLM at every turn.

    Keep this minimal - only what's needed for immediate context.
    Everything here counts against token budget every turn.
    """

    # Current task context
    task_type: Optional[TaskType] = None
    current_objective: str = ""

    # Recent history (compressed)
    recent_actions: List[str] = field(default_factory=list)
    recent_errors: List[str] = field(default_factory=list)

    # Active study summary
    study_name: Optional[str] = None
    study_status: str = "unknown"
    trials_completed: int = 0
    trials_total: int = 0
    best_value: Optional[float] = None
    best_trial: Optional[int] = None

    # Playbook excerpt (most relevant items)
    active_playbook_items: List[str] = field(default_factory=list)

    # Constraints for context size
    MAX_ACTIONS: int = 10
    MAX_ERRORS: int = 5
    MAX_PLAYBOOK_ITEMS: int = 15


@dataclass
class IsolatedState:
    """
    State isolated from LLM - accessed selectively.

    This data is NOT included in every context window.
    Load specific fields when explicitly needed.
    """

    # Full optimization history (can be large)
    full_trial_history: List[Dict[str, Any]] = field(default_factory=list)

    # NX session state (heavy, complex)
    nx_model_path: Optional[str] = None
    nx_expressions: Dict[str, Any] = field(default_factory=dict)
    nx_sim_path: Optional[str] = None

    # Neural network cache
    neural_predictions: Dict[str, float] = field(default_factory=dict)
    surrogate_model_path: Optional[str] = None

    # Full playbook (loaded on demand)
    full_playbook_path: Optional[str] = None

    # Debug information
    last_solver_output: str = ""
    last_f06_content: str = ""
    last_solver_returncode: Optional[int] = None

    # Configuration snapshots
    optimization_config: Dict[str, Any] = field(default_factory=dict)
    study_config: Dict[str, Any] = field(default_factory=dict)


@dataclass
class AtomizerSessionState:
    """
    Complete session state with exposure control.

    The exposed state is automatically injected into every LLM context.
    The isolated state is accessed only when explicitly needed.

    Usage:
        session = AtomizerSessionState(session_id="session_001")
        session.exposed.task_type = TaskType.CREATE_STUDY
        session.add_action("Created study directory")

        # Get context for LLM
        context = session.get_llm_context()

        # Access isolated data when needed
        f06 = session.load_isolated_data("last_f06_content")
    """

    session_id: str
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    last_updated: str = field(default_factory=lambda: datetime.now().isoformat())

    exposed: ExposedState = field(default_factory=ExposedState)
    isolated: IsolatedState = field(default_factory=IsolatedState)

    def get_llm_context(self) -> str:
        """
        Generate context string for LLM consumption.

        Only includes exposed state - isolated state requires
        explicit access via load_isolated_data().

        Returns:
            Formatted markdown context string
        """
        lines = [
            "## Current Session State",
            "",
            f"**Task**: {self.exposed.task_type.value if self.exposed.task_type else 'Not set'}",
            f"**Objective**: {self.exposed.current_objective or 'None specified'}",
            "",
        ]

        # Study context
        if self.exposed.study_name:
            progress = ""
            if self.exposed.trials_total > 0:
                pct = (self.exposed.trials_completed / self.exposed.trials_total) * 100
                progress = f" ({pct:.0f}%)"

            lines.extend([
                f"### Active Study: {self.exposed.study_name}",
                f"- Status: {self.exposed.study_status}",
                f"- Trials: {self.exposed.trials_completed}/{self.exposed.trials_total}{progress}",
            ])

            if self.exposed.best_value is not None:
                lines.append(f"- Best: {self.exposed.best_value:.6g} (trial #{self.exposed.best_trial})")
            lines.append("")

        # Recent actions
        if self.exposed.recent_actions:
            lines.append("### Recent Actions")
            for action in self.exposed.recent_actions[-5:]:
                lines.append(f"- {action}")
            lines.append("")

        # Recent errors (highlight these)
        if self.exposed.recent_errors:
            lines.append("### Recent Errors (address these)")
            for error in self.exposed.recent_errors:
                lines.append(f"- {error}")
            lines.append("")

        # Relevant playbook items
        if self.exposed.active_playbook_items:
            lines.append("### Relevant Knowledge")
            for item in self.exposed.active_playbook_items:
                lines.append(f"- {item}")
            lines.append("")

        return "\n".join(lines)

    def add_action(self, action: str) -> None:
        """
        Record an action (auto-compresses old actions).

        Args:
            action: Description of the action taken
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        self.exposed.recent_actions.append(f"[{timestamp}] {action}")

        # Compress if over limit
        if len(self.exposed.recent_actions) > self.exposed.MAX_ACTIONS:
            # Keep first, summarize middle, keep last 5
            first = self.exposed.recent_actions[0]
            last_five = self.exposed.recent_actions[-5:]
            middle_count = len(self.exposed.recent_actions) - 6

            self.exposed.recent_actions = (
                [first] +
                [f"... ({middle_count} earlier actions)"] +
                last_five
            )

        self.last_updated = datetime.now().isoformat()

    def add_error(self, error: str, error_type: str = "") -> None:
        """
        Record an error for LLM attention.

        Errors are preserved more aggressively than actions
        because they need to be addressed.

        Args:
            error: Error message
            error_type: Optional error classification
        """
        prefix = f"[{error_type}] " if error_type else ""
        self.exposed.recent_errors.append(f"{prefix}{error}")

        # Keep most recent errors
        self.exposed.recent_errors = self.exposed.recent_errors[-self.exposed.MAX_ERRORS:]
        self.last_updated = datetime.now().isoformat()

    def clear_errors(self) -> None:
        """Clear all recorded errors (after they're addressed)."""
        self.exposed.recent_errors = []
        self.last_updated = datetime.now().isoformat()

    def update_study_status(
        self,
        name: str,
        status: str,
        trials_completed: int,
        trials_total: int,
        best_value: Optional[float] = None,
        best_trial: Optional[int] = None
    ) -> None:
        """
        Update the study status in exposed state.

        Args:
            name: Study name
            status: Current status (running, completed, failed, etc.)
            trials_completed: Number of completed trials
            trials_total: Total planned trials
            best_value: Best objective value found
            best_trial: Trial number with best value
        """
        self.exposed.study_name = name
        self.exposed.study_status = status
        self.exposed.trials_completed = trials_completed
        self.exposed.trials_total = trials_total
        self.exposed.best_value = best_value
        self.exposed.best_trial = best_trial
        self.last_updated = datetime.now().isoformat()

    def set_playbook_items(self, items: List[str]) -> None:
        """
        Set the active playbook items for context.

        Args:
            items: List of playbook item context strings
        """
        self.exposed.active_playbook_items = items[:self.exposed.MAX_PLAYBOOK_ITEMS]
        self.last_updated = datetime.now().isoformat()

    def load_isolated_data(self, key: str) -> Any:
        """
        Explicitly load isolated data when needed.

        Use this when you need access to heavy data that
        shouldn't be in every context window.

        Args:
            key: Attribute name in IsolatedState

        Returns:
            The isolated data value, or None if not found
        """
        return getattr(self.isolated, key, None)

    def set_isolated_data(self, key: str, value: Any) -> None:
        """
        Set isolated data.

        Args:
            key: Attribute name in IsolatedState
            value: Value to set
        """
        if hasattr(self.isolated, key):
            setattr(self.isolated, key, value)
            self.last_updated = datetime.now().isoformat()

    def add_trial_to_history(self, trial_data: Dict[str, Any]) -> None:
        """
        Add a trial to the full history (isolated state).

        Args:
            trial_data: Dictionary with trial information
        """
        trial_data["recorded_at"] = datetime.now().isoformat()
        self.isolated.full_trial_history.append(trial_data)
        self.last_updated = datetime.now().isoformat()

    def get_trial_history_summary(self, last_n: int = 10) -> List[Dict[str, Any]]:
        """
        Get summary of recent trials from isolated history.

        Args:
            last_n: Number of recent trials to return

        Returns:
            List of trial summary dictionaries
        """
        return self.isolated.full_trial_history[-last_n:]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "session_id": self.session_id,
            "created_at": self.created_at,
            "last_updated": self.last_updated,
            "exposed": {
                "task_type": self.exposed.task_type.value if self.exposed.task_type else None,
                "current_objective": self.exposed.current_objective,
                "recent_actions": self.exposed.recent_actions,
                "recent_errors": self.exposed.recent_errors,
                "study_name": self.exposed.study_name,
                "study_status": self.exposed.study_status,
                "trials_completed": self.exposed.trials_completed,
                "trials_total": self.exposed.trials_total,
                "best_value": self.exposed.best_value,
                "best_trial": self.exposed.best_trial,
                "active_playbook_items": self.exposed.active_playbook_items
            },
            "isolated": {
                "nx_model_path": self.isolated.nx_model_path,
                "nx_sim_path": self.isolated.nx_sim_path,
                "surrogate_model_path": self.isolated.surrogate_model_path,
                "full_playbook_path": self.isolated.full_playbook_path,
                "trial_history_count": len(self.isolated.full_trial_history)
            }
        }

    def save(self, path: Path) -> None:
        """
        Save session state to JSON.

        Note: Full trial history is saved to a separate file
        to keep the main state file manageable.

        Args:
            path: Path to save state file
        """
        path.parent.mkdir(parents=True, exist_ok=True)

        # Save main state
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(self.to_dict(), f, indent=2)

        # Save trial history separately if large
        if len(self.isolated.full_trial_history) > 0:
            history_path = path.with_suffix('.history.json')
            with open(history_path, 'w', encoding='utf-8') as f:
                json.dump(self.isolated.full_trial_history, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> "AtomizerSessionState":
        """
        Load session state from JSON.

        Args:
            path: Path to state file

        Returns:
            Loaded session state (or new state if file doesn't exist)
        """
        if not path.exists():
            return cls(session_id=f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}")

        with open(path, encoding='utf-8') as f:
            data = json.load(f)

        state = cls(
            session_id=data.get("session_id", "unknown"),
            created_at=data.get("created_at", datetime.now().isoformat()),
            last_updated=data.get("last_updated", datetime.now().isoformat())
        )

        # Load exposed state
        exposed = data.get("exposed", {})
        if exposed.get("task_type"):
            state.exposed.task_type = TaskType(exposed["task_type"])
        state.exposed.current_objective = exposed.get("current_objective", "")
        state.exposed.recent_actions = exposed.get("recent_actions", [])
        state.exposed.recent_errors = exposed.get("recent_errors", [])
        state.exposed.study_name = exposed.get("study_name")
        state.exposed.study_status = exposed.get("study_status", "unknown")
        state.exposed.trials_completed = exposed.get("trials_completed", 0)
        state.exposed.trials_total = exposed.get("trials_total", 0)
        state.exposed.best_value = exposed.get("best_value")
        state.exposed.best_trial = exposed.get("best_trial")
        state.exposed.active_playbook_items = exposed.get("active_playbook_items", [])

        # Load isolated state metadata
        isolated = data.get("isolated", {})
        state.isolated.nx_model_path = isolated.get("nx_model_path")
        state.isolated.nx_sim_path = isolated.get("nx_sim_path")
        state.isolated.surrogate_model_path = isolated.get("surrogate_model_path")
        state.isolated.full_playbook_path = isolated.get("full_playbook_path")

        # Load trial history from separate file if exists
        history_path = path.with_suffix('.history.json')
        if history_path.exists():
            with open(history_path, encoding='utf-8') as f:
                state.isolated.full_trial_history = json.load(f)

        return state


# Convenience functions for session management
_active_session: Optional[AtomizerSessionState] = None


def get_session() -> AtomizerSessionState:
    """
    Get the active session state.

    Creates a new session if none exists.

    Returns:
        The active AtomizerSessionState
    """
    global _active_session
    if _active_session is None:
        _active_session = AtomizerSessionState(
            session_id=f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        )
    return _active_session


def set_session(session: AtomizerSessionState) -> None:
    """
    Set the active session.

    Args:
        session: Session state to make active
    """
    global _active_session
    _active_session = session


def clear_session() -> None:
    """Clear the active session."""
    global _active_session
    _active_session = None