# Copyright (c) Alibaba, Inc. and its affiliates.
"""
数学 Agentic RL 的 Rollout 处理器（trace_client 版 — 推荐）。

本模块实现了基于 MCP 工具 + LangGraph 的 agent 推理流程：
  1. 启动一个 MCP Server 进程，暴露 python_interpreter 工具
  2. 通过 LangGraph 构建 agent → tools 循环图
  3. 每次 rollout 请求：构造消息 → 调用策略模型 → 若模型发起工具调用则执行 → 循环直至模型给出最终回答
  4. 收集对话轨迹和指标（延迟、工具调用次数/耗时）作为 RolloutOutput 返回

可观测性：
  - trace_client(llm)  : 在 ChatOpenAI 底层 completions.create() 处自动打点，
                          input/output/token_usage 均由 SDK 自动填充，零封装
  - trace_tool(tools)   : 在 BaseTool.invoke/ainvoke 处自动打点，每次工具调用独立 span
  - @observe_processor  : 追踪整个 process() 调用的 span

"""

import asyncio
import json
import logging
import multiprocessing
import socket
import time
from typing import Any, Dict, List

from langchain_core.messages import (
    AIMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
    convert_to_openai_messages,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langgraph.graph import END, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode
from mcp.server.fastmcp import FastMCP

from functions.rollout.code_interpreter import execute_python

from dashscope.finetune.reinforcement import RolloutInput, RolloutOutput
from dashscope.finetune.reinforcement.component.data.base_data_model import (
    AgentOutput,
    TaskStatus,
)
from dashscope.finetune.reinforcement.component.processor.abstract_rollout_processor import (
    AbstractRolloutProcessor,
)

from dashscope.finetune.reinforcement.component.observability import (
    observe_processor,
    trace_client,
    trace_tool,
)

logger = logging.getLogger(__name__)

MCP_PORT = 10086

# ============================================================================ #
#                             MCP SERVER（工具服务端）                           #
# ============================================================================ #

mcp = FastMCP("math_tools", port=MCP_PORT)


@mcp.tool()
async def python_interpreter(code: str) -> str:
    """执行 Python 脚本并返回 stdout + stderr。

    适用于单步算术之外的所有计算：多步推导、模运算、方程求解、
    枚举搜索、数值模拟、优化等。

    预导入模块: math, cmath, decimal, fractions, statistics, random,
      string, re, json, collections, itertools, functools, operator,
      io, sys, os, copy, bisect, heapq, typing, numpy (as np)。
    不可用: scipy, sympy, pandas 等第三方包。

    超时: 30 秒。输出上限: 8000 字符。
    """
    return await asyncio.to_thread(execute_python, code, 30)


def _run_mcp_server():
    """在子进程中启动 MCP Server（SSE 传输协议）。"""
    mcp.run(transport="sse")


def _wait_for_port(port: int, host: str = "localhost", timeout: float = 30.0) -> bool:
    """阻塞等待指定端口可连接，用于确保 MCP Server 就绪后再初始化客户端。"""
    start_time = time.perf_counter()
    while True:
        try:
            with socket.create_connection((host, port), timeout=1):
                return True
        except OSError:
            time.sleep(0.5)
            if time.perf_counter() - start_time >= timeout:
                raise TimeoutError(
                    f"[DashSystem] Cannot connect to {port} in {timeout} s!"
                )


# ============================================================================ #
#                          ROLLOUT PROCESSOR（推理处理器）                       #
# ============================================================================ #

DEFAULT_SYSTEM_PROMPT = (
    "You are a math assistant. You have one tool:\n"
    "- `python_interpreter`: run a Python script (numpy available as np; "
    "no scipy/sympy) for computation, equation solving, enumeration, "
    "modular arithmetic, etc.\n\n"
    "Always verify your calculations by calling the tool instead of computing "
    "in your head. Think step by step, then put your final answer in \\boxed{}."
)

DATASET_SYSTEM_PROMPTS: Dict[str, str] = {}


class CalcXRolloutProcessor(AbstractRolloutProcessor):
    """基于 MCP 工具的数学 Rollout 处理器（trace_client 版）。

    可观测性由 trace_client / trace_tool 在底层自动完成，
    业务代码无需任何额外封装。
    """

    def _is_port_in_use(self, port: int, host: str = "127.0.0.1") -> bool:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(1.0)
            return s.connect_ex((host, port)) == 0

    def __init__(self):
        super().__init__()
        self._init_lock = asyncio.Lock()

        if self._is_port_in_use(MCP_PORT):
            logger.info(
                f"MCP Server already running on port {MCP_PORT}. Skipping initialization."
            )
            self.mcp_process = None
        else:
            logger.info(f"Port {MCP_PORT} is free. Starting MCP Server process...")
            self.mcp_process = multiprocessing.Process(
                target=_run_mcp_server, daemon=True
            )
            self.mcp_process.start()
            try:
                _wait_for_port(MCP_PORT)
            except TimeoutError as e:
                logger.error(f"Failed to start MCP Server: {e}")
                raise

        self._shared_mcp_client = None
        self._shared_tools = None
        self._shared_graph = None

    async def _async_setup(self):
        if not self._shared_mcp_client:
            start_time = time.perf_counter()
            async with self._init_lock:
                if not self._shared_mcp_client:
                    await self._init_resources_async()
                    elapsed = time.perf_counter() - start_time
                    logger.info(f"[DEBUG] _async_setup: {elapsed:.4f}s")

    def _get_mcp_servers(self) -> dict:
        return {
            "calculator": {
                "transport": "sse",
                "url": f"http://localhost:{MCP_PORT}/sse",
            },
        }

    async def _init_resources_async(self, max_retries: int = 5, base_delay: float = 1.0):
        logger.info(
            "Initializing shared MCP Client and Graph for this worker instance..."
        )
        last_exc = None
        for attempt in range(1, max_retries + 1):
            try:
                client = MultiServerMCPClient(self._get_mcp_servers())
                tools = await client.get_tools()
                trace_tool(tools)
                self._shared_tool_node = ToolNode(tools)
                graph = self._build_graph()

                self._shared_tools = tools
                self._shared_graph = graph
                self._shared_mcp_client = client
                logger.info("Shared MCP Client successfully initialized. Tools cached.")
                return
            except Exception as e:
                last_exc = e
                delay = base_delay * attempt
                logger.warning(
                    f"MCP Client init attempt {attempt}/{max_retries} failed: {e}. "
                    f"Retrying in {delay}s..."
                )
                await asyncio.sleep(delay)
        raise RuntimeError(
            f"Failed to initialize MCP Client after {max_retries} attempts"
        ) from last_exc

    # ------------------------------------------------------------------ #
    #  LangGraph 状态机构建                                                #
    # ------------------------------------------------------------------ #

    async def _call_model(self, state: MessagesState, config: RunnableConfig):
        """Agent 节点：调用绑定了工具的 LLM。

        trace_client 已在 _build_llm 中 patch 了底层 client，
        每次 model.ainvoke → completions.create() 都会自动产生 LLM span。
        """
        model = config["configurable"]["model"]
        try:
            response = await model.ainvoke(state["messages"], config=config)
            return {"messages": [response]}
        except Exception as e:
            logger.error(f"Model invoke error: {e}")
            return {"messages": [("ai", f"Error: Model invocation failed. {e}")]}

    def _should_continue(self, state: MessagesState):
        last_message = state["messages"][-1]
        if hasattr(last_message, "tool_calls") and len(last_message.tool_calls) > 0:
            return "tools"
        return END

    def _build_graph(self) -> StateGraph:
        workflow = StateGraph(MessagesState)
        workflow.add_node("agent", self._call_model)
        workflow.add_node("tools", self._execute_tools)
        workflow.set_entry_point("agent")
        workflow.add_conditional_edges(
            "agent", self._should_continue, {"tools": "tools", END: END}
        )
        workflow.add_edge("tools", "agent")
        return workflow.compile()

    async def _execute_tools(self, state: MessagesState, config: RunnableConfig):
        """Tools 节点：执行工具调用并记录耗时。

        trace_tool 已在初始化时 patch 到每个 BaseTool 上，
        每次工具调用自动产生独立的 execute_tool span。
        """
        last_message = state["messages"][-1]
        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            for call in last_message.tool_calls:
                logger.info(
                    f"[Tool Call] Tool: {call.get('name')}, Args: {call.get('args')}"
                )
        tool_node = config["configurable"]["tool_node"]
        t0 = time.time()
        result = await tool_node.ainvoke(state, config)
        durations = config["configurable"].get("tool_call_durations")
        if durations is not None:
            durations.append(time.time() - t0)
        return result

    # ------------------------------------------------------------------ #
    #  辅助方法                                                            #
    # ------------------------------------------------------------------ #

    def _build_llm(self, input: RolloutInput) -> ChatOpenAI:
        resource = input.model_resource
        params = input.sampling_params or {}
        api_key = (
            resource.api_key.get_secret_value()
            if hasattr(resource.api_key, "get_secret_value")
            else resource.api_key
        )
        extra_kwargs = {
            k: v
            for k, v in params.items()
            if k not in ("temperature", "max_tokens", "timeout", "max_turns")
        }
        llm = ChatOpenAI(
            model=resource.model_name,
            openai_api_key=api_key,
            openai_api_base=resource.base_url,
            temperature=params.get("temperature", 0.7),
            max_tokens=params.get("max_tokens", 8192),
            request_timeout=params.get("timeout", None),
            model_kwargs=extra_kwargs,
            streaming=False,
        )
        trace_client(llm)
        return llm

    @staticmethod
    def _to_langchain_messages(messages: List[Dict]) -> list:
        lc_messages = []
        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role == "system":
                lc_messages.append(SystemMessage(content=content))
            elif role == "user":
                lc_messages.append(HumanMessage(content=content))
            elif role == "assistant":
                lc_messages.append(AIMessage(content=content))
            else:
                lc_messages.append(ChatMessage(role=role, content=content))
        return lc_messages

    @staticmethod
    def _get_system_prompt(rollout_extra: Dict[str, Any] | None) -> str:
        if rollout_extra:
            if "system_prompt" in rollout_extra:
                return rollout_extra["system_prompt"]
            dataset = rollout_extra.get("dataset", "")
            if dataset in DATASET_SYSTEM_PROMPTS:
                return DATASET_SYSTEM_PROMPTS[dataset]
        return DEFAULT_SYSTEM_PROMPT

    # ------------------------------------------------------------------ #
    #  核心 Rollout 逻辑                                                   #
    # ------------------------------------------------------------------ #

    async def _async_process(self, input: RolloutInput) -> RolloutOutput:
        rollout_id = "unknown"
        if input.request_metadata and input.request_metadata.rollout_id:
            rollout_id = input.request_metadata.rollout_id
        elif input.rollout_extra and "rollout_id" in input.rollout_extra:
            rollout_id = input.rollout_extra["rollout_id"]

        logger.info(
            f"[CalcXRolloutProcessor] Starting rollout {rollout_id} "
            f"with model {input.model_resource.model_name} at {input.model_resource.base_url}"
        )

        start_time = time.time()
        tool_call_durations: list[float] = []

        try:
            tools = self._shared_tools
            graph = self._shared_graph

            llm = self._build_llm(input)
            model_with_tools = llm.bind_tools(tools)

            system_prompt = self._get_system_prompt(input.rollout_extra)
            messages = [SystemMessage(content=system_prompt)] + \
                self._to_langchain_messages(input.messages)

            max_turns = (input.sampling_params or {}).get("max_turns", 10)
            config = RunnableConfig(
                recursion_limit=max_turns * 2 + 5,
                configurable={
                    "model": model_with_tools,
                    "tool_node": self._shared_tool_node,
                    "tool_call_durations": tool_call_durations,
                },
                metadata={
                    "rollout_id": rollout_id,
                    "model_name": input.model_resource.model_name,
                    "base_url": input.model_resource.base_url,
                },
            )

            final_state = await graph.ainvoke({"messages": messages}, config=config)
            final_messages = final_state["messages"]
            output_content = str(final_messages[-1].content)

            reward = 0.0
            latency = round(time.time() - start_time, 4)
            tool_total = sum(tool_call_durations)
            tool_avg = (tool_total / len(tool_call_durations)) if tool_call_durations else 0.0
            agent_output = AgentOutput(
                messages=convert_to_openai_messages(final_messages),
                rollout_extra=input.rollout_extra,
                rollout_metrics={
                    "latency": latency,
                    "tool_call_count": len(tool_call_durations),
                    "tool_call_total_time": round(tool_total, 4),
                    "tool_call_avg_time": round(tool_avg, 4),
                },
                reward_score=reward,
            )

            logger.info(
                f"Rollout result {rollout_id} | latency: {latency}s | "
                f"tool_calls: {len(tool_call_durations)} | "
                f"reward: {reward} | content: {json.dumps(output_content)}"
            )
            return RolloutOutput(
                agent_output=agent_output,
                status=TaskStatus.SUCCESS,
                error=None,
            )

        except Exception as e:
            logger.error(f"Graph execution failed: {e}", exc_info=True)
            latency = round(time.time() - start_time, 4)
            tool_total = sum(tool_call_durations)
            tool_avg = (tool_total / len(tool_call_durations)) if tool_call_durations else 0.0

            return RolloutOutput(
                agent_output=AgentOutput(
                    messages=[{"role": "system", "content": f"Error: {e}"}],
                    rollout_extra=input.rollout_extra,
                    rollout_metrics={
                        "latency": latency,
                        "tool_call_count": len(tool_call_durations),
                        "tool_call_total_time": round(tool_total, 4),
                        "tool_call_avg_time": round(tool_avg, 4),
                    },
                    reward_score=0.0,
                ),
                status=TaskStatus.FAILED,
                error=str(e),
            )

    @observe_processor
    async def process(self, input: RolloutInput) -> RolloutOutput:
        await self._async_setup()
        return await self._async_process(input)

    def __del__(self):
        if (
            hasattr(self, "mcp_process")
            and self.mcp_process is not None
            and self.mcp_process.is_alive()
        ):
            try:
                self.mcp_process.terminate()
                self.mcp_process.join(timeout=1)
                logger.info("MCP Server process terminated.")
            except Exception as e:
                logger.warning(f"Error terminating MCP Server process: {e}")


if __name__ == "__main__":
    from dashscope.finetune.reinforcement.component.data.base_data_model import (
        ModelResource,
        RequestMetadata,
    )
    import os
    os.environ.setdefault("NO_PROXY", "localhost,127.0.0.1,::1")

    rollout_input = RolloutInput(
        messages=[
            {"role": "user", "content": "6.6 minus x (3/2) times equals 5.6."},
        ],
        ground_truth="2/3",
        model_resource=ModelResource(
            model_name="qwen-plus",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            api_key="sk-XXXX",
        ),
        sampling_params={
            "temperature": 0.6,
            "top_p": 0.8,
            "max_tokens": 2048,
            "max_turns": 10,
        }
    )
    print(json.dumps(rollout_input.model_dump(mode="json"), indent=4))

    processor = CalcXRolloutProcessor()
    result = asyncio.run(processor.process(rollout_input))
    print(result)
