"""
ToolCallTracker - 工具调用追踪器
"""

import json
from dataclasses import dataclass, field
from typing import Dict, Optional


@dataclass
class ToolCallInfo:
    """工具调用信息"""
    id: str
    name: str
    args: Dict = field(default_factory=dict)
    emitted: bool = False
    args_complete: bool = False
    _json_buffer: str = ""


class ToolCallTracker:
    """工具调用追踪器"""

    def __init__(self):
        self._calls: Dict[str, ToolCallInfo] = {}
        self._last_tool_id: Optional[str] = None

    def update(
        self,
        tool_id: str,
        name: Optional[str] = None,
        args: Optional[Dict] = None,
        args_complete: bool = False,
    ) -> None:
        """更新工具调用信息"""
        if tool_id not in self._calls:
            self._calls[tool_id] = ToolCallInfo(
                id=tool_id,
                name=name or "",
                args=args or {},
                args_complete=args_complete,
            )
            self._last_tool_id = tool_id
        else:
            info = self._calls[tool_id]
            if name:
                info.name = name
            if args:
                info.args = args
            if args_complete:
                info.args_complete = True

    def append_json_delta(self, partial_json: str, index: int = 0) -> None:
        """累积 input_json_delta 片段"""
        tool_id = self._last_tool_id
        if tool_id and tool_id in self._calls:
            self._calls[tool_id]._json_buffer += partial_json

    def finalize_all(self) -> None:
        """最终化所有工具调用"""
        for info in self._calls.values():
            if info._json_buffer:
                try:
                    info.args = json.loads(info._json_buffer)
                except json.JSONDecodeError:
                    pass
                info._json_buffer = ""
            info.args_complete = True

    def is_ready(self, tool_id: str) -> bool:
        """检查工具调用是否准备好发送"""
        if tool_id not in self._calls:
            return False
        info = self._calls[tool_id]
        return bool(info.name) and not info.emitted

    def get_all(self) -> list[ToolCallInfo]:
        """获取所有工具调用"""
        return list(self._calls.values())

    def mark_emitted(self, tool_id: str) -> None:
        """标记已发送"""
        if tool_id in self._calls:
            self._calls[tool_id].emitted = True

    def get(self, tool_id: str) -> Optional[ToolCallInfo]:
        """获取工具调用信息"""
        return self._calls.get(tool_id)

    def get_pending(self) -> list[ToolCallInfo]:
        """获取所有未发送的工具调用"""
        return [info for info in self._calls.values() if not info.emitted]

    def emit_all_pending(self) -> list[ToolCallInfo]:
        """发送所有待处理的工具调用并标记"""
        pending = self.get_pending()
        for info in pending:
            info.emitted = True
        return pending

    def clear(self) -> None:
        """清空追踪器"""
        self._calls.clear()
