from __future__ import annotations
from pydantic import BaseModel, Field
from typing import Literal, Optional, List, Dict, Any
from enum import Enum


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class FSMState(str, Enum):
    IDLE          = "IDLE"
    INVESTIGATING = "INVESTIGATING"
    REMEDIATING   = "REMEDIATING"
    VERIFYING     = "VERIFYING"
    RESOLVED      = "RESOLVED"
    FAILED        = "FAILED"


class Difficulty(str, Enum):
    EASY   = "easy"
    MEDIUM = "medium"
    HARD   = "hard"


# ---------------------------------------------------------------------------
# Action models
# ---------------------------------------------------------------------------

class DiagnosticQueryAction(BaseModel):
    action_type:         Literal["diagnostic_query"]
    metric_id:           str
    service:             str
    time_window_minutes: int = 10


class LogInspectionAction(BaseModel):
    action_type:  Literal["log_inspection"]
    service:      str
    tail_lines:   int = 50
    grep_pattern: Optional[str] = None


class RemediationAction(BaseModel):
    action_type:    Literal["remediation"]
    operation_type: Literal["restart", "rollback", "scale_up", "kill_pid", "update_config"]
    target_service: str
    parameters:     Dict[str, Any] = Field(default_factory=dict)


class SubmitResolutionAction(BaseModel):
    action_type:        Literal["submit_resolution"]
    root_cause_service: str
    explanation:        str


# Single union type the agent always submits
Action = (
    DiagnosticQueryAction
    | LogInspectionAction
    | RemediationAction
    | SubmitResolutionAction
)


# ---------------------------------------------------------------------------
# Observation model
# ---------------------------------------------------------------------------

class GoldenSignals(BaseModel):
    latency_p95_ms: float = 0.0
    error_rate:     float = 0.0
    traffic_rps:    float = 0.0
    saturation_pct: float = 0.0


class Observation(BaseModel):
    command_stdout:   str            = ""
    command_stderr:   str            = ""
    exit_code:        int            = 0
    active_alerts:    List[str]      = Field(default_factory=list)
    golden_signals:   GoldenSignals  = Field(default_factory=GoldenSignals)
    rolling_summary:  str            = ""
    system_context:   str            = ""        # injected by memory layer
    system_timestamp: float          = 0.0


# ---------------------------------------------------------------------------
# Reward model
# ---------------------------------------------------------------------------

class Reward(BaseModel):
    value:           float                   # bounded [-0.5, 1.5]
    health_delta:    float = 0.0             # alpha * delta_H
    milestone_bonus: float = 0.0             # beta * M
    efficiency:      float = 0.0             # lambda * E
    penalty:         float = 0.0             # gamma * P
    timestep_cost:   float = 0.01            # delta (constant)
    breakdown:       Dict[str, float] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# Episode state
# ---------------------------------------------------------------------------

class EpisodeState(BaseModel):
    task_id:          str
    fsm_state:        FSMState        = FSMState.IDLE
    step:             int             = 0
    max_steps:        int             = 15
    done:             bool            = False
    total_reward:     float           = 0.0
    health_score:     float           = 0.0
    active_alerts:    List[str]       = Field(default_factory=list)
    golden_signals:   GoldenSignals   = Field(default_factory=GoldenSignals)
    memory_backend:   Literal["pgvector", "disabled"] = "pgvector"
    info:             Dict[str, Any]  = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# API request / response wrappers
# ---------------------------------------------------------------------------

class StepRequest(BaseModel):
    action: Dict[str, Any]


class StepResponse(BaseModel):
    observation: Observation
    reward:      Reward
    done:        bool
    info:        Dict[str, Any] = Field(default_factory=dict)


class ResetRequest(BaseModel):
    task_id:        str = "task_1"
    memory_backend: Literal["pgvector", "disabled"] = "pgvector"
    seed:           int = 42


class ResetResponse(BaseModel):
    observation: Observation
    task_id:     str
    max_steps:   int


# ---------------------------------------------------------------------------
# Memory models
# ---------------------------------------------------------------------------

class EpisodeMemory(BaseModel):
    episode_id:   str
    task_id:      str
    task_success: bool
    total_reward: float
    steps_taken:  int                    # tracked for ablation metric
    actions:      List[Dict[str, Any]]   = Field(default_factory=list)
    state_emb:    List[float]            = Field(default_factory=list)
    summary:      str                    = ""


# ---------------------------------------------------------------------------
# Ablation / metrics models
# ---------------------------------------------------------------------------

class AblationRecord(BaseModel):
    epoch:                    int
    task_id:                  str
    memory_backend:           Literal["pgvector", "disabled"]
    mean_reward:              float
    mean_steps_to_resolution: float      # point 7 — steps-to-resolution metric
    task_success_rate:        float


# ---------------------------------------------------------------------------
# Adversarial log test model  (point 8)
# ---------------------------------------------------------------------------

class QuarantineResult(BaseModel):
    raw_log:          str
    sanitised_log:    str
    injection_detected: bool
    truncated:        bool