"""Pre-response LLM routing for Workflow Agent edges.

The router deliberately uses a separate, short completion.  Its only output is
a required function choice, so the current Agent cannot speak before the graph
has decided whether the user turn belongs to another node.
"""

from __future__ import annotations

import json
from collections.abc import Callable
from typing import Any

from loguru import logger
from models import AssistantConfig
from openai import AsyncOpenAI


STAY_ON_CURRENT_AGENT = "workflow_stay_on_current_agent"
MAX_ROUTING_HISTORY_ENTRIES = 20


class WorkflowLLMRouter:
    """Select one LLM edge before the conversational LLM is allowed to reply."""

    def __init__(self, cfg: AssistantConfig):
        self._cfg = cfg

    async def select_edge(
        self,
        *,
        node_name: str,
        node_prompt: str,
        edges: list[dict[str, Any]],
        history: list[dict[str, str]],
        variables: dict[str, Any],
        edge_name: Callable[[dict[str, Any]], str],
        edge_description: Callable[[dict[str, Any]], str],
    ) -> str | None:
        """Return an edge function name, STAY, or None when routing failed."""
        if not edges:
            return STAY_ON_CURRENT_AGENT

        names = {edge_name(edge) for edge in edges}
        stay_name = STAY_ON_CURRENT_AGENT
        while stay_name in names:
            stay_name = f"_{stay_name}"

        tools = [
            {
                "type": "function",
                "function": {
                    "name": edge_name(edge),
                    "description": edge_description(edge),
                    "parameters": {"type": "object", "properties": {}},
                },
            }
            for edge in edges
        ]
        tools.append(
            {
                "type": "function",
                "function": {
                    "name": stay_name,
                    "description": "所有转移条件都不满足，继续由当前 Agent 处理用户消息。",
                    "parameters": {"type": "object", "properties": {}},
                },
            }
        )

        ordered_conditions = "\n".join(
            f"{index + 1}. {edge_description(edge)}"
            for index, edge in enumerate(edges)
        )
        router_prompt = (
            "你是工作流路由器，不是对话助手。收到一轮完整用户输入后，"
            "必须且只能调用一个提供的函数，禁止输出任何口头回复。\n"
            "按给出的顺序判断转移条件；选择第一个明确满足的转移函数。"
            "如果没有条件满足，调用留在当前 Agent 的函数。\n\n"
            f"当前节点：{node_name}\n"
            f"当前节点任务：{node_prompt or '未配置'}\n"
            f"转移条件：\n{ordered_conditions}"
        )
        recent_history = history[-MAX_ROUTING_HISTORY_ENTRIES:]
        routing_input = json.dumps(
            {
                "conversation": recent_history,
                "session_variables": variables,
            },
            ensure_ascii=False,
            separators=(",", ":"),
        )
        extra_body = self._cfg.llm_values.get("extraBody")
        request_extra = (
            {"extra_body": extra_body} if isinstance(extra_body, dict) else {}
        )
        client = AsyncOpenAI(
            api_key=self._cfg.llm_api_key,
            base_url=self._cfg.llm_base_url,
            timeout=15.0,
        )
        try:
            response = await client.chat.completions.create(
                model=self._cfg.model,
                messages=[
                    {"role": "system", "content": router_prompt},
                    {"role": "user", "content": routing_input},
                ],
                tools=tools,
                tool_choice="required",
                temperature=0,
                **request_extra,
            )
            tool_calls = response.choices[0].message.tool_calls or []
            if not tool_calls:
                logger.warning("Workflow 路由 LLM 未返回函数调用，留在当前 Agent")
                return STAY_ON_CURRENT_AGENT
            selected = str(tool_calls[0].function.name or "")
            if selected == stay_name:
                return STAY_ON_CURRENT_AGENT
            if selected not in names:
                logger.warning(f"Workflow 路由 LLM 返回未知函数:{selected}")
                return STAY_ON_CURRENT_AGENT
            return selected
        except Exception as exc:  # noqa: BLE001 - routing failure must not end the call
            logger.warning(f"Workflow LLM 边判断失败，留在当前 Agent:{exc}")
            return None
        finally:
            await client.close()