devops_env/state_manager.py

"""
State Manager — Manages observation state for the DevOps RL environment.

Tracks error logs, command history, step counts, and error classifications
across episode steps.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional

from fingerprint.classifier import ErrorFingerprinter, FingerprintResult


@dataclass
class EnvironmentState:
    """Complete state of the environment at a given step.

    Attributes:
        error_log: Last N lines of terminal output (max 2000 chars).
        command_history: Last 10 commands issued.
        step_count: Current step number.
        scenario_id: Identifier for the active scenario.
        error_type: Classified error type from fingerprinting.
        error_confidence: Confidence of the error classification.
        is_terminal: Whether this is a terminal state.
        solved: Whether the scenario was successfully resolved.
    """

    error_log: str = ""
    command_history: List[str] = field(default_factory=list)
    step_count: int = 0
    scenario_id: str = ""
    error_type: str = "unknown"
    error_confidence: float = 0.0
    is_terminal: bool = False
    solved: bool = False

    def to_observation(self) -> Dict:
        """Convert state to an OpenEnv-compatible observation dict.

        Returns:
            Dict with keys: error_log, command_history, step_count,
            scenario_id, error_type, error_confidence, is_terminal, solved.
        """
        return {
            "error_log": self.error_log[:2000],
            "command_history": list(self.command_history[-StateManager.MAX_HISTORY:]),
            "step_count": self.step_count,
            "scenario_id": self.scenario_id,
            "error_type": self.error_type,
            "error_confidence": self.error_confidence,
            "is_terminal": self.is_terminal,
            "solved": self.solved,
        }


class StateManager:
    """Manages environment state transitions across episode steps.

    Handles error log updates, command history tracking, and
    error fingerprinting on each state transition.

    Usage:
        manager = StateManager()
        manager.reset("missing_flask", initial_error_log)
        manager.update(command, new_error_log)
        obs = manager.get_observation()
    """

    MAX_HISTORY: int = 10
    MAX_ERROR_LOG_CHARS: int = 2000

    def __init__(self) -> None:
        """Initialize the state manager."""
        self._state = EnvironmentState()
        self._fingerprinter = ErrorFingerprinter()
        self._prev_error_log: str = ""

    def reset(self, scenario_id: str, initial_error_log: str) -> Dict:
        """Reset state for a new episode.

        Args:
            scenario_id: ID of the scenario being loaded.
            initial_error_log: The initial error output.

        Returns:
            Initial observation dict.
        """
        fp_result = self._fingerprinter.classify(initial_error_log)
        self._state = EnvironmentState(
            error_log=initial_error_log[:self.MAX_ERROR_LOG_CHARS],
            command_history=[],
            step_count=0,
            scenario_id=scenario_id,
            error_type=fp_result.error_type,
            error_confidence=fp_result.confidence,
        )
        self._prev_error_log = initial_error_log
        return self._state.to_observation()

    def update(
        self,
        command: str,
        new_error_log: str,
        is_terminal: bool = False,
        solved: bool = False,
    ) -> Dict:
        """Update state after an action is taken.

        Args:
            command: The command that was executed.
            new_error_log: New terminal output after execution.
            is_terminal: Whether the episode has ended.
            solved: Whether the scenario was solved.

        Returns:
            Updated observation dict.
        """
        self._prev_error_log = self._state.error_log

        # Update command history
        self._state.command_history.append(command)
        if len(self._state.command_history) > self.MAX_HISTORY:
            self._state.command_history = self._state.command_history[-self.MAX_HISTORY:]

        # Update error log and re-classify
        self._state.error_log = new_error_log[:self.MAX_ERROR_LOG_CHARS]
        fp_result = self._fingerprinter.classify(new_error_log)
        self._state.error_type = fp_result.error_type
        self._state.error_confidence = fp_result.confidence

        # Update step and terminal info
        self._state.step_count += 1
        self._state.is_terminal = is_terminal
        self._state.solved = solved

        return self._state.to_observation()

    def get_observation(self) -> Dict:
        """Get the current observation.

        Returns:
            Current observation dict.
        """
        return self._state.to_observation()

    def get_prev_error_log(self) -> str:
        """Get the previous step's error log (for reward computation).

        Returns:
            Previous error log string.
        """
        return self._prev_error_log

    @property
    def state(self) -> EnvironmentState:
        """Access the full state object."""
        return self._state