"""Conversation-scoped bridge for tools implemented by the connected client."""

from __future__ import annotations

import asyncio
from typing import Any, Literal, Protocol

from loguru import logger
from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    InputTransportMessageFrame,
    OutputTransportMessageUrgentFrame,
    StopFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor


class ClientToolError(RuntimeError):
    """Raised when a client tool cannot be delivered or completed."""


ClientToolResponseWaitMode = Literal["timeout", "session"]


class ClientToolPort(Protocol):
    async def call(
        self,
        function_name: str,
        arguments: dict[str, Any],
        *,
        timeout_seconds: float,
        wait_for_response: bool = True,
        response_wait_mode: ClientToolResponseWaitMode = "timeout",
    ) -> dict[str, Any]: ...


class ClientToolBroker(FrameProcessor):
    """Send client tool calls and correlate their app-message results."""

    def __init__(self) -> None:
        super().__init__()
        self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._closed_message: str | None = None

    async def call(
        self,
        function_name: str,
        arguments: dict[str, Any],
        *,
        timeout_seconds: float,
        wait_for_response: bool = True,
        response_wait_mode: ClientToolResponseWaitMode = "timeout",
    ) -> dict[str, Any]:
        from uuid import uuid4

        if self._closed_message is not None:
            raise ClientToolError(self._closed_message)

        tool_call_id = f"client_{uuid4().hex}"
        message = {
            "type": "client-tool-call",
            "tool_call_id": tool_call_id,
            "function_name": function_name,
            "arguments": dict(arguments),
            "wait_for_response": wait_for_response,
        }
        if not wait_for_response:
            try:
                await self.push_frame(
                    OutputTransportMessageUrgentFrame(message=message)
                )
            except Exception as exc:
                raise ClientToolError(
                    f"客户端工具调用发送失败: {function_name}"
                ) from exc
            return {
                "status": "ok",
                "data": {"dispatched": True},
            }

        loop = asyncio.get_running_loop()
        future: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[tool_call_id] = future
        try:
            await self.push_frame(
                OutputTransportMessageUrgentFrame(message=message)
            )
            if response_wait_mode == "session":
                return await future
            return await asyncio.wait_for(future, timeout=timeout_seconds)
        except TimeoutError as exc:
            raise ClientToolError(f"客户端工具调用超时: {function_name}") from exc
        except ClientToolError:
            raise
        except Exception as exc:
            raise ClientToolError(f"客户端工具调用失败: {function_name}") from exc
        finally:
            self._pending.pop(tool_call_id, None)

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, (EndFrame, CancelFrame)):
            self._close(
                "会话已结束" if isinstance(frame, EndFrame) else "会话已取消"
            )
            await self.push_frame(frame, direction)
            return

        if isinstance(frame, StopFrame):
            # StopFrame keeps processors reusable, so release only the calls
            # owned by the stopped run without permanently closing the broker.
            self._fail_pending("管线已停止")
            await self.push_frame(frame, direction)
            return

        if not isinstance(frame, InputTransportMessageFrame):
            await self.push_frame(frame, direction)
            return

        message = frame.message
        if not isinstance(message, dict) or message.get("type") != "client-tool-result":
            await self.push_frame(frame, direction)
            return

        tool_call_id = str(message.get("tool_call_id") or "")
        future = self._pending.get(tool_call_id)
        if future is None or future.done():
            logger.debug(f"忽略未知或过期的客户端工具结果: {tool_call_id}")
            return

        status = str(message.get("status") or "error")
        if status == "ok":
            future.set_result(
                {
                    "status": "ok",
                    "data": message.get("data"),
                }
            )
        else:
            future.set_result(
                {
                    "status": "error",
                    "message": str(message.get("message") or "客户端工具执行失败"),
                    "data": message.get("data"),
                }
            )

    def _fail_pending(self, message: str) -> None:
        for future in self._pending.values():
            if not future.done():
                future.set_exception(ClientToolError(message))
        self._pending.clear()

    def _close(self, message: str) -> None:
        """Close once and reject both current and future client-tool calls."""
        if self._closed_message is None:
            self._closed_message = message
        self._fail_pending(self._closed_message)

    async def cleanup(self) -> None:
        self._close("客户端工具通道已关闭")
        await super().cleanup()