AI-Driven Remediation Testing

"""Finite State Machine for scenario execution."""

import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

from ..models.scenario import Scenario


class State(Enum):
    """Scenario execution states."""

    INIT = "INIT"
    PRECHECK = "PRECHECK"
    FAULT_INJECT = "FAULT_INJECT"
    STABILIZE = "STABILIZE"
    ASSISTANT_RCA = "ASSISTANT_RCA"
    EVAL_RCA = "EVAL_RCA"
    ASSISTANT_REMEDY = "ASSISTANT_REMEDY"
    EVAL_REMEDY = "EVAL_REMEDY"
    EXECUTE_REMEDY = "EXECUTE_REMEDY"
    VERIFY = "VERIFY"
    PASS = "PASS"
    FAIL = "FAIL"
    CLEANUP = "CLEANUP"


@dataclass
class StepResult:
    """Result of a single execution step."""

    state: State
    success: bool
    message: str
    score: Optional[float] = None
    artifacts: Dict[str, str] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    timestamp: float = field(default_factory=time.time)


@dataclass
class ScenarioContext:
    """Execution context for a scenario run."""

    run_id: str
    scenario: Scenario
    bindings: Dict[str, str]

    # Execution state
    current_state: State = State.INIT
    start_time: float = field(default_factory=time.time)
    end_time: Optional[float] = None

    # Results tracking
    step_results: List[StepResult] = field(default_factory=list)
    final_state: Optional[State] = None

    # Remediation tracking
    fault_id: Optional[str] = None
    thread_id: Optional[str] = None
    interrupt_id: Optional[str] = None
    rca_response: Optional[str] = None
    remedy_response: Optional[str] = None

    def add_result(self, result: StepResult) -> None:
        """Add a step result to context."""
        self.step_results.append(result)
        self.current_state = result.state

    def set_final_state(self, state: State) -> None:
        """Set the final state and end time."""
        self.final_state = state
        self.end_time = time.time()

    @property
    def duration(self) -> float:
        """Get execution duration in seconds."""
        end = self.end_time or time.time()
        return end - self.start_time

    @property
    def passed(self) -> bool:
        """Check if scenario passed."""
        return self.final_state == State.PASS

    def to_dict(self) -> Dict[str, Any]:
        """Convert context to dictionary."""
        return {
            "run_id": self.run_id,
            "scenario_id": self.scenario.meta.id,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration": self.duration,
            "current_state": self.current_state.value,
            "final_state": self.final_state.value if self.final_state else None,
            "passed": self.passed,
            "step_results": [
                {
                    "state": r.state.value,
                    "success": r.success,
                    "message": r.message,
                    "score": r.score,
                    "timestamp": r.timestamp,
                }
                for r in self.step_results
            ],
        }


@dataclass
class ScenarioResult:
    """Final result of scenario execution."""

    run_id: str
    scenario_id: str
    passed: bool
    final_state: State
    duration: float
    step_results: List[StepResult]
    artifacts: Dict[str, str] = field(default_factory=dict)
    message: str = ""

    @classmethod
    def from_context(cls, context: ScenarioContext) -> "ScenarioResult":
        """Create result from context."""
        # Collect all artifacts
        artifacts = {}
        for step in context.step_results:
            artifacts.update(step.artifacts)

        return cls(
            run_id=context.run_id,
            scenario_id=context.scenario.meta.id,
            passed=context.passed,
            final_state=context.final_state or context.current_state,
            duration=context.duration,
            step_results=context.step_results,
            artifacts=artifacts,
            message=f"Scenario {'PASSED' if context.passed else 'FAILED'}",
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary."""
        return {
            "run_id": self.run_id,
            "scenario_id": self.scenario_id,
            "passed": self.passed,
            "final_state": self.final_state.value,
            "duration": self.duration,
            "message": self.message,
            "artifacts": self.artifacts,
            "step_results": [
                {
                    "state": r.state.value,
                    "success": r.success,
                    "message": r.message,
                    "score": r.score,
                    "timestamp": r.timestamp,
                }
                for r in self.step_results
            ],
        }