scenarios/registry.py

"""
Scenario Registry — Central registry for all DevOps troubleshooting scenarios.

Each scenario defines a broken environment state, success conditions,
and optimal fix sequences for reward shaping.
"""

from __future__ import annotations

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


@dataclass
class Scenario:
    """A single DevOps troubleshooting scenario.

    Attributes:
        id: Unique string identifier (e.g., 'missing_flask').
        level: Difficulty level (1=single-step, 2=two-step, 3=multi-step).
        description: Human-readable description of the broken state.
        initial_state: Dict defining files, env vars, and processes to set up.
        success_condition: Callable that takes last output and returns True if fixed.
        hint_commands: Optimal fix sequence (used for reward shaping).
        error_fingerprint: Regex or keyword to classify the error type.
        setup_commands: Shell commands to create the broken state in the sandbox.
        initial_error_log: The error output the agent sees on reset.
        verification_command: Command to run to check if the fix worked.
    """

    id: str
    level: int
    description: str
    initial_state: Dict
    success_condition: Callable[[str], bool]
    hint_commands: List[str]
    error_fingerprint: str
    setup_commands: List[str] = field(default_factory=list)
    initial_error_log: str = ""
    verification_command: str = ""


class ScenarioRegistry:
    """Central registry holding all scenarios, queryable by level and ID.

    Usage:
        registry = ScenarioRegistry()
        registry.register_defaults()
        scenario = registry.get_random(level=1)
    """

    def __init__(self) -> None:
        """Initialize an empty registry."""
        self._scenarios: Dict[str, Scenario] = {}

    def register(self, scenario: Scenario) -> None:
        """Register a scenario in the registry.

        Args:
            scenario: The Scenario to add.

        Raises:
            ValueError: If a scenario with the same ID already exists.
        """
        if scenario.id in self._scenarios:
            raise ValueError(f"Scenario '{scenario.id}' already registered")
        self._scenarios[scenario.id] = scenario

    def get(self, scenario_id: str) -> Scenario:
        """Get a scenario by its ID.

        Args:
            scenario_id: The unique scenario identifier.

        Returns:
            The matching Scenario.

        Raises:
            KeyError: If the scenario ID is not found.
        """
        if scenario_id not in self._scenarios:
            raise KeyError(f"Scenario '{scenario_id}' not found in registry")
        return self._scenarios[scenario_id]

    def get_by_level(self, level: int) -> List[Scenario]:
        """Get all scenarios at a given difficulty level.

        Args:
            level: Difficulty level (1, 2, or 3).

        Returns:
            List of scenarios at the specified level.
        """
        return [s for s in self._scenarios.values() if s.level == level]

    def get_random(self, level: Optional[int] = None) -> Scenario:
        """Get a random scenario, optionally filtered by level.

        Args:
            level: If provided, filter to this level only.

        Returns:
            A randomly selected Scenario.

        Raises:
            ValueError: If no scenarios match the criteria.
        """
        candidates = self.get_by_level(level) if level else list(self._scenarios.values())
        if not candidates:
            raise ValueError(f"No scenarios available for level={level}")
        return random.choice(candidates)

    def get_all(self) -> List[Scenario]:
        """Get all registered scenarios.

        Returns:
            List of all scenarios in the registry.
        """
        return list(self._scenarios.values())

    def list_ids(self) -> List[str]:
        """Get all registered scenario IDs.

        Returns:
            List of scenario ID strings.
        """
        return list(self._scenarios.keys())

    def register_defaults(self) -> None:
        """Register all built-in scenarios (Levels 1-3).

        This loads all default scenarios from the level modules.
        """
        from scenarios.level1.scenarios import get_level1_scenarios
        from scenarios.level2.scenarios import get_level2_scenarios
        from scenarios.level3.scenarios import get_level3_scenarios

        for scenario in get_level1_scenarios():
            self.register(scenario)
        for scenario in get_level2_scenarios():
            self.register(scenario)
        for scenario in get_level3_scenarios():
            self.register(scenario)

    @property
    def count(self) -> int:
        """Total number of registered scenarios."""
        return len(self._scenarios)