"""
型別定義 - OpenAI Chat Completions API 格式

遷移自 Anthropic Messages API，改用 OpenAI Chat Completions 格式
透過 LiteLLM 路由支援 Claude / Gemini / OpenAI 等多模型
"""

import json
import re
import uuid
from dataclasses import dataclass, field
from typing import Any, Literal, Optional


# ============================================================================
# Anthropic 格式的訊息結構
# ============================================================================

@dataclass
class ToolUseBlock:
    """工具使用區塊（client-side tool）"""
    type: Literal["tool_use"] = "tool_use"
    id: str = ""
    name: str = ""
    input: dict = field(default_factory=dict)


@dataclass
class ServerToolUseBlock:
    """
    Server-side 工具使用區塊

    用於 Anthropic 伺服器執行的工具（如 web_search_20250305）
    這類工具不需要 client 執行，結果會在 content 中以 web_search_tool_result 形式回傳
    """
    type: Literal["server_tool_use"] = "server_tool_use"
    id: str = ""
    name: str = ""
    input: dict = field(default_factory=dict)


@dataclass
class WebSearchToolResultBlock:
    """
    Web Search 結果區塊

    Anthropic server-side web search 的搜尋結果
    """
    type: Literal["web_search_tool_result"] = "web_search_tool_result"
    tool_use_id: str = ""
    content: list[dict] = field(default_factory=list)  # 搜尋結果列表


@dataclass
class ToolResultBlock:
    """工具結果區塊"""
    type: Literal["tool_result"] = "tool_result"
    tool_use_id: str = ""
    content: str = ""
    is_error: bool = False


@dataclass
class TextBlock:
    """文字區塊"""
    type: Literal["text"] = "text"
    text: str = ""


# Content block 可以是文字、工具使用、工具結果或 server tool 相關
ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock | ServerToolUseBlock | WebSearchToolResultBlock


@dataclass
class Message:
    """訊息物件（Anthropic 格式）"""
    role: Literal["user", "assistant"]
    content: list[ContentBlock] | str = field(default_factory=list)

    def to_dict(self) -> dict:
        """轉換為 dict（用於 API 請求）"""
        d = {"role": self.role}

        if isinstance(self.content, str):
            d["content"] = self.content
        else:
            content_list = []
            for block in self.content:
                if isinstance(block, TextBlock):
                    content_list.append({
                        "type": "text",
                        "text": block.text
                    })
                elif isinstance(block, ToolUseBlock):
                    content_list.append({
                        "type": "tool_use",
                        "id": block.id,
                        "name": block.name,
                        "input": block.input
                    })
                elif isinstance(block, ToolResultBlock):
                    content_list.append({
                        "type": "tool_result",
                        "tool_use_id": block.tool_use_id,
                        "content": block.content,
                        "is_error": block.is_error
                    })
            d["content"] = content_list

        return d


@dataclass
class ToolDefinition:
    """工具定義（Anthropic 格式）"""
    name: str = ""
    description: str = ""
    input_schema: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "description": self.description,
            "input_schema": self.input_schema,
        }


# ============================================================================
# SDK 使用的訊息類型（對外介面）
# ============================================================================

@dataclass
class AgentMessage:
    """
    Agent 回應訊息

    可以直接 print() 或 str() 取得內容
    """
    type: Literal["text", "tool_use", "tool_result", "server_tool_use", "web_search_result", "error"] = "text"
    content: str = ""
    tool_name: Optional[str] = None
    tool_input: Optional[dict] = None
    tool_result: Optional[str] = None
    search_results: Optional[list[dict]] = None  # Web search 結果

    def __str__(self) -> str:
        if self.type == "text":
            return self.content
        elif self.type == "tool_use":
            return f"[Tool: {self.tool_name}]"
        elif self.type == "tool_result":
            result_preview = (self.tool_result or "")[:100]
            return f"[Result: {self.tool_name}] {result_preview}..."
        elif self.type == "server_tool_use":
            return f"[Server Tool: {self.tool_name}]"
        elif self.type == "web_search_result":
            count = len(self.search_results) if self.search_results else 0
            return f"[Web Search: {count} results]"
        elif self.type == "error":
            return f"[Error] {self.content}"
        return self.content


