# -*- coding: utf-8 -*-
"""
Agent API 相关的类型定义
"""

from datetime import datetime
from typing import Any, Dict, List, Optional, Union
from uuid import uuid4

try:
    from typing import Self
except ImportError:
    from typing_extensions import Self

from pydantic import BaseModel, Field, field_validator


class MessageType:
    MESSAGE = "message"
    FUNCTION_CALL = "function_call"
    FUNCTION_CALL_OUTPUT = "function_call_output"
    PLUGIN_CALL = "plugin_call"
    PLUGIN_CALL_OUTPUT = "plugin_call_output"
    COMPONENT_CALL = "component_call"
    COMPONENT_CALL_OUTPUT = "component_call_output"
    MCP_LIST_TOOLS = "mcp_list_tools"
    MCP_APPROVAL_REQUEST = "mcp_approval_request"
    MCP_TOOL_CALL = "mcp_call"
    MCP_APPROVAL_RESPONSE = "mcp_approval_response"
    REASONING = "reasoning"
    HEARTBEAT = "heartbeat"
    ERROR = "error"


class ContentType:
    TEXT = "text"
    DATA = "data"
    IMAGE = "image"
    AUDIO = "audio"
    FILE = "file"
    REFUSAL = "refusal"


class Role:
    ASSISTANT = "assistant"
    USER = "user"
    SYSTEM = "system"


class RunStatus:
    """
    Enum class for agent event message.
    """

    Created = "created"
    InProgress = "in_progress"
    Completed = "completed"
    Cancelled = "cancelled"
    Failed = "failed"
    Rejected = "rejected"
    Unknown = "unknown"
    Queued = "queued"
    Incomplete = "incomplete"


class FunctionParameters(BaseModel):
    type: str
    """The type of the parameters object. Must be `object`."""

    properties: Dict[str, Any]
    """The properties of the parameters object."""

    required: Optional[List[str]]
    """The names of the required properties."""


class FunctionTool(BaseModel):
    """
    Model class for prompt message tool.
    """

    name: str
    """The name of the function to be called. """

    description: str
    """A description of what the function does, used by the model to choose
    when and how to call the function.
    """

    parameters: Union[FunctionParameters, Dict[str, Any]]
    """The parameters the functions accepts, described as a JSON Schema object.

    """


class Tool(BaseModel):
    """
    Model class for assistant prompt message tool call.
    """

    type: Optional[str] = None
    """The type of the tool. Currently, only `function` is supported."""

    function: Optional[FunctionTool] = None
    """The function that the model called."""


class FunctionCall(BaseModel):
    """
    Model class for assistant prompt message tool call function.
    """

    call_id: Optional[str] = None
    """The ID of the tool call."""

    name: Optional[str] = None
    """The name of the function to call."""

    arguments: Optional[str] = None
    """The arguments to call the function with, as generated by the model in
    JSON format.

    Note that the model does not always generate valid JSON, and may
    hallucinate  parameters not defined by your function schema. Validate
    the arguments in your code before calling your function.
    """


class FunctionCallOutput(BaseModel):
    """
    Model class for assistant prompt message tool call function.
    """

    call_id: str
    """The ID of the tool call."""

    output: str
    """The result of the function."""


class Error(BaseModel):
    code: str
    """The error code of the message."""

    message: str
    """The error message of the message."""


class Event(BaseModel):
    sequence_number: Optional[int] = None
    """sequence number of event, starting from 0"""

    object: str
    """The identity of the content part."""

    status: Optional[str] = None
    """The status of the message. in_progress, completed, or incomplete"""

    error: Optional[Error] = None
    """response error for output"""

    def created(self) -> Self:
        """
        Set the message status to 'created'.
        """
        self.status = RunStatus.Created
        return self

    def in_progress(self) -> Self:
        """
        Set the message status to 'in_progress'.
        """
        self.status = RunStatus.InProgress
        return self

    def completed(self) -> Self:
        """
        Set the message status to 'completed'.
        """
        self.status = RunStatus.Completed
        return self

    def failed(self, error: Error) -> Self:
        """
        Set the message status to 'failed'.
        """
        self.status = RunStatus.Failed
        self.error = error
        return self

    def rejected(self) -> Self:
        """
        Set the message status to 'rejected'.
        """
        self.status = RunStatus.Rejected
        return self

    def canceled(self) -> Self:
        """
        Set the message status to 'canceled'.
        """
        self.status = RunStatus.Cancelled
        return self


class Content(Event):
    type: str
    """The type of the content part."""

    object: str = "content"
    """The identity of the content part."""

    index: Optional[int] = None
    """the content index in message's content list"""

    delta: Optional[bool] = False
    """Whether this content is a delta."""

    msg_id: Optional[str] = None
    """message unique id"""


class ImageContent(Content):
    type: str = ContentType.IMAGE
    """The type of the content part."""

    image_url: Optional[str] = None
    """The image URL details."""


class TextContent(Content):
    type: str = ContentType.TEXT
    """The type of the content part."""

    text: Optional[str] = None
    """The text content."""


class DataContent(Content):
    type: str = ContentType.DATA
    """The type of the content part."""

    data: Optional[Dict] = None
    """The data content."""


class AudioContent(Content):
    type: str = ContentType.AUDIO
    """The type of the content part."""

    data: Optional[str] = None
    """The audio data details."""

    format: Optional[str] = None
    """
    The format of the audio data.
    """


class FileContent(Content):
    type: str = ContentType.FILE
    """The type of the content part."""

    file_url: Optional[str] = None
    """The file URL details."""

    file_id: Optional[str] = None
    """The file ID details."""

    filename: Optional[str] = None
    """The file name details."""

    file_data: Optional[str] = None
    """The file data details."""


class RefusalContent(Content):
    type: str = ContentType.REFUSAL
    """The type of the content part."""

    refusal: Optional[str] = None
    """The refusal content."""


class Message(Event):
    id: str = Field(default_factory=lambda: "msg_" + str(uuid4()))
    """message unique id"""

    object: str = "message"
    """message identity"""

    type: str = "message"
    """The type of the message."""

    status: str = RunStatus.Created
    """The status of the message. in_progress, completed, or incomplete"""

    role: Optional[str] = None
    """The role of the messages author, should be in `user`,`system`,
    'assistant'."""

    content: Optional[
        List[
            Union[
                TextContent,
                ImageContent,
                DataContent,
                RefusalContent,
                FileContent,
                AudioContent,
            ]
        ]
    ] = None
    """The contents of the message."""

    code: Optional[str] = None
    """The error code of the message."""

    message: Optional[str] = None
    """The error message of the message."""

    metadata: Optional[Dict] = None
    """The metadata of the message. """


class BaseRequest(BaseModel):
    """agent request"""

    input: List[Message]
    """
    input messages
    """

    stream: bool = False
    """If set, partial message deltas will be sent, like in ChatGPT. """


class AgentRequest(BaseRequest):
    """agent request"""

    model: Optional[str] = None
    """
    model id
    """

    top_p: Optional[float] = None
    """Nucleus sampling, between (0, 1.0],  where the model considers the
    results of the tokens with top_p probability  mass.

    So 0.1 means only the tokens comprising the top 10% probability mass are
    considered.

    We generally recommend altering this or `temperature` but not both.
    """

    temperature: Optional[float] = None
    """What sampling temperature to use, between 0 and 2.

    Higher values like 0.8 will make the output more random, while lower values
    like 0.2 will make it more focused and deterministic.

    We generally recommend altering this or `top_p` but not both.
    """

    frequency_penalty: Optional[float] = None
    """Positive values penalize new tokens based on their existing frequency in
    the text so far, decreasing the model's likelihood to repeat the same line
    verbatim.

    """

    presence_penalty: Optional[float] = None
    """Number between -2.0 and 2.0.

    Positive values penalize new tokens based on whether they appear in the
    text so far, increasing the model's likelihood to talk about new topics.

    """

    max_tokens: Optional[int] = None
    """The maximum number of [tokens](/tokenizer) that can be generated in the
    chat completion.

    The total length of input tokens and generated tokens is limited by the
    model's context length.
    """

    stop: Optional[Union[Optional[str], List[str]]] = None
    """Up to 4 sequences where the API will stop generating further tokens."""

    n: Optional[int] = Field(default=1, ge=1, le=5)
    """How many chat completion choices to generate for each input message.

    Note that you will be charged based on the number of generated tokens
    across all of the choices. Keep `n` as `1` to minimize costs.
    """

    seed: Optional[int] = None
    """If specified, system will make a best effort to sample
    deterministically, such that repeated requests with the same `seed` and
    parameters should return the same result.
    """

    tools: Optional[List[Union[Tool, Dict]]] = None
    """
    tool call list
    """

    session_id: Optional[str] = None
    """conversation id for dialog"""

    previous_response_id: Optional[str] = None
    """previous response id for context linking"""

    background: Optional[bool] = None
    """whether to run in background"""

    instructions: Optional[str] = None
    """instructions for the agent"""

    tools: Optional[List[Union[Tool]]] = None
    """tools for the agent"""

    biz_params: Optional[Dict] = None
    """
    business parameters for workflow execution
    """


class BaseResponse(Event):
    id: Optional[str] = Field(
        default_factory=lambda: "response_" + str(uuid4()),
    )
    """response unique id"""

    @field_validator("id", mode="before")
    @classmethod
    def validate_id(cls, v):
        if v is None:
            return "response_" + str(uuid4())
        return v

    object: str = "response"
    """response identity"""

    status: str = RunStatus.Created
    """response run status"""

    created_at: int = int(datetime.now().timestamp())
    """request start time"""

    completed_at: Optional[int] = None
    """request completed time"""

    output: Optional[List[Message]] = None
    """response data for output"""

    usage: Optional[Dict] = None
    """response usage for output"""


class AgentResponse(BaseResponse):
    """agent response"""

    session_id: Optional[str] = None
    """conversation id for dialog"""

    model: Optional[str] = None
    """model id"""
