devops_env/env.py

"""
DevOpsEnv — OpenEnv-style RL environment for terminal troubleshooting.

The agent observes broken Linux/Python environment states, issues shell commands,
and receives multi-signal rewards. Episodes are bounded by max steps, success,
or dangerous command detection.
"""

from __future__ import annotations

import random
from typing import Any, Dict, List, Optional, Tuple

from devops_env.state_manager import StateManager
from executor.docker_executor import DockerExecutor, ExecutionResult
from fingerprint.classifier import ErrorFingerprinter
from rewards.engine import RewardEngine
from scenarios.registry import Scenario, ScenarioRegistry


class DevOpsEnv:
    """OpenEnv-style environment for DevOps troubleshooting with RL.

    The agent receives an error log and command history as observations,
    outputs a shell command, and gets a reward based on whether the
    command moved toward fixing the issue.

    Attributes:
        metadata: Environment metadata dict.
        max_steps: Maximum steps per episode before truncation.
    """

    metadata = {"render_modes": ["human"]}

    def __init__(
        self,
        scenario_registry: ScenarioRegistry | None = None,
        executor: DockerExecutor | None = None,
        max_steps: int = 10,
        render_mode: str | None = None,
        target_level: int | None = None,
        target_scenario: str | None = None,
    ) -> None:
        """Initialize the DevOps environment.

        Args:
            scenario_registry: Registry of available scenarios. Creates default if None.
            executor: Docker executor for running commands. Creates default if None.
            max_steps: Maximum steps per episode.
            render_mode: Render mode.
            target_level: If set, only sample scenarios from this level.
            target_scenario: If set, always use this specific scenario.
        """
        self.max_steps = max_steps
        self.render_mode = render_mode
        self.target_level = target_level
        self.target_scenario = target_scenario

        # Initialize components
        if scenario_registry is None:
            self.registry = ScenarioRegistry()
            self.registry.register_defaults()
        else:
            self.registry = scenario_registry

        self.executor = executor or DockerExecutor(use_local_fallback=True)
        self.state_manager = StateManager()
        self.reward_engine = RewardEngine()
        self.fingerprinter = ErrorFingerprinter()

        # Episode state
        self._current_scenario: Optional[Scenario] = None
        self._step_count: int = 0
        self._episode_reward: float = 0.0
        self._episode_steps: List[Dict] = []
        self._done: bool = False

        # OpenEnv schemas (documented shape constraints for API clients)
        self.observation_schema: Dict[str, str] = {
            "error_log": "str(max=2000)",
            "command_history": "List[str](max_items=10)",
            "step_count": f"int(0..{max_steps})",
            "scenario_id": "str(max=100)",
            "error_type": "str(max=50)",
            "error_confidence": "float(0.0..1.0)",
            "is_terminal": "bool",
            "solved": "bool",
        }
        self.action_schema: Dict[str, str] = {
            "command": "str(max=500)",
        }

    def reset(
        self,
        seed: int | None = None,
        options: Dict[str, Any] | None = None,
    ) -> Tuple[Dict, Dict]:
        """Reset the environment for a new episode.

        Loads a random scenario (or the target scenario), sets up the
        Docker sandbox, and returns the initial observation.

        Args:
            seed: Random seed for reproducibility.
            options: Additional options (e.g., {"scenario_id": "missing_flask"}).

        Returns:
            Tuple of (observation, info_dict).
        """
        if seed is not None:
            random.seed(seed)

        # Select scenario
        scenario_id = None
        if options and "scenario_id" in options:
            scenario_id = options["scenario_id"]
        elif self.target_scenario:
            scenario_id = self.target_scenario

        if scenario_id:
            self._current_scenario = self.registry.get(scenario_id)
        else:
            self._current_scenario = self.registry.get_random(level=self.target_level)

        # Reset episode state
        self._step_count = 0
        self._episode_reward = 0.0
        self._episode_steps = []
        self._done = False

        # Set up Docker sandbox
        try:
            self.executor.stop_container()
            self.executor.start_container(self._current_scenario.setup_commands)
        except Exception:
            # Continue with local fallback
            pass

        # Initialize state with the scenario's error log
        obs = self.state_manager.reset(
            scenario_id=self._current_scenario.id,
            initial_error_log=self._current_scenario.initial_error_log,
        )

        info = {
            "scenario_id": self._current_scenario.id,
            "level": self._current_scenario.level,
            "description": self._current_scenario.description,
            "error_type": obs["error_type"],
        }

        return obs, info

    def step(self, action: str) -> Tuple[Dict, float, bool, bool, Dict]:
        """Execute one step in the environment.

        Args:
            action: Shell command to execute.

        Returns:
            Tuple of (observation, reward, terminated, truncated, info).
        """
        if self._done:
            raise RuntimeError("Episode is done. Call reset() first.")

        assert self._current_scenario is not None

        self._step_count += 1
        action = action.strip()

        # Execute command in sandbox
        result = self.executor.execute(action)

        # Build new error log from execution output
        if result.blocked:
            new_error_log = f"COMMAND BLOCKED: {result.block_reason}"
        elif result.timed_out:
            new_error_log = "COMMAND TIMED OUT after 30 seconds."
        else:
            new_error_log = ""
            if result.stdout:
                new_error_log += result.stdout
            if result.stderr:
                new_error_log += ("\n" if new_error_log else "") + result.stderr
            if not new_error_log:
                new_error_log = f"Command completed with exit code {result.exit_code}"

        # Get previous error log for reward computation
        prev_error_log = self.state_manager.get_prev_error_log()

        # Compute reward
        all_commands = list(self.state_manager.state.command_history) + [action]
        reward, reward_breakdown = self.reward_engine.compute_reward(
            action=action,
            result=result,
            scenario=self._current_scenario,
            step_count=self._step_count,
            command_history=all_commands,
            prev_error_log=prev_error_log,
            curr_error_log=new_error_log,
        )

        # Check termination conditions
        combined_output = f"{result.stdout}\n{result.stderr}".strip()
        solved = self._current_scenario.success_condition(combined_output)
        is_dangerous_block = result.blocked and "dangerous" in result.block_reason.lower()
        terminated = solved or is_dangerous_block
        truncated = self._step_count >= self.max_steps

        # Update state
        obs = self.state_manager.update(
            command=action,
            new_error_log=new_error_log,
            is_terminal=terminated or truncated,
            solved=solved,
        )

        # Track episode
        self._episode_reward += reward
        self._episode_steps.append({
            "step": self._step_count,
            "action": action,
            "result": {
                "stdout": result.stdout[:1000],
                "stderr": result.stderr[:1000],
                "exit_code": result.exit_code,
                "timed_out": result.timed_out,
                "blocked": result.blocked,
            },
            "reward": reward,
            "reward_breakdown": reward_breakdown,
            "error_type": obs["error_type"],
            "observation": {
                "error_log": obs["error_log"][:500],
                "command_history": obs["command_history"],
                "step_count": obs["step_count"],
            },
        })

        self._done = terminated or truncated

        info = {
            "scenario_id": self._current_scenario.id,
            "level": self._current_scenario.level,
            "solved": solved,
            "step_count": obs["step_count"],
            "episode_reward": self._episode_reward,
            "reward_breakdown": reward_breakdown,
            "error_type": obs["error_type"],
            "execution_result": {
                "exit_code": result.exit_code,
                "blocked": result.blocked,
                "timed_out": result.timed_out,
            },
        }

        if self._done:
            info["episode_steps"] = self._episode_steps

        return obs, reward, terminated, truncated, info

    def get_episode_summary(self) -> Dict:
        """Get a summary of the current/last episode.

        Returns:
            Dict with episode metadata and step details.
        """
        return {
            "scenario_id": self._current_scenario.id if self._current_scenario else None,
            "level": self._current_scenario.level if self._current_scenario else None,
            "steps": self._episode_steps,
            "total_reward": self._episode_reward,
            "solved": self.state_manager.state.solved,
            "total_steps": self._step_count,
        }

    def render(self) -> None:
        """Render the current environment state (human-readable)."""
        if self.render_mode != "human":
            return
        state = self.state_manager.state
        print(f"\n{'='*60}")
        print(f"Scenario: {state.scenario_id} | Step: {state.step_count}")
        print(f"Error Type: {state.error_type}")
        print(f"{'─'*60}")
        print(f"Error Log:\n{state.error_log[:500]}")
        print(f"{'─'*60}")
        if state.command_history:
            print(f"Commands: {state.command_history}")
        print(f"{'='*60}\n")

    def close(self) -> None:
        """Clean up resources."""
        self.executor.stop_container()