"""Resolve and apply one Workflow Agent's complete stage configuration."""

from __future__ import annotations

from copy import deepcopy

from models import AssistantConfig
from pipecat.flows import ContextStrategy, ContextStrategyConfig, NodeConfig
from pipecat.frames.frames import LLMUpdateSettingsFrame
from pipecat.services.settings import LLMSettings

from services.brains.base import BrainRuntime
from services.runtime_variables import DynamicVariableStore
from services.vision import VISION_SYSTEM_HINT
from services.workflow_engine import WorkflowEngine
from services.workflow_router import WorkflowLLMRouter


AGENT_STAGE_INSTRUCTION = (
    "工作流路由已在用户一轮输入结束时完成。只执行当前阶段任务，"
    "不要自行解释、模拟或宣布节点切换。"
)


class WorkflowAgentStage:
    """Translate graph-level Agent data into Pipecat Flows configuration."""

    def __init__(
        self,
        *,
        cfg: AssistantConfig,
        engine: WorkflowEngine,
        store: DynamicVariableStore,
        runtime: BrainRuntime,
    ) -> None:
        self._cfg = cfg
        self._engine = engine
        self._store = store
        self._runtime = runtime

    def role_message(self, node_id: str) -> str:
        stage_prompt = self._engine.prompt_for(node_id, self._store)
        stage = self._engine.agent_stage_config(node_id)
        if stage.vision_enabled:
            stage_prompt = f"{stage_prompt}\n\n[视觉能力]\n{VISION_SYSTEM_HINT}"
        return (
            f"{stage_prompt}\n\n[工作流执行规则]\n"
            f"{AGENT_STAGE_INSTRUCTION}"
        )

    async def refresh_prompt(self, node_id: str) -> None:
        await self._runtime.queue_frame(
            LLMUpdateSettingsFrame(
                delta=LLMSettings(
                    system_instruction=self.role_message(node_id)
                )
            )
        )

    def router_for_node(
        self,
        node_id: str,
        default_router: WorkflowLLMRouter,
    ) -> WorkflowLLMRouter:
        resource_id = self._engine.agent_stage_config(node_id).llm_resource_id
        resource = self._cfg.workflow_model_resources.get(resource_id)
        if not resource:
            return default_router
        from services.pipecat.service_factory import config_with_resource

        return WorkflowLLMRouter(config_with_resource(self._cfg, resource))

    async def apply(self, node_id: str) -> None:
        stage = self._engine.agent_stage_config(node_id)
        if self._runtime.set_input_enabled:
            self._runtime.set_input_enabled(True)
        if self._runtime.apply_turn_config:
            await self._runtime.apply_turn_config(
                stage.enable_interrupt,
                stage.turn_config,
            )
        if self._runtime.switch_services:
            await self._runtime.switch_services(
                stage.llm_resource_id or None,
                stage.asr_resource_id or None,
                stage.tts_resource_id or None,
            )
        if self._runtime.set_knowledge_scope:
            self._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,
                }
            )
        if self._runtime.set_vision_scope:
            self._runtime.set_vision_scope(
                {
                    "enabled": stage.vision_enabled,
                    "vision_model_resource_id": stage.vision_model_resource_id,
                    "llm_resource_id": stage.llm_resource_id,
                }
            )

    def node_config(
        self,
        node_id: str,
        *,
        functions: list,
        greeting_context_message: dict[str, str] | None,
        leading_messages: list[dict[str, str]] | None = None,
    ) -> 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 ""))
        strategy = (
            ContextStrategy.RESET
            if data.get("contextPolicy") == "fresh"
            else ContextStrategy.APPEND
        )
        greeting_messages = (
            [deepcopy(greeting_context_message)]
            if strategy == ContextStrategy.RESET and greeting_context_message
            else []
        )
        fixed_reply_messages = (
            [{"role": "assistant", "content": entry_speech}]
            if entry_mode == "fixed_speech" and entry_speech
            else []
        )
        return {
            "name": node_id,
            "role_message": self.role_message(node_id),
            "task_messages": [
                *greeting_messages,
                *(leading_messages or []),
                *fixed_reply_messages,
            ],
            "functions": functions,
            "context_strategy": ContextStrategyConfig(strategy=strategy),
            # Direct node activations let WorkflowRuntime decide whether to run
            # the LLM. A Pipecat function transition may explicitly override
            # this because FlowManager must coordinate the tool result first.
            "respond_immediately": False,
        }