"""Pipecat Flows-backed Workflow v3 brain."""

from __future__ import annotations

from typing import Any

from loguru import logger
from models import AssistantConfig, RuntimeTool
from db.session import SessionLocal
from pipecat.flows import (
    ContextStrategy,
    ContextStrategyConfig,
    FlowManager,
    FlowsFunctionSchema,
    NodeConfig,
)
from pipecat.frames.frames import (
    LLMUpdateSettingsFrame,
    OutputTransportMessageUrgentFrame,
    TTSSpeakFrame,
)
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.services.settings import LLMSettings

from services.brains.base import BaseBrain, BrainRuntime, BrainSpec
from services.knowledge import search as search_knowledge
from services.runtime_variables import DynamicVariableStore
from services.tool_executor import ToolExecutionError, ToolExecutor
from services.workflow_engine import WorkflowEngine


MAX_AUTOMATIC_HOPS = 50
AGENT_STAGE_INSTRUCTION = (
    "完成当前阶段任务。需要流转时必须调用对应的转移函数；"
    "不要在调用转移函数后继续生成口头回复。"
)


class WorkflowBrain(BaseBrain):
    spec = BrainSpec(
        type="workflow",
        supported_runtime_modes=frozenset({"pipeline"}),
        owns_context=True,
    )

    def __init__(self, cfg_or_graph: AssistantConfig | dict[str, Any]):
        cfg = cfg_or_graph if isinstance(cfg_or_graph, AssistantConfig) else None
        graph = cfg.graph if cfg is not None else cfg_or_graph
        self._engine = WorkflowEngine(graph or {})
        if not self._engine.has_graph() or not self._engine.start_id:
            raise ValueError("WorkflowBrain 缺少有效的 Start 节点")
        self._cfg = cfg
        self._store = DynamicVariableStore.from_config(cfg or AssistantConfig(type="workflow"))
        self._tools = ToolExecutor(self._store)
        self._tool_by_id: dict[str, RuntimeTool] = {
            tool.id: tool for tool in (cfg.tools if cfg else [])
        }
        self._runtime: BrainRuntime | None = None
        self._manager: FlowManager | None = None
        self._ended = False

    async def greeting(self, cfg: AssistantConfig) -> str:
        return self._engine.greeting(self._store) or cfg.greeting

    def system_prompt(self, cfg: AssistantConfig) -> str:
        return self._store.render(self._engine.global_prompt())

    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:
        if runtime.worker is None or runtime.context_aggregator is None:
            raise RuntimeError("WorkflowBrain 需要 PipelineWorker 和 context aggregator pair")
        self._cfg = cfg
        self._runtime = runtime
        self._store = DynamicVariableStore.from_config(cfg)
        self._tools = ToolExecutor(self._store)
        self._tool_by_id = {tool.id: tool for tool in cfg.tools}
        self._manager = FlowManager(
            worker=runtime.worker,
            llm=runtime.llm,
            context_aggregator=runtime.context_aggregator,
            transport=runtime.transport,
            global_functions=runtime.flow_global_functions,
        )
        self._manager.state["variables"] = self._store.values

    async def on_connected(self) -> None:
        await self._emit_node_active(self._engine.start_id)
        edge = self._engine.deterministic_edge(
            self._engine.start_id,
            self._store,
            include_default=True,
        )
        if not edge:
            raise RuntimeError("Start 初始化后没有命中的表达式边或默认边")
        node_config = await self._follow_edge(edge)
        if self._manager is None:
            raise RuntimeError("Workflow FlowManager 尚未初始化")
        await self._manager.initialize(node_config)
        logger.info(f"工作流模式启用: 当前节点={self._manager.current_node}")

    def record_user_message(self, content: str) -> None:
        if content and not self._ended:
            self._store.record("user", content)

    async def on_assistant_text_end(
        self,
        _turn_id: str,
        content: str,
        interrupted: bool,
    ) -> None:
        if not content or interrupted or self._ended:
            return
        self._store.record("agent", content, completed_agent_turn=True)
        manager = self._require_manager()
        current = manager.current_node
        if not current or self._engine.node_type(current) != "agent":
            return
        await self._refresh_agent_prompt(current)
        edge = self._engine.deterministic_edge(
            current,
            self._store,
            include_default=False,
        )
        if edge and manager.current_node == current:
            next_config = await self._follow_edge(edge)
            await manager.set_node_from_config(next_config)

    async def _refresh_agent_prompt(self, node_id: str) -> None:
        runtime = self._require_runtime()
        await runtime.queue_frame(
            LLMUpdateSettingsFrame(
                delta=LLMSettings(
                    system_instruction=self._agent_role_message(node_id)
                )
            )
        )

    def _agent_role_message(self, node_id: str) -> str:
        """Build one provider-compatible system instruction for an Agent stage."""
        stage_prompt = self._engine.prompt_for(node_id, self._store)
        return f"{stage_prompt}\n\n[工作流执行规则]\n{AGENT_STAGE_INSTRUCTION}"

    async def _apply_agent_stage(self, node_id: str) -> None:
        data = self._engine.data(node_id)
        await self._emit_node_active(node_id)
        if self._runtime and self._runtime.set_input_enabled:
            self._runtime.set_input_enabled(True)
        asr_id = str(
            data.get("asrResourceId")
            or self._engine.settings.get("defaultAsrResourceId")
            or ""
        )
        tts_id = str(
            data.get("ttsResourceId")
            or self._engine.settings.get("defaultTtsResourceId")
            or ""
        )
        runtime = self._require_runtime()
        if runtime.switch_services:
            await runtime.switch_services(asr_id or None, tts_id or None)
        if runtime.set_knowledge_scope:
            runtime.set_knowledge_scope(
                {
                    "knowledge_base_id": data.get("knowledgeBaseId"),
                    "mode": data.get("knowledgeMode", "disabled"),
                    "top_n": int(data.get("knowledgeTopN") or 5),
                    "score_threshold": float(data.get("knowledgeScoreThreshold") or 0.0),
                }
            )

    def _agent_config(self, node_id: str) -> NodeConfig:
        data = self._engine.data(node_id)
        strategy = (
            ContextStrategy.RESET
            if data.get("contextPolicy") == "fresh"
            else ContextStrategy.APPEND
        )
        functions: list[FlowsFunctionSchema] = []
        for tool_id in data.get("toolIds") or []:
            tool = self._tool_by_id.get(str(tool_id))
            if tool and tool.type == "http":
                functions.append(self._flow_tool(tool, node_id))
        knowledge_function = self._knowledge_function(node_id)
        if knowledge_function:
            functions.append(knowledge_function)
        for edge in self._engine.llm_edges(node_id):
            functions.append(self._flow_edge(edge))
        return {
            "name": node_id,
            "role_message": self._agent_role_message(node_id),
            "task_messages": [],
            "functions": functions,
            "context_strategy": ContextStrategyConfig(strategy=strategy),
            "respond_immediately": True,
        }

    def _terminal_config(self, node_id: str) -> NodeConfig:
        return {
            "name": node_id,
            "role_message": self._store.render(self._engine.global_prompt()),
            "task_messages": [],
            "functions": [],
            "context_strategy": ContextStrategyConfig(strategy=ContextStrategy.APPEND),
            "respond_immediately": False,
        }

    def _flow_tool(self, tool: RuntimeTool, node_id: str) -> FlowsFunctionSchema:
        properties, required = self._tools.schema_parts(tool)
        self._tools.register_secrets(tool)

        async def handler(args, _flow_manager):
            try:
                result = await self._tools.execute(tool, dict(args or {}))
            except ToolExecutionError as exc:
                return {"status": "error", "message": str(exc)}
            if result.get("updated_variables"):
                await self._refresh_agent_prompt(node_id)
                edge = self._engine.deterministic_edge(
                    node_id,
                    self._store,
                    include_default=False,
                )
                if edge:
                    return result, await self._follow_edge(edge)
            return result

        return FlowsFunctionSchema(
            name=tool.function_name,
            description=tool.description or f"调用 {tool.name}",
            properties=properties,
            required=required,
            handler=handler,
        )

    def _flow_edge(self, edge: dict) -> FlowsFunctionSchema:
        async def handler(_args, _flow_manager):
            return None, await self._follow_edge(edge)

        return FlowsFunctionSchema(
            name=self._engine.edge_fn_name(edge),
            description=self._engine.edge_description(edge),
            properties={},
            required=[],
            handler=handler,
        )

    def _knowledge_function(self, node_id: str) -> FlowsFunctionSchema | None:
        data = self._engine.data(node_id)
        knowledge_id = str(data.get("knowledgeBaseId") or "")
        if not knowledge_id or data.get("knowledgeMode") != "on_demand":
            return None
        cfg = self._cfg or AssistantConfig(type="workflow")
        knowledge = cfg.workflow_knowledge_bases.get(knowledge_id)
        description = "在当前 Agent 绑定的知识库中检索资料。"
        if knowledge:
            description += f"知识库：{knowledge.name}。{knowledge.description}"

        async def handler(args, _flow_manager):
            query = str((args or {}).get("query") or "").strip()
            if not query:
                return {"status": "error", "message": "检索问题为空"}
            try:
                async with SessionLocal() as session:
                    results = await search_knowledge(
                        session,
                        knowledge_id,
                        query,
                        top_k=int(data.get("knowledgeTopN") or 5),
                        score_threshold=float(data.get("knowledgeScoreThreshold") or 0.0),
                    )
                return {"status": "ok", "results": results}
            except Exception as exc:  # noqa: BLE001 - tool errors are returned to the LLM
                logger.warning(f"Workflow 知识库检索失败:{exc}")
                return {"status": "error", "message": "知识库检索暂时不可用"}

        return FlowsFunctionSchema(
            name="search_knowledge_base",
            description=description,
            properties={
                "query": {"type": "string", "description": "完整问题或检索关键词"}
            },
            required=["query"],
            handler=handler,
        )

    async def _follow_edge(self, edge: dict) -> NodeConfig:
        speech = self._engine.edge_transition_speech(edge)
        if speech:
            await self._require_runtime().queue_frame(
                TTSSpeakFrame(self._store.render(speech), append_to_context=False)
            )
        return await self._resolve_path(str(edge.get("target") or ""))

    async def _resolve_path(self, node_id: str) -> NodeConfig:
        for _ in range(MAX_AUTOMATIC_HOPS):
            node_type = self._engine.node_type(node_id)
            if node_type == "agent":
                await self._apply_agent_stage(node_id)
                return self._agent_config(node_id)
            if node_type == "end":
                await self._enter_end(node_id)
                return self._terminal_config(node_id)
            if node_type == "action":
                await self._enter_action(node_id)
            elif node_type == "handoff":
                await self._enter_handoff(node_id)
            elif node_type == "start":
                await self._emit_node_active(node_id)
            else:
                raise RuntimeError(f"工作流指向未知节点:{node_id}")
            edge = self._engine.deterministic_edge(
                node_id,
                self._store,
                include_default=True,
            )
            if not edge:
                raise RuntimeError(f"自动节点 {node_id} 没有命中的表达式边或默认边")
            speech = self._engine.edge_transition_speech(edge)
            if speech:
                await self._require_runtime().queue_frame(
                    TTSSpeakFrame(self._store.render(speech), append_to_context=False)
                )
            node_id = str(edge.get("target") or "")
        raise RuntimeError("工作流连续自动跳转超过安全上限")

    async def _enter_action(self, node_id: str) -> None:
        await self._emit_node_active(node_id)
        data = self._engine.data(node_id)
        tool_id = str(data.get("toolId") or "")
        tool = self._tool_by_id.get(tool_id)
        if not tool:
            self._store.values["system__last_action_status"] = "error"
            self._store.values["system__last_action_error"] = f"工具不存在:{tool_id}"
            return
        try:
            arguments = self._store.render_data(data.get("arguments") or {})
            await self._tools.execute(
                tool,
                arguments,
                result_assignments=data.get("resultAssignments") or {},
            )
            self._store.values["system__last_action_status"] = "ok"
            self._store.values["system__last_action_error"] = ""
        except (ToolExecutionError, ValueError) as exc:
            self._store.values["system__last_action_status"] = "error"
            self._store.values["system__last_action_error"] = str(exc)[:2048]

    async def _enter_handoff(self, node_id: str) -> None:
        await self._emit_node_active(node_id)
        data = self._engine.data(node_id)
        message = self._store.render(str(data.get("message") or ""))
        await self._require_runtime().queue_frame(
            OutputTransportMessageUrgentFrame(
                message={
                    "type": "handoff-requested",
                    "nodeId": node_id,
                    "targetType": data.get("targetType", "human"),
                    "target": data.get("target", ""),
                    "message": message,
                }
            )
        )
        if message:
            await self._require_runtime().queue_frame(
                TTSSpeakFrame(message, append_to_context=False)
            )
        self._store.values["system__handoff_status"] = "requested"

    async def _enter_end(self, node_id: str) -> None:
        self._ended = True
        await self._emit_node_active(node_id)
        runtime = self._require_runtime()
        if runtime.set_knowledge_scope:
            runtime.set_knowledge_scope({"mode": "disabled"})
        if runtime.set_input_enabled:
            runtime.set_input_enabled(False)
        data = self._engine.data(node_id)
        message = self._store.render(str(data.get("message") or ""))
        scope = str(data.get("scope") or "session")
        if scope == "flow":
            await runtime.queue_frame(
                OutputTransportMessageUrgentFrame(
                    message={"type": "flow-ended", "nodeId": node_id}
                )
            )
            if message:
                await runtime.queue_frame(TTSSpeakFrame(message, append_to_context=False))
            return
        runtime.call_end.begin("workflow_completed")
        if message:
            runtime.call_end.arm_after_speech()
            await runtime.queue_frame(TTSSpeakFrame(message, append_to_context=False))
        else:
            await runtime.call_end.finish()

    async def _emit_node_active(self, node_id: str | None) -> None:
        if node_id:
            await self._require_runtime().queue_frame(
                OutputTransportMessageUrgentFrame(
                    message={"type": "node-active", "nodeId": node_id}
                )
            )

    def _require_runtime(self) -> BrainRuntime:
        if self._runtime is None:
            raise RuntimeError("WorkflowBrain 尚未绑定 pipeline runtime")
        return self._runtime

    def _require_manager(self) -> FlowManager:
        if self._manager is None:
            raise RuntimeError("Workflow FlowManager 尚未初始化")
        return self._manager