# ============================================================================
# API 回應結構（OpenAI Chat Completions API）
# ============================================================================

@dataclass
class Usage:
    """Token 使用量（OpenAI 格式，保留 Anthropic 相容別名）"""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

    @property
    def input_tokens(self) -> int:
        """Anthropic 相容別名"""
        return self.prompt_tokens

    @property
    def output_tokens(self) -> int:
        """Anthropic 相容別名"""
        return self.completion_tokens


@dataclass
class FunctionCall:
    """OpenAI function call"""
    name: str = ""
    arguments: str = ""  # JSON 字串

    def get_parsed_arguments(self) -> dict:
        """解析 arguments JSON 字串為 dict"""
        if not self.arguments:
            return {}
        return json.loads(self.arguments)


@dataclass
class ToolCall:
    """OpenAI tool call"""
    id: str = ""
    type: str = "function"
    function: FunctionCall = field(default_factory=FunctionCall)


@dataclass
class ChatCompletionChoice:
    """OpenAI choice"""
    index: int = 0
    message: dict = field(default_factory=dict)
    finish_reason: Optional[str] = None


# Harmony 格式 tool call 的正則（GPT-OSS-120B 的非標準輸出）
_HARMONY_PATTERN = re.compile(
    r'<\|start\|>assistant<\|channel\|>.+? to=functions\.(\w+).*?<\|message\|>(.*?)<\|call\|>',
    re.DOTALL,
)


def _parse_harmony_tool_calls(content: str) -> list[dict]:
    """
    解析 Harmony 格式的 tool call（GPT-OSS-120B 的非標準輸出）

    Harmony 格式範例：
    <|start|>assistant<|channel|>commentary to=functions.Glob<|message|>{"directory": "src", "pattern": "**/*.c"}<|call|>

    回傳標準 OpenAI tool_calls 格式的 list[dict]
    """
    matches = _HARMONY_PATTERN.findall(content)
    if not matches:
        return []

    tool_calls = []
    for tool_name, arguments in matches:
        tool_calls.append({
            "id": f"harmony_{uuid.uuid4().hex[:8]}",
            "type": "function",
            "function": {
                "name": tool_name,
                "arguments": arguments.strip(),
            },
        })
    return tool_calls


@dataclass
class ChatCompletionResponse:
    """OpenAI Chat Completion 回應"""
    id: str = ""
    object: str = "chat.completion"
    created: int = 0
    model: str = ""
    choices: list[ChatCompletionChoice] = field(default_factory=list)
    usage: Optional[Usage] = None

    @classmethod
    def from_dict(cls, data: dict) -> "ChatCompletionResponse":
        """從 API 回應 dict 解析"""
        usage_data = data.get("usage")
        usage = None
        if usage_data:
            usage = Usage(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            )

        choices = []
        for c in data.get("choices", []):
            choices.append(ChatCompletionChoice(
                index=c.get("index", 0),
                message=c.get("message", {}),
                finish_reason=c.get("finish_reason"),
            ))

        return cls(
            id=data.get("id", ""),
            object=data.get("object", "chat.completion"),
            created=data.get("created", 0),
            model=data.get("model", ""),
            choices=choices,
            usage=usage,
        )

    def get_text(self) -> str:
        """取得回覆文字（自動清理 Harmony 格式的 tool call raw text）"""
        if self.choices:
            content = self.choices[0].message.get("content") or ""
            if content and _parse_harmony_tool_calls(content):
                # 移除 Harmony 格式部分，保留其餘文字
                cleaned = _HARMONY_PATTERN.sub('', content).strip()
                return cleaned
            return content
        return ""

    def get_tool_calls(self) -> list[ToolCall]:
        """取得工具呼叫（支援標準 tool_calls 和 Harmony 格式 fallback）"""
        if not self.choices:
            return []

        # 優先使用標準 tool_calls 欄位
        raw_calls = self.choices[0].message.get("tool_calls") or []

        # Fallback: 如果 tool_calls 為空，檢查 content 是否含 Harmony 格式
        if not raw_calls:
            content = self.choices[0].message.get("content") or ""
            raw_calls = _parse_harmony_tool_calls(content)

        return [ToolCall(
            id=c["id"],
            type=c.get("type", "function"),
            function=FunctionCall(
                name=c["function"]["name"],
                arguments=c["function"]["arguments"],
            ),
        ) for c in raw_calls]

    def get_local_tool_calls(self) -> list[ToolCall]:
        """取得需要本地執行的工具呼叫（過濾掉 server-side 的 web_search）"""
        return [tc for tc in self.get_tool_calls() if tc.function.name != "web_search"]

    def has_tool_calls(self) -> bool:
        """是否有需要本地執行的工具呼叫"""
        return len(self.get_local_tool_calls()) > 0

    def to_assistant_message_dict(self) -> dict:
        """轉為可加入 messages 的 assistant dict"""
        if not self.choices:
            return {"role": "assistant", "content": ""}

        msg = {**self.choices[0].message}

        # 如果有解析出 Harmony tool calls，用標準格式替換
        if not msg.get("tool_calls"):
            content = msg.get("content") or ""
            harmony_calls = _parse_harmony_tool_calls(content)
            if harmony_calls:
                msg["tool_calls"] = harmony_calls
                # 清理 content 中的 Harmony raw text
                msg["content"] = _HARMONY_PATTERN.sub('', content).strip() or None

        return {"role": "assistant", **msg}

    @property
    def finish_reason(self) -> Optional[str]:
        """取得 finish_reason"""
        if self.choices:
            return self.choices[0].finish_reason
        return None


# ============================================================================
# 相容類型（用於 chat.py 等消費端）
# ============================================================================

@dataclass
class AssistantMessage:
    """
    與 claude_agent_sdk.types.AssistantMessage 相容的類型

    用於 ClaudeSDKClient.receive_response() 回傳的訊息
    """
    content: list[ContentBlock] = field(default_factory=list)
    role: Literal["assistant"] = "assistant"

    @classmethod
    def from_response(cls, response: ChatCompletionResponse) -> "AssistantMessage":
        """從 ChatCompletionResponse 建立 AssistantMessage

        將 OpenAI 格式轉為內部 ContentBlock 格式，保持消費端介面不變
        """
        content_blocks: list[ContentBlock] = []

        # 文字內容
        text = response.get_text()
        if text:
            content_blocks.append(TextBlock(type="text", text=text))

        # tool_calls 轉為 ToolUseBlock（保持消費端相容）
        for tc in response.get_local_tool_calls():
            content_blocks.append(ToolUseBlock(
                type="tool_use",
                id=tc.id,
                name=tc.function.name,
                input=tc.function.get_parsed_arguments(),
            ))

        return cls(content=content_blocks)


@dataclass
class UserMessage:
    """
    與 claude_agent_sdk.types.UserMessage 相容的類型
    """
    content: list[ContentBlock] | str = ""
    role: Literal["user"] = "user"


@dataclass
class SystemMessage:
    """
    與 claude_agent_sdk.types.SystemMessage 相容的類型
    """
    content: str = ""
    role: Literal["system"] = "system"


@dataclass
class ResultMessage:
    """
    與 claude_agent_sdk.types.ResultMessage 相容的類型

    包含完成訊息和使用量資訊
    """
    content: str = ""
    usage: Optional[Usage] = None
    stop_reason: Optional[str] = None
