# https://docs.anthropic.com/en/api/messages
import time
from typing import Literal, Optional, List, Dict, Any, Union

from pydantic import BaseModel, Field, model_serializer

class ErrorDetail(BaseModel):
    type: str
    message: str

class ErrorResponse(BaseModel):
    error: ErrorDetail

def now_time():
    return time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())

class ModelItem(BaseModel):
    id: str
    type: Literal['model']
    display_name: str
    created_at: str = Field(default_factory=now_time)

class ModelResponse(BaseModel):
    data: List[ModelItem] = []
    first_id: Optional[str] = None
    last_id: Optional[str] = None
    has_more: bool

    @model_serializer
    def _serialize(self):
        return { k: v for k, v in self if v is not None }

class RequestCacheControl(BaseModel):
    type: Literal['ephemeral']
    ttl: Optional[Literal['5m', '1h']] = '5m'

class RequestMessageContentBlockText(BaseModel):
    text: str
    type: str = 'text'
    cache_control: Optional[RequestCacheControl] = None
    citations: Optional[List[Dict[str, Any]]] = None

class RequestMessageContentBlockImage(BaseModel):
    type: Literal['image']
    source: Dict[str, Any]
    cache_control: Optional[RequestCacheControl] = None

class RequestMessageContentBlockDocument(BaseModel):
    type: Literal['document']
    source: Dict[str, Any]
    cache_control: Optional[RequestCacheControl] = None
    citations: Optional[List[Dict[str, Any]]] = None
    context: Optional[str] = None
    title: Optional[str] = None

RequestMessageContentBlock = Union[
    RequestMessageContentBlockText,
    RequestMessageContentBlockImage,
    RequestMessageContentBlockDocument,
    # TODO: ... add more types
    Dict[str, Any]
]

class RequestMessageItem(BaseModel):
    role: Literal['user', 'assistant']
    content: Union[str, List[RequestMessageContentBlock]]

class McpServerToolConfiguration:
    allowed_tools: Optional[List[str]] = None
    enabled: Optional[bool] = None

class McpServerItem(BaseModel):
    name: str
    type: Literal['url']
    url: str
    authorization_token: Optional[str] = None

RequestSystem = Union[str, List[RequestMessageContentBlockText]]

class RequestThinkingEnabled(BaseModel):
    type: Literal['enabled']
    budget_tokens: int # should more than 1024

class RequestThinkingDisabled(BaseModel):
    type: Literal['disabled']

RequestThinking = Union[
    RequestThinkingEnabled,
    RequestThinkingDisabled,
]

class RequestMetadata(BaseModel):
    user_id: Optional[str] = None

class RequestToolChoiceAuto(BaseModel):
    type: Literal['auto']
    disable_parallel_tool_use: Optional[bool] = None

class RequestToolChoiceAny(BaseModel):
    type: Literal['any']
    disable_parallel_tool_use: Optional[bool] = None

class RequestToolChoiceTool(BaseModel):
    type: Literal['tool']
    name: str
    disable_parallel_tool_use: Optional[bool] = None

class RequestToolChoiceNone(BaseModel):
    type: Literal['none']
    disable_parallel_tool_use: Optional[bool] = None

RequestToolChoice = Union[
    RequestToolChoiceAuto,
    RequestToolChoiceAny,
    RequestToolChoiceTool,
    RequestToolChoiceNone,
    Dict[str, Any],
]

class RequestToolCustomInputSchema(BaseModel):
    type: Literal['object']
    properties: Dict[str, Any]
    required: Optional[List[str]] = None

class RequestToolCustom(BaseModel):
    type: Optional[Literal['custom']] = None
    name: str
    input_schema: RequestToolCustomInputSchema
    description: Optional[str] = None
    cache_control: Optional[RequestCacheControl] = None

class RequestToolBash(BaseModel):
    type: Literal['bash_20241022', 'bash_20250124']
    name: Literal['bash']
    cache_control: Optional[RequestCacheControl] = None

class RequestToolCodeExecution(BaseModel):
    type: Literal['code_execution_20250522']
    name: Literal['code_execution']
    cache_control: Optional[RequestCacheControl] = None

class RequestToolComputer(BaseModel):
    type: Literal['computer_20241022', 'computer_20250124']
    name: Literal['computer']
    display_height_px: int
    display_width_px: int
    display_number: int
    cache_control: Optional[RequestCacheControl] = None

class RequestToolTextEditor(BaseModel):
    type: Literal['text_editor_20241022', 'text_editor_20250124', 'text_editor_20250429']
    name: Literal['str_replace_editor', 'str_replace_based_edit_tool']
    cache_control: Optional[RequestCacheControl] = None

class RequestToolWebSearchUserLocation(BaseModel):
    type: Literal['approximate']
    city: Optional[str] = None
    country: Optional[str] = None
    region: Optional[str] = None
    timezone: Optional[str] = None

class RequestToolWebSearch(BaseModel):
    type: Literal['web_search_20250305']
    name: Literal['web_search']
    allowed_domains: Optional[List[str]] = None
    blocked_domains: Optional[List[str]] = None
    max_uses: Optional[int] = None
    user_location: Optional[RequestToolWebSearchUserLocation] = None
    cache_control: Optional[RequestCacheControl] = None

RequestTool = Union[
    RequestToolCustom,
    RequestToolBash,
    RequestToolCodeExecution,
    RequestToolComputer,
    RequestToolTextEditor,
    RequestToolWebSearch,
    Dict[str, Any],
]

class MessagesRequest(BaseModel):
    model: str
    messages: List[RequestMessageItem]
    max_tokens: int = 256
    temperature: Optional[float] = 1.0
    container: Optional[str] = None
    mcp_servers: Optional[List[McpServerItem]] = None
    metadata: Optional[RequestMetadata] = None
    service_tier: Optional[Literal['auto', 'standard_only']] = None
    stop_sequences: Optional[List[str]] = None
    stream: Optional[bool] = False
    system: Optional[RequestSystem] = None
    thinking: Optional[RequestThinking] = None
    tool_choice: Optional[RequestToolChoice] = None
    tools: Optional[List[RequestTool]] = None
    top_k: Optional[int] = -1
    top_p: Optional[float] = 1.0

class ResponseMessageContentBlockText(BaseModel):
    type: Literal['text', 'text_delta']
    text: str

class ResponseMessageContentBlockThinking(BaseModel):
    type: Literal['thinking']
    signature: str
    thinking: str

class ResponseMessageContentBlockRedactedThinking(BaseModel):
    type: Literal['redacted_thinking']
    data: str

class ResponseMessageContentBlockToolUse(BaseModel):
    type: Literal['tool_use']
    id: str
    name: str
    input: str

class ResponseMessageContentBlockToolServerToolUse(BaseModel):
    type: Literal['server_tool_use']
    id: str
    name: str
    input: str

class ResponseMessageContentBlockToolWebSearchToolResult(BaseModel):
    type: Literal['web_search_tool_result']
    tool_use_id: str
    content: Dict[str, Any]

class ResponseMessageContentBlockToolCodeExecutionToolResult(BaseModel):
    type: Literal['code_execution_tool_result']
    tool_use_id: str
    content: Dict[str, Any]

class ResponseMessageContentBlockToolMcpToolUse(BaseModel):
    type: Literal['mcp_tool_use']
    id: str
    name: str
    input: str
    server_name: str

class ResponseMessageContentBlockToolMcpToolResult(BaseModel):
    type: Literal['mcp_tool_result']
    tool_use_id: str
    is_error: bool = False
    content: Union[str, List[Dict[str, Any]]]

class ResponseMessageContentBlockToolContainerUpload(BaseModel):
    type: Literal['container_upload']
    file_id: str


ResponseMessageContentBlock = Union[
    ResponseMessageContentBlockText,
    ResponseMessageContentBlockThinking,
    ResponseMessageContentBlockRedactedThinking,
    ResponseMessageContentBlockToolUse,
    ResponseMessageContentBlockToolServerToolUse,
    ResponseMessageContentBlockToolWebSearchToolResult,
    ResponseMessageContentBlockToolCodeExecutionToolResult,
    ResponseMessageContentBlockToolMcpToolUse,
    ResponseMessageContentBlockToolMcpToolResult,
    ResponseMessageContentBlockToolContainerUpload,
    Dict[str, Any],
]

