"""LLM Protocol - 数据类与抽象基类"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any


@dataclass
class TokenUsage:
    """Token 使用量"""

    prompt_tokens: int = 0
    completion_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens


@dataclass
class ToolCall:
    """工具调用"""

    id: str
    name: str
    arguments: dict[str, Any]


@dataclass
class LLMRequest:
    """LLM 请求"""

    messages: list[dict[str, str]]
    model: str
    tools: list[dict[str, Any]] | None = None
    tool_choice: str = "auto"
    temperature: float = 0.7
    max_tokens: int = 2000

    def __init__(
        self,
        messages: list[dict[str, str]],
        model: str,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str = "auto",
        temperature: float = 0.7,
        max_tokens: int = 2000,
        **kwargs: Any,
    ):
        self.messages = messages
        self.model = model
        self.tools = tools
        self.tool_choice = tool_choice
        self.temperature = temperature
        self.max_tokens = max_tokens
        self._extra = kwargs


@dataclass
class StreamChunk:
    """LLM 流式响应块"""

    content: str  # Delta content
    model: str
    tool_calls: list[ToolCall] = field(default_factory=list)  # Accumulated tool calls (only in final chunk)
    usage: TokenUsage | None = None  # Only in final chunk
    is_final: bool = False  # True for the last chunk


@dataclass
class LLMResponse:
    """LLM 响应"""

    content: str
    model: str
    usage: TokenUsage
    tool_calls: list[ToolCall] = field(default_factory=list)
    latency_ms: float = 0.0

    @property
    def has_tool_calls(self) -> bool:
        return len(self.tool_calls) > 0


class LLMProvider(ABC):
    """LLM Provider 抽象基类"""

    @abstractmethod
    async def chat(self, request: LLMRequest) -> LLMResponse:
        """发送 chat 请求并返回响应"""
        ...

    async def chat_stream(self, request: LLMRequest):
        """Stream chat response. Override in subclasses that support streaming.

        Yields StreamChunk objects. Default implementation falls back to
        non-streaming chat and yields a single chunk.
        """
        response = await self.chat(request)
        yield StreamChunk(
            content=response.content,
            model=response.model,
            tool_calls=response.tool_calls,
            usage=response.usage,
            is_final=True,
        )
