from pydantic import BaseModel, validator
from typing import List, Optional, Dict, Literal

# What agent sees
class Observation(BaseModel):
    code: str
    intent: str
    constraints: List[str]
    task_id: str
    difficulty: Literal["easy", "medium", "hard"]
    step_index: int

# What agent sends
class Action(BaseModel):
    action_type: Literal["detect", "risk", "fix"]
    detected_issues: Optional[List[str]] = None
    risk_level: Optional[Literal["low", "medium", "high"]] = None
    suggested_fix: Optional[str] = None

# Structured agent response for RL training
class AgentResponse(BaseModel):
    """Structured response from LLM agent with confidence and reasoning"""
    issues_detected: List[str] = []
    risk_level: Literal["low", "medium", "high"] = "low"
    fix: str = ""
    confidence: float = 0.5
    reasoning: str = ""
    reward_breakdown: Dict[str, float] = {}
    
    @validator('confidence')
    def clamp_confidence(cls, v):
        return max(0.0, min(1.0, v))

# What env returns
class Reward(BaseModel):
    value: float
    feedback: str
    partial_scores: Dict[str, float]
    breakdown: Dict[str, float] = {}  # Detailed reward breakdown per step
    
    @validator('value')
    def clamp_value(cls, v):
        return max(0.0, min(1.0, v))

# Agent state for RL learning
class AgentState(BaseModel):
    """State tracking for RL agent learning"""
    episode: int = 0
    total_reward: float = 0.0
    best_score: float = 0.0
    prompt_strategy: str = "default"
    past_failures: List[str] = []
    learning_rate: float = 0.1
    improvement_history: List[float] = []