"""Pydantic models for the OpenEnv Customer Support environment."""

from __future__ import annotations

from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, Field


# ── Enumerations ──────────────────────────────────────────────────────

class IssueCategory(str, Enum):
    DELIVERY_DELAY = "delivery_delay"
    DAMAGED_PRODUCT = "damaged_product"
    REFUND_REQUEST = "refund_request"
    BILLING_ERROR = "billing_error"
    PRODUCT_INQUIRY = "product_inquiry"
    ACCOUNT_ISSUE = "account_issue"


class Sentiment(str, Enum):
    ANGRY = "angry"
    NEUTRAL = "neutral"
    HAPPY = "happy"


class ActionType(str, Enum):
    CLASSIFY = "classify"
    RESPOND = "respond"
    REFUND = "refund"
    ESCALATE = "escalate"
    RESOLVE = "resolve"


class TicketStatus(str, Enum):
    OPEN = "open"
    CLASSIFIED = "classified"
    RESPONDED = "responded"
    RESOLVED = "resolved"
    ESCALATED = "escalated"


# ── Ticket ────────────────────────────────────────────────────────────

class Ticket(BaseModel):
    """A customer support ticket."""
    ticket_id: str
    customer_name: str
    issue_text: str
    category: IssueCategory
    sentiment: Sentiment
    order_value: float = Field(ge=0.0)
    priority: str = "normal"  # low / normal / high / critical


# ── Observation ───────────────────────────────────────────────────────

class Observation(BaseModel):
    """Observation returned to the agent at each step."""
    ticket: Ticket
    status: TicketStatus = TicketStatus.OPEN
    conversation_history: List[str] = Field(default_factory=list)
    step_number: int = 0
    max_steps: int = 10
    classification: Optional[str] = None
    response_sent: Optional[str] = None
    action_taken: Optional[str] = None
    done: bool = False
    score: float = 0.0


# ── Action ────────────────────────────────────────────────────────────

class Action(BaseModel):
    """Action submitted by the agent."""
    action_type: ActionType
    classification: Optional[str] = None
    response_text: Optional[str] = None
    refund_amount: Optional[float] = None
    escalation_reason: Optional[str] = None


# ── Reward ────────────────────────────────────────────────────────────

class Reward(BaseModel):
    """Reward breakdown returned after each step."""
    total: float = Field(ge=0.0, le=1.0)
    classification_accuracy: float = Field(default=0.0, ge=0.0, le=1.0)
    response_quality: float = Field(default=0.0, ge=0.0, le=1.0)
    action_correctness: float = Field(default=0.0, ge=0.0, le=1.0)
    penalty_rude_tone: float = Field(default=0.0, ge=0.0, le=1.0)
    penalty_invalid_action: float = Field(default=0.0, ge=0.0, le=1.0)
    penalty_looping: float = Field(default=0.0, ge=0.0, le=1.0)


# ── Step Result ───────────────────────────────────────────────────────

class StepResult(BaseModel):
    """Result of a single environment step."""
    observation: Observation
    reward: Reward
    done: bool
    info: dict = Field(default_factory=dict)