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

from __future__ import annotations

from copy import deepcopy
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 (
    LLMRunFrame,
    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 pipecat.utils.time import time_now_iso8601

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
from services.workflow_router import STAY_ON_CURRENT_AGENT, WorkflowLLMRouter


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 = deepcopy(cfg.graph if cfg is not None else cfg_or_graph)
        if cfg is not None:
            # Graph v3 owns Workflow defaults. Keep older saved graphs compatible
            # by filling the new interaction settings from the assistant row.
            settings = graph.setdefault("settings", {})
            settings.setdefault("enableInterrupt", cfg.enableInterrupt)
            settings.setdefault("turnConfig", deepcopy(cfg.turnConfig))
        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._router = WorkflowLLMRouter(cfg or AssistantConfig(type="workflow"))
        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._router = WorkflowLLMRouter(cfg)
        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)
        await self._emit_variables(
            reason="initialized",
            node_id=self._engine.start_id,
        )
        edge = self._engine.deterministic_edge(
            self._engine.start_id,
            self._store,
            include_default=True,
        )
        if not edge and self._engine.has_outgoing(self._engine.start_id):
            raise RuntimeError("Start 初始化后没有命中的表达式边或默认边")
        node_config = (
            await self._follow_edge(edge)
            if edge
            else self._passive_node_config(self._engine.start_id)
        )
        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_user_turn_end(self, content: str) -> bool:
        """Route a complete user turn before any Agent is allowed to reply."""
        if not content or self._ended:
            return True
        self.record_user_message(content)
        manager = self._require_manager()
        current = manager.current_node
        if not current or self._engine.node_type(current) != "agent":
            return True

        edge = self._engine.deterministic_edge(
            current,
            self._store,
            include_default=False,
        )
        outgoing = self._engine.outgoing(current)
        llm_edges = [
            candidate
            for candidate in outgoing
            if self._engine.edge_mode(candidate) == "llm"
        ]
        default_edge = next(
            (
                candidate
                for candidate in outgoing
                if self._engine.edge_mode(candidate) == "always"
            ),
            None,
        )

        if edge is None and llm_edges:
            selected = await self._router_for_node(current).select_edge(
                node_name=self._engine.name(current),
                node_prompt=self._engine.prompt_for(current, self._store),
                edges=llm_edges,
                history=self._store.history,
                variables={
                    key: value
                    for key, value in self._store.values.items()
                    if not key.startswith("system__")
                },
                edge_name=self._engine.edge_fn_name,
                edge_description=self._engine.edge_description,
            )
            if selected and selected != STAY_ON_CURRENT_AGENT:
                edge = next(
                    (
                        candidate
                        for candidate in llm_edges
                        if self._engine.edge_fn_name(candidate) == selected
                    ),
                    None,
                )
            elif selected == STAY_ON_CURRENT_AGENT:
                edge = default_edge
        elif edge is None and not llm_edges:
            edge = default_edge

        if edge and manager.current_node == current:
            next_config = await self._follow_edge(edge)
            await manager.set_node_from_config(next_config)
            return True

        # The incoming LLMContextFrame is intentionally suppressed by the
        # pipeline router. Queue prompt refresh + inference in this order so
        # this user turn is answered with the current Agent's latest variables.
        await self._refresh_agent_prompt(current)
        await self._require_runtime().queue_frame(LLMRunFrame())
        return True

    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)

    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}"

    def _router_for_node(self, node_id: str) -> WorkflowLLMRouter:
        stage = self._engine.agent_stage_config(node_id)
        resource_id = stage.llm_resource_id
        cfg = self._cfg
        resource = cfg.workflow_model_resources.get(resource_id) if cfg else None
        if not cfg or not resource:
            return self._router
        from services.pipecat.service_factory import config_with_resource

        return WorkflowLLMRouter(config_with_resource(cfg, resource))

    async def _apply_agent_stage(self, node_id: str) -> None:
        stage = self._engine.agent_stage_config(node_id)
        await self._emit_node_active(node_id)
        if self._runtime and self._runtime.set_input_enabled:
            self._runtime.set_input_enabled(True)
        runtime = self._require_runtime()
        if runtime.apply_turn_config:
            await runtime.apply_turn_config(
                stage.enable_interrupt,
                stage.turn_config,
            )
        if runtime.switch_services:
            await runtime.switch_services(
                stage.llm_resource_id or None,
                stage.asr_resource_id or None,
                stage.tts_resource_id or None,
            )
        if runtime.set_knowledge_scope:
            runtime.set_knowledge_scope(
                {
                    "knowledge_base_id": stage.knowledge_base_id,
                    "mode": stage.knowledge_mode,
                    "top_n": stage.knowledge_top_n,
                    "score_threshold": stage.knowledge_score_threshold,
                }
            )

    def _agent_config(self, node_id: str) -> NodeConfig:
        data = self._engine.data(node_id)
        entry_mode = str(data.get("entryMode") or "wait_user")
        entry_speech = self._store.render(str(data.get("entrySpeech") or ""))
        fixed_reply_messages = (
            [{"role": "assistant", "content": entry_speech}]
            if entry_mode == "fixed_speech" and entry_speech
            else []
        )
        strategy = (
            ContextStrategy.RESET
            if data.get("contextPolicy") == "fresh"
            else ContextStrategy.APPEND
        )
        stage = self._engine.agent_stage_config(node_id)
        functions: list[FlowsFunctionSchema] = []
        for tool_id in stage.tool_ids:
            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)
        config: NodeConfig = {
            "name": node_id,
            "role_message": self._agent_role_message(node_id),
            # Flows writes task_messages into the Pipecat LLM context. The
            # pre-action below is responsible only for display, persistence,
            # dynamic conversation history, and TTS playback.
            "task_messages": fixed_reply_messages,
            "functions": functions,
            "context_strategy": ContextStrategyConfig(strategy=strategy),
            "respond_immediately": entry_mode == "generate",
        }
        if entry_mode == "fixed_speech":
            config["pre_actions"] = [
                {
                    "type": "workflow_fixed_speech",
                    "text": entry_speech,
                    "node_id": node_id,
                    "handler": self._play_fixed_speech,
                }
            ]
        return config

    async def _play_fixed_speech(self, action: dict, _flow_manager: FlowManager) -> None:
        """Play and persist Agent entry speech without creating an LLM turn."""
        await self._queue_visible_speech(
            str(action.get("text") or ""),
            source="workflow-fixed-reply",
            node_id=str(action.get("node_id") or "") or None,
        )

    async def _queue_visible_speech(
        self,
        text: str,
        *,
        source: str = "workflow-speech",
        node_id: str | None = None,
    ) -> None:
        """Show and persist fixed workflow speech before sending it to TTS."""
        content = text.strip()
        if not content:
            return
        self._store.record("agent", content)
        runtime = self._require_runtime()
        await runtime.queue_frame(
            OutputTransportMessageUrgentFrame(
                message={
                    "type": "transcript",
                    "role": "assistant",
                    "content": content,
                    "timestamp": time_now_iso8601(),
                    "source": source,
                    **({"nodeId": node_id} if node_id else {}),
                }
            )
        )
        await runtime.queue_frame(TTSSpeakFrame(content, append_to_context=False))

    def _passive_node_config(self, node_id: str) -> NodeConfig:
        """Keep a non-conversational terminal node active without ending the call."""
        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)}
            updated_variables = list(result.get("updated_variables") or [])
            if updated_variables:
                await self._emit_variables(
                    reason="tool",
                    node_id=node_id,
                    changed=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,
            cancel_on_interruption=True,
        )

    def _knowledge_function(self, node_id: str) -> FlowsFunctionSchema | None:
        stage = self._engine.agent_stage_config(node_id)
        knowledge_id = str(stage.knowledge_base_id or "")
        if not knowledge_id or stage.knowledge_mode != "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=stage.knowledge_top_n,
                        score_threshold=stage.knowledge_score_threshold,
                    )
                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,
            cancel_on_interruption=True,
        )

    async def _follow_edge(self, edge: dict) -> NodeConfig:
        speech = self._engine.edge_transition_speech(edge)
        if speech:
            await self._queue_visible_speech(self._store.render(speech))
        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._passive_node_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}")
            if not self._engine.has_outgoing(node_id):
                return self._passive_node_config(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._queue_visible_speech(self._store.render(speech))
            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 {})
            result = await self._tools.execute(
                tool,
                arguments,
                result_assignments=data.get("resultAssignments") or {},
            )
            updated_variables = list(result.get("updated_variables") or [])
            if updated_variables:
                await self._emit_variables(
                    reason="action",
                    node_id=node_id,
                    changed=updated_variables,
                )
            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._queue_visible_speech(message)
        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 self._queue_visible_speech(message)
            return
        runtime.call_end.begin("workflow_completed")
        if message:
            runtime.call_end.arm_after_speech()
            await self._queue_visible_speech(message)
        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 _public_variables(self) -> dict[str, str | int | float | bool]:
        """Return the browser-safe part of this session's variable state."""
        return {
            name: value
            for name, value in self._store.values.items()
            if not name.startswith(("system__", "secret__"))
            and isinstance(value, (str, int, float, bool))
        }

    async def _emit_variables(
        self,
        *,
        reason: str,
        node_id: str | None,
        changed: list[str] | None = None,
    ) -> None:
        """Publish a safe snapshot so Workflow debug mirrors runtime state."""
        message: dict[str, Any] = {
            "type": "workflow-variables",
            "reason": reason,
            "variables": self._public_variables(),
        }
        if node_id:
            message["nodeId"] = node_id
        if changed:
            message["changed"] = [
                name
                for name in changed
                if not name.startswith(("system__", "secret__"))
            ]
        await self._require_runtime().queue_frame(
            OutputTransportMessageUrgentFrame(message=message)
        )

    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