# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Customer Support Ticket Management Environment.

This environment simulates real-world customer support operations where an AI agent
must handle incoming tickets by categorizing, prioritizing, routing, and responding.
"""

from typing import Optional, List, Literal
from openenv.core.env_server.types import Action, Observation
from pydantic import Field


class CustomerHistory(Observation):
    """Customer historical data embedded in observations"""
    account_age_days: int = Field(description="Number of days since account creation")
    total_tickets: int = Field(description="Total number of previous tickets")
    resolved_tickets: int = Field(description="Number of successfully resolved tickets")
    satisfaction_score: float = Field(ge=0.0, le=5.0, description="Average satisfaction rating (0-5)")
    is_premium: bool = Field(description="Whether customer has premium status")
    lifetime_value: float = Field(description="Customer lifetime value in USD")


class TicketMetadata(Observation):
    """Metadata about the current ticket"""
    ticket_id: str = Field(description="Unique ticket identifier")
    timestamp: str = Field(description="Ticket creation timestamp")
    customer_id: str = Field(description="Customer identifier")
    channel: Literal["email", "chat", "phone", "social"] = Field(description="Communication channel")


class CustomerSupportObservation(Observation):
    """
    Observation returned by the environment.
    Contains all information visible to the agent about the current ticket.
    """
    # Ticket information
    ticket_id: str = Field(description="Unique ticket identifier")
    timestamp: str = Field(description="Ticket creation timestamp")
    customer_id: str = Field(description="Customer identifier")
    channel: Literal["email", "chat", "phone", "social"] = Field(description="Communication channel")

    # Customer message
    customer_message: str = Field(description="The customer's support request message")

    # Customer history
    account_age_days: int = Field(description="Number of days since account creation")
    total_tickets: int = Field(description="Total number of previous tickets")
    resolved_tickets: int = Field(description="Number of successfully resolved tickets")
    satisfaction_score: float = Field(ge=0.0, le=5.0, description="Average satisfaction rating (0-5)")
    is_premium: bool = Field(description="Whether customer has premium status")
    lifetime_value: float = Field(description="Customer lifetime value in USD")

    # Additional context
    previous_interactions: List[str] = Field(
        default_factory=list,
        description="List of previous messages in this ticket thread"
    )
    attachments: List[str] = Field(
        default_factory=list,
        description="List of attachment filenames (if any)"
    )

    # Task context
    task_id: str = Field(default="easy", description="Current task difficulty level")


class CustomerSupportAction(Action):
    """
    Action taken by the agent.
    Represents the agent's decision about how to handle the ticket.
    """
    category: Literal["billing", "technical", "account", "shipping", "general"] = Field(
        description="Ticket category classification"
    )
    priority: Literal["low", "medium", "high", "critical"] = Field(
        description="Assigned priority level"
    )
    assigned_team: Literal["tier1", "tier2", "billing", "technical", "management"] = Field(
        description="Team to handle the ticket"
    )
    response_draft: str = Field(
        min_length=10,
        description="Draft response to the customer (minimum 10 characters)"
    )
    internal_notes: Optional[str] = Field(
        default=None,
        description="Internal notes for the assigned team"
    )
    escalate: bool = Field(
        default=False,
        description="Whether to escalate to management"
    )