"""Data models for the SRE Bench environment."""

from typing import Any, Dict, List, Literal, Optional, Set

from pydantic import BaseModel, Field

try:
    from openenv.core.env_server.types import Action, Observation, State
except ImportError:
    from openenv.core.env_server.types import Action, Observation, State


ActionType = Literal["restart", "scale_up", "rollback", "investigate", "lookup_runbook", "noop"]


class SREAction(Action):
    """
    Action for the SRE Bench environment.

    action_type:
        - restart:         Restart a service (clears crash/memory_leak faults)
        - scale_up:        Scale up a service (clears latency_spike/high CPU+memory)
        - rollback:        Roll back a service config (clears config_corruption)
        - investigate:     Reveal the hidden fault type for a service (costs a step)
        - lookup_runbook:  Pull the SRE runbook for a service — returns symptom
                           checklist and recommended fix procedure (costs a step)
        - noop:            Do nothing
    """

    action_type: ActionType = Field(default="noop", description="Type of action")
    service: Optional[str] = Field(
        default=None,
        description="Target service name (required for all actions except noop)",
    )


class ServiceMetrics(BaseModel):
    """Public metrics for a single service."""

    status: str = Field(default="healthy", description="healthy | degraded | down")
    cpu: float = Field(default=20.0, description="CPU usage %")
    memory: float = Field(default=30.0, description="Memory usage %")
    error_rate: float = Field(default=0.0, description="Fraction of requests failing (0-1)")
    latency_ms: float = Field(default=50.0, description="P99 latency in ms")
    restart_count: int = Field(default=0, description="Total restarts this episode")


class SREObservation(Observation):
    """Observation returned at each step."""

    services: Dict[str, Dict[str, Any]] = Field(
        default_factory=dict,
        description=(
            "Public metrics per service: status, cpu, memory, error_rate, latency_ms, "
            "restart_count, plus trend signals: memory_trend, latency_trend, error_trend "
            "(rising | falling | stable)."
        ),
    )
    alerts: List[str] = Field(
        default_factory=list,
        description="Recent alert log (last 5 events, with severity prefix CRITICAL/WARNING/INFO)",
    )
    step: int = Field(default=0, description="Current step number")
    max_steps: int = Field(default=50, description="Episode length")
    uptime_score: float = Field(
        default=1.0,
        description="SLA-weighted uptime score so far (0-1). Final grader score at done=True.",
    )
    investigate_result: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Populated when action_type='investigate'",
    )
    runbook_result: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Populated when action_type='lookup_runbook': symptom checklist + fix procedure",
    )
    task_id: str = Field(default="", description="Active task identifier")
    task_description: str = Field(default="", description="Human-readable task objective")
    incident_severity: Optional[str] = Field(
        default=None,
        description="P1/P2/P3/OK — severity based on which SLA-critical services are down",
    )
    sla_breach_in: Optional[int] = Field(
        default=None,
        description="For incident_response task: steps remaining before SLA breach penalty kicks in",
    )


class SREState(State):
    """Full server-side state (includes hidden fault info)."""

    services: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
    faults: Dict[str, Optional[str]] = Field(
        default_factory=dict,
        description="Hidden fault per service: crash | memory_leak | latency_spike | config_corruption | None",
    )
    alerts: List[str] = Field(default_factory=list)
    step_count: int = Field(default=0)
    max_steps: int = Field(default=50)
    task_id: str = Field(default="")
    total_reward: float = Field(default=0.0)
    # SLA-weighted cumulative health (replaces simple cumulative_healthy count)
    cumulative_weighted_health: float = Field(default=0.0)
    # For cascading_failure grader: tracks which root causes the agent healed
    root_causes_healed: List[str] = Field(default_factory=list)
    # For single_fault grader: first step when all services were healthy
    steps_to_full_health: Optional[int] = Field(default=None)
    # Metric history for trend computation: {svc: [{"memory": f, "latency_ms": f, "error_rate": f}, ...]}
    metric_history: Dict[str, List[Dict[str, Any]]] = Field(default_factory=dict)
    # For incident_response grader: step at which api was restored
    api_restored_at: Optional[int] = Field(default=None)


__all__ = ["SREAction", "SREObservation", "SREState", "ServiceMetrics", "ActionType"]