from __future__ import annotations

import json
from collections import defaultdict
from typing import Any

from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage


class ToolCallAccumulator:
    """将流式 tool call 分片累积为完整的 tool call。

    处理流式场景中常见的多分片 tool call 参数，
    按索引拼接部分 JSON 字符串。
    """

    def __init__(self) -> None:
        self.content: str = ""
        self.reasoning: str = ""
        self._tool_call_slots: dict[int, dict[str, Any]] = defaultdict(
            lambda: {"index": 0, "id": "", "name": "", "args": ""}
        )
        self.usage: TokenUsage | None = None

    def add_chunk(self, chunk: ChatChunk) -> None:
        """处理一个 ChatChunk，累积 content/reasoning/tool-calls/usage。"""
        if chunk.delta_content:
            self.content += chunk.delta_content
        if chunk.delta_reasoning:
            self.reasoning += chunk.delta_reasoning
        if chunk.usage is not None:
            self.usage = chunk.usage
        for tc in chunk.tool_call_chunks:
            idx = tc.get("index", 0)
            if idx not in self._tool_call_slots:
                self._tool_call_slots[idx] = {"index": idx, "id": "", "name": "", "args": ""}
            slot = self._tool_call_slots[idx]
            # The first chunk for a given index carries id/name; subsequent
            # arg-only chunks have them as None and should NOT overwrite.
            if tc.get("id"):
                slot["id"] = tc["id"]
            if tc.get("name"):
                slot["name"] = tc["name"]
            slot["args"] += tc.get("args", "")

    def drain(self) -> list[dict[str, Any]]:
        """返回所有已累积的 tool call 并重置内部槽位。"""
        result: list[dict[str, Any]] = []
        for idx in sorted(self._tool_call_slots.keys()):
            slot = self._tool_call_slots[idx]
            result.append(
                {
                    "index": slot["index"],
                    "id": slot["id"] or f"call_{idx}",
                    "name": slot["name"],
                    "args": slot["args"],
                }
            )
        self._tool_call_slots.clear()
        return result


class StreamHandler:
    """收集流式分片并生成最终的 ChatResult。"""

    def __init__(self) -> None:
        self._accumulator = ToolCallAccumulator()
        self._chunks: list[ChatChunk] = []

    def add_chunk(self, chunk: ChatChunk) -> None:
        """收集一个流式分片。"""
        self._chunks.append(chunk)
        self._accumulator.add_chunk(chunk)

    @property
    def chunks(self) -> list[ChatChunk]:
        """返回所有已收集分片的副本。"""
        return list(self._chunks)

    def to_result(self, model_name: str = "") -> ChatResult:
        """从所有累积分片构建最终的 ChatResult。"""
        raw_tool_calls = self._accumulator.drain()
        parsed_calls: list[dict[str, Any]] = []
        for tc in raw_tool_calls:
            args_raw = tc.get("args", "{}")
            try:
                parsed_args = json.loads(args_raw)
            except json.JSONDecodeError:
                parsed_args = {"_raw": args_raw}
            parsed_calls.append(
                {
                    "name": tc["name"],
                    "args": parsed_args,
                    "id": tc["id"],
                }
            )
        return ChatResult(
            content=self._accumulator.content,
            reasoning=self._accumulator.reasoning,
            tool_calls=parsed_calls,
            usage=self._accumulator.usage,
            model_name=model_name,
        )


class StreamFactory:
    """流处理器和累积器的工厂。"""

    @staticmethod
    def create_handler() -> StreamHandler:
        return StreamHandler()

    @staticmethod
    def create_accumulator() -> ToolCallAccumulator:
        return ToolCallAccumulator()