class MessagesResponseUsage(BaseModel):
    cache_creation: Optional[Dict[str, Any]] = None
    cache_creation_input_tokens: Optional[int] = None
    cache_read_input_tokens: Optional[int] = None
    server_tool_use: Optional[Dict[str, Any]] = None
    service_tier: Optional[Literal['standard', 'priority', 'batch']] = None
    input_tokens: int
    output_tokens: int

    @model_serializer
    def _serialize(self):
        return { k: v for k, v in self if v is not None }

class MessagesResponseContainer(BaseModel):
    expires_at: str
    id: str

ResponseStopReason = Literal[
    'end_turn',
    'max_tokens',
    'stop_sequence',
    'tool_use',
    'pause_turn',
    'refusal',
]

class MessagesResponse(BaseModel):
    id: str
    type: Literal['message']
    role: Literal['assistant']
    content: List[ResponseMessageContentBlock]
    model: str
    stop_reason: Optional[ResponseStopReason]
    stop_sequence: Optional[str]
    usage: MessagesResponseUsage
    container: Optional[MessagesResponseContainer] = None

    @model_serializer
    def _serialize(self):
        return { k: v for k, v in self if v is not None or k == 'stop_sequence' }

class StreamMessagesResponseMessage(BaseModel):
    type: Literal['message']
    id: str
    role: Literal['assistant']
    content: List[ResponseMessageContentBlock]
    model: str
    stop_reason: Optional[ResponseStopReason] = None
    stop_sequence: Optional[str] = None
    usage: MessagesResponseUsage

# only one
class StreamMessagesResponseMessageStart(BaseModel):
    type: Literal['message_start']
    message: StreamMessagesResponseMessage

# only one for each index
class StreamMessagesResponseContentBlockStart(BaseModel):
    type: Literal['content_block_start']
    index: int
    content_block: ResponseMessageContentBlock

# zero or more for each index
class StreamMessagesResponsePing(BaseModel):
    type: Literal['ping']

# one or more for each index
class StreamMessagesResponseContentBlockDelta(BaseModel):
    type: Literal['content_block_delta']
    index: int
    content_block: ResponseMessageContentBlock

# only one for each index
class StreamMessagesResponseContentBlockStop(BaseModel):
    type: Literal['content_block_stop']
    index: int

class StreamMessagesResponseMessageDeltaDelta(BaseModel):
    stop_reason: Optional[ResponseStopReason] = None
    stop_sequence: Optional[str] = None

# only one (partial messages response)
class StreamMessagesResponseMessageDelta(BaseModel):
    type: Literal['message_delta']
    delta: Union[StreamMessagesResponseMessageDeltaDelta, Dict[str, Any]]
    usage: MessagesResponseUsage # Cumulative

# only one
class StreamMessagesResponseMessageStop(BaseModel):
    type: Literal['message_stop']

class CountTokensRequest(BaseModel):
    messages: List[RequestMessageItem]
    model: str
    mcp_servers: Optional[List[McpServerItem]] = None
    system: Optional[RequestSystem] = None
    thinking: Optional[RequestThinking] = None
    tool_choice: Optional[RequestToolChoice] = None
    tools: Optional[List[RequestTool]] = None

class CountTokensResponse(BaseModel):
    input_tokens: int

class CompleteRequest(BaseModel):
    prompt: str
    model: str
    max_tokens_to_sample: int = 256
    temperature: Optional[float] = 1.0
    stop_sequences: Optional[List[str]] = None
    top_p: Optional[float] = 1.0
    top_k: Optional[int] = -1
    metadata: Optional[RequestMetadata] = None
    stream: Optional[bool] = False

class CompleteResponse(BaseModel):
    id: str
    completion: str
    model: str
    stop_reason: Optional[Literal['stop_sequence', 'max_tokens']] = None
    type: Literal['completion']

# one or more
class StreamCompleteResponseCompletion(BaseModel):
    type: Literal['completion']
    completion: str
    stop_reason: Optional[Literal['stop_sequence', 'max_tokens']] = None
    model: str

# zero or more
class StreamMessagesResponsePing(BaseModel):
    type: Literal['ping']

# zero or more
class StreamCompleteResponseError(BaseModel):
    error: ErrorDetail
