"""Local prompt assistant, including prompt-only reusable tools."""

from __future__ import annotations

from uuid import uuid4

from models import AssistantConfig
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.frames.frames import OutputTransportMessageUrgentFrame, TTSSpeakFrame
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.services.llm_service import (
    FunctionCallParams,
    FunctionCallResultProperties,
)
from pipecat.utils.time import time_now_iso8601

from services.brains.base import BaseBrain, BrainRuntime, BrainSpec


class PromptBrain(BaseBrain):
    spec = BrainSpec(
        type="prompt",
        supported_runtime_modes=frozenset({"pipeline", "realtime"}),
        owns_context=True,
    )

    def build_llm(self, cfg: AssistantConfig, context: LLMContext) -> FrameProcessor:
        from services.pipecat.service_factory import create_llm

        return create_llm(cfg)

    async def setup(self, cfg: AssistantConfig, runtime: BrainRuntime) -> None:
        schemas: list[FunctionSchema] = []
        for tool in cfg.tools:
            if tool.type != "end_call":
                continue
            schema, handler = self._make_end_call_tool(tool, runtime)
            schemas.append(schema)
            runtime.llm.register_function(tool.function_name, handler)
        runtime.set_tools(schemas)

    @staticmethod
    def _make_end_call_tool(tool, runtime: BrainRuntime):
        config = (tool.definition or {}).get("config") or {}
        message_type = str(config.get("message_type") or "none")
        custom_message = str(config.get("custom_message") or "").strip()
        capture_reason = bool(config.get("capture_reason", True))

        async def end_call(params: FunctionCallParams) -> None:
            reason = str(params.arguments.get("reason") or "end_call_tool").strip()
            runtime.call_end.begin(reason)
            await params.result_callback(
                {"status": "success", "action": "ending_call"},
                properties=FunctionCallResultProperties(run_llm=False),
            )

            if message_type != "custom" or not custom_message:
                await runtime.call_end.finish()
                return

            turn_id = uuid4().hex
            timestamp = time_now_iso8601()
            for message in (
                {
                    "type": "assistant-text-start",
                    "turn_id": turn_id,
                    "timestamp": timestamp,
                },
                {
                    "type": "assistant-text-delta",
                    "turn_id": turn_id,
                    "delta": custom_message,
                },
                {
                    "type": "assistant-text-end",
                    "turn_id": turn_id,
                    "content": custom_message,
                    "interrupted": False,
                },
            ):
                await runtime.queue_frame(
                    OutputTransportMessageUrgentFrame(message=message)
                )
            runtime.call_end.arm_after_speech()
            await runtime.queue_frame(
                TTSSpeakFrame(custom_message, append_to_context=False)
            )

        properties = (
            {
                "reason": {
                    "type": "string",
                    "description": "结束本次通话的简短原因。",
                }
            }
            if capture_reason
            else {}
        )
        schema = FunctionSchema(
            name=tool.function_name,
            description=tool.description or "结束当前通话。",
            properties=properties,
            required=["reason"] if capture_reason else [],
        )
        return schema, end_call