from enum import Enum
from typing import Optional
from pydantic import BaseModel, field_validator
from openenv.core import Action, Observation, State


class ActionType(str, Enum):
    read_doc = "read_doc"
    flag_suspicious = "flag_suspicious"
    unflag_doc = "unflag_doc"
    submit_answer = "submit_answer"


class ContextCorruptionAction(Action):
    action_type: ActionType
    doc_id: Optional[int] = None
    answer: Optional[str] = None
    confidence: Optional[float] = None

    @field_validator("confidence")
    @classmethod
    def confidence_range(cls, v):
        if v is not None and not (0.0 <= v <= 1.0):
            raise ValueError("confidence must be between 0.0 and 1.0")
        return v


class Document(BaseModel):
    id: int
    title: str
    content: str
    is_flagged: bool = False


class EpisodeObservation(Observation):
    question: str = ""
    documents: list[Document] = []
    flagged_ids: list[int] = []
    budget_remaining: int = 0
    turn: int = 0
    message: Optional[str] = None
    # `done` and `reward` inherited from Observation


class ContextCorruptionState(State):
    question: str = ""
    ground_truth: str = ""
    corrupt_ids: list[int] = []
    flagged_ids: list[int] = []
    budget_used: int = 0
    done: bool = False
    reward: Optional[float] = None
    breakdown: Optional[dict] = None