"""Typed models for the SupportDesk OpenEnv environment."""

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field

from supportdesk_env.openenv_compat import Action, Observation, State


class KnowledgeSnippet(BaseModel):
    """A policy or runbook excerpt the agent can use during triage."""

    article_id: str
    title: str
    content: str


class SupportTicket(BaseModel):
    """Static task input representing the inbound support ticket."""

    customer_name: str
    customer_tier: Literal["free", "pro", "enterprise"]
    company: str
    subject: str
    body: str
    region: str
    affected_users: int | None = None
    sla_minutes_remaining: int | None = None
    business_impact: str | None = None
    secondary_concerns: list[str] = Field(default_factory=list)
    attachments: list[str] = Field(default_factory=list)


class ActionHistoryEntry(BaseModel):
    """A concise trace entry used in observations and state dumps."""

    step: int
    operation: str
    summary: str
    reward_delta: float = 0.0


class SupportCaseProgress(BaseModel):
    """Mutable case state that graders score against."""

    queue: str | None = None
    priority: str | None = None
    issue_type: str | None = None
    status: str = "new"
    resolution_code: str | None = None
    requested_fields: list[str] = Field(default_factory=list)
    reply: str | None = None
    internal_note: str | None = None


class SupportDeskAction(Action):
    """One structured action the agent can take at each step."""

    operation: Literal["classify", "request_info", "draft_reply", "add_internal_note", "submit"]
    queue: str | None = None
    priority: str | None = None
    issue_type: str | None = None
    status: str | None = None
    resolution_code: str | None = None
    requested_fields: list[str] = Field(default_factory=list)
    reply: str | None = None
    internal_note: str | None = None


class SupportDeskObservation(Observation):
    """Observation emitted to the agent after reset and each step."""

    task_id: str
    difficulty: Literal["easy", "medium", "hard"]
    objective: str
    ticket: SupportTicket
    knowledge_base: list[KnowledgeSnippet]
    available_queues: list[str]
    available_priorities: list[str]
    available_statuses: list[str]
    available_issue_types: list[str]
    case: SupportCaseProgress
    action_history: list[ActionHistoryEntry] = Field(default_factory=list)
    feedback: str = ""
    remaining_steps: int = 0


class SupportDeskState(State):
    """Current environment state returned by the OpenEnv state() API."""

    episode_id: str | None = None
    task_id: str
    difficulty: Literal["easy", "medium", "hard"]
    step_count: int = 0
    reward: float = 0.0
    done: bool = False
    current_score: float = 0.0
    max_steps: int = 0
    case: SupportCaseProgress
    action_history: list[ActionHistoryEntry] = Field(default_factory=list)
    completed_milestones: list[str] = Field(default_factory=list)
    last_feedback: str = ""