from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Dict, Any

class Resume(BaseModel):
    candidate_id: str
    name: str
    email: str
    skills: List[str]
    experience_years: int
    education: Literal["High School", "Bachelor's", "Master's", "PhD"]
    previous_roles: List[str]
    # Protected attributes (for bias testing only; hidden in blind_mode)
    name_gender_proxy: Literal["M", "F", "N"]
    name_ethnicity_proxy: Literal["White", "Black", "Asian", "Hispanic", "Other"]
    graduation_year: Optional[int] = None

class JobDescription(BaseModel):
    job_id: str
    title: str
    required_skills: List[str]
    preferred_skills: List[str]
    min_experience: int
    max_experience: Optional[int] = None
    education_requirement: Literal["High School", "Bachelor's", "Master's", "PhD", "Any"]
    gender_coded_terms: List[str] = []  # Auto-detected (e.g., "ninja", "rockstar")

class Action(BaseModel):
    action_type: Literal["shortlist", "reject", "flag_bias", "request_clarification"]
    candidate_id: Optional[str] = None
    rank: Optional[int] = Field(None, ge=1, le=50)
    bias_reason: Optional[Literal["name_bias", "age_bias", "gender_coded_language", "education_elitism"]] = None
    clarification_field: Optional[str] = None

class Observation(BaseModel):
    current_resume: Optional[Resume] = None
    job_description: JobDescription
    skill_match_score: float = Field(ge=0.0, le=1.0)
    bias_risk_score: float = Field(ge=0.0, le=1.0)
    ml_fit_prob: Optional[float] = None
    detailed_analysis: Optional[Dict[str, Any]] = None
    shortlist_so_far: List[str]
    remaining_candidates: int
    step_count: int
    bias_metrics: Optional[Dict[str, float]] = None  # Populated after step()

class State(BaseModel):
    episode_id: str
    task_name: Literal["easy_shortlist", "medium_rank", "hard_fair_screen"]
    step_count: int
    total_candidates: int
    shortlist_complete: bool
    cumulative_reward: float
    bias_audit: Optional[Dict[str, float]] = None  # Final bias report