"""Shared Pydantic models used by the workshop code and notebook."""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field

##### Shared label types #####

IntentLabel = Literal[
    "refund_request",
    "order_status",
    "billing_issue",
    "account_access",
    "damaged_item",
    "subscription_cancellation",
    "unknown",
]

RouteLabel = Literal[
    "general",
    "logistics",
    "finance",
    "account_support",
    "retention",
    "escalations",
]


##### Ticket input models #####

class TicketMetadata(BaseModel):
    """Capture ticket-related IDs passed into the system."""

    ticket_id: str
    customer_id: str | None = None
    account_id: str | None = None
    order_id: str | None = None


class TicketInput(BaseModel):
    """Represent one incoming support ticket."""

    input: str
    metadata: TicketMetadata


##### Evaluation dataset models #####

class ExpectedOutcome(BaseModel):
    """Store the expected structured labels for an evaluation example."""

    intent: IntentLabel
    route: RouteLabel
    should_escalate: bool


class EvalExample(BaseModel):
    """Represent one row in the JSONL evaluation dataset."""

    id: str
    input: str
    metadata: TicketMetadata
    expected: ExpectedOutcome
    reference_response: str


##### Tool and workflow intermediate models #####

class ToolResult(BaseModel):
    """Normalize tool responses into a consistent structured payload."""

    tool_name: str
    ok: bool = True
    payload: dict[str, Any]


class ClassificationResult(BaseModel):
    """Store the workflow classifier output."""

    intent: IntentLabel
    rationale: str = Field(default="")


class DraftResponseResult(BaseModel):
    """Store the drafted customer-facing response."""

    response: str


class WorkflowState(BaseModel):
    """Track intermediate workflow state across fixed steps."""

    ticket_text: str
    metadata: TicketMetadata
    intent: IntentLabel | None = None
    route: RouteLabel | None = None
    should_escalate: bool | None = None
    context: dict[str, Any] = Field(default_factory=dict)
    response: str | None = None
    used_tools: list[str] = Field(default_factory=list)


##### Final system output models #####

class SystemOutput(BaseModel):
    """Represent the final structured system answer shown in the workshop."""

    ticket_id: str
    intent: IntentLabel
    route: RouteLabel
    should_escalate: bool
    response: str
    used_tools: list[str] = Field(default_factory=list)
    tool_call_count: int = 0


##### Evaluation result models #####

class EvalResult(BaseModel):
    """Store one system prediction plus its evaluation scores."""

    example_id: str
    system_name: str
    input_text: str
    prediction: SystemOutput
    expected: ExpectedOutcome
    reference_response: str
    scores: dict[str, Any]


##### Judge output models #####

class JudgeResult(BaseModel):
    """Store the optional LLM-as-a-judge grading output."""

    helpfulness_score: int
    correctness_score: int
    policy_alignment_score: int
    overall_pass: bool
    rationale: str
