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

from __future__ import annotations

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,
        leading_messages: list[dict[str, str]] | None = None,
    ) -> NodeConfig:
        data = self._engine.data(node_id)
        strategy = (
            ContextStrategy.RESET
            if data.get("contextPolicy") == "fresh"
            else ContextStrategy.APPEND
        )
        return {
            "name": node_id,
            "role_message": self.role_message(node_id),
            "task_messages": list(leading_messages or []),
            "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,
        }