"""Rollout Processor for multi-turn code generation with tool use.

The model receives a competitive programming problem and can iteratively
debug its solution using an ``execute_code`` tool that runs the code against
public (sample) test cases.  After debugging, the model's final solution
(last ``python`` code block across all assistant messages) is scored by the
reward processor against ALL test cases.

Architecture:
- LangGraph agent/tools loop
- Dynamic ``StructuredTool`` per rollout, binding the current problem's
  public test cases via closure (no MCP server needed)
- Supports both stdin-based and functional (call-based) problem types
"""
from __future__ import annotations

import json
import logging
import time
from typing import Any, Dict, List, Optional

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, StructuredTool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode

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,
)

from functions.reward.code_testing import (
    extract_code,
    run_public_tests,
)

logger = logging.getLogger(__name__)

# ============================================================================ #
#  System prompts
# ============================================================================ #

SYSTEM_PROMPT_STDIN = """\
You are an expert competitive programmer. You have an `execute_code` tool \
that runs your Python solution against sample test cases and reports \
pass/fail, actual vs expected output, and any errors.

Workflow:
1. Analyze the problem and identify edge cases.
2. Write an initial solution.
3. Call `execute_code` to test it — do NOT skip this step.
4. If any test fails, read the diff carefully, fix, and retest.
5. Once all sample tests pass, output your final solution in a \
```python``` code block.

Rules:
- Read from stdin, print to stdout.
- Sample tests are a subset; handle all edge cases.
- You MUST call `execute_code` at least once before finalizing.
- To submit your final answer, write the complete solution in a \
```python``` code block in your response."""

SYSTEM_PROMPT_FUNCTIONAL = """\
You are an expert competitive programmer. You have an `execute_code` tool \
that runs your Python solution against sample test cases and reports \
pass/fail, actual vs expected output, and any errors.

Workflow:
1. Analyze the problem and identify edge cases.
2. Complete the required function/method inside a `Solution` class.
3. Call `execute_code` to test it — do NOT skip this step.
4. If any test fails, read the diff carefully, fix, and retest.
5. Once all sample tests pass, output your final solution in a \
```python``` code block.

Rules:
- Include necessary imports (e.g. `from typing import *`).
- Sample tests are a subset; handle all edge cases.
- You MUST call `execute_code` at least once before finalizing.
- To submit your final answer, write the complete solution in a \
```python``` code block in your response."""


# ============================================================================ #
#  Tool builder
# ============================================================================ #

def _resolve_public_tests(rollout_extra: Dict[str, Any]) -> list[dict]:
    """Extract public test cases list from rollout_extra."""
    raw = rollout_extra.get("public_test_cases", "[]")
    if isinstance(raw, str):
        raw = json.loads(raw)
    return raw


def _resolve_fn_name(rollout_extra: Dict[str, Any]) -> Optional[str]:
    """Try to extract the function name for call-based problems."""
    metadata = rollout_extra.get("metadata", "{}")
    if isinstance(metadata, str):
        metadata = json.loads(metadata)
    return metadata.get("func_name")


def _build_execute_code_tool(
    public_tests: list[dict],
    fn_name: Optional[str],
) -> StructuredTool:
    """Create a StructuredTool that captures the current problem's tests."""

    def _execute_code(code: str) -> str:
        """Run your Python code against sample test cases and return results.

        Use this to verify and debug your solution before finalizing.
        Pass your COMPLETE Python solution as the `code` argument.
        """
        extracted = extract_code(code)
        actual_code = extracted if extracted is not None else code
        return run_public_tests(actual_code, public_tests, fn_name=fn_name)

    return StructuredTool.from_function(
        func=_execute_code,
        name="execute_code",
        description=(
            "Run your Python solution against sample test cases.\n"
            "Returns per-case results: PASS/FAIL, actual vs expected output, "
            "and any runtime errors or timeouts.\n\n"
            "Usage: pass your COMPLETE Python solution as the `code` argument "
            "(you may wrap it in a ```python``` block or pass raw code).\n\n"
            "Limitations:\n"
            "- Only sample (public) test cases are run; the final score uses "
            "all test cases including hidden ones.\n"
            "- Each case has a 10-second timeout.\n"
            "- Always test at least once before giving your final answer."
        ),
    )


# ============================================================================ #
#  LangGraph agent
# ============================================================================ #

class CodeRolloutProcessor(AbstractRolloutProcessor):
    """Multi-turn rollout for code generation with tool-assisted debugging."""

    def __init__(self):
        super().__init__()

    # ------------------------------------------------------------------ #
    #  LangGraph construction
    # ------------------------------------------------------------------ #

    async def _call_model(self, state: MessagesState, config: RunnableConfig):
        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):
        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')}, "
                    f"Args length: {len(str(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

    # ------------------------------------------------------------------ #
    #  Helpers
    # ------------------------------------------------------------------ #

    @staticmethod
    def _build_llm(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", 16384),
            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 and rollout_extra.get("starter_code"):
            return SYSTEM_PROMPT_FUNCTIONAL
        return SYSTEM_PROMPT_STDIN

    def _build_tools(self, rollout_extra: Dict[str, Any]) -> List[BaseTool]:
        """Create the execute_code tool bound to this problem's public tests."""
        public_tests = _resolve_public_tests(rollout_extra)
        fn_name = _resolve_fn_name(rollout_extra)
        tool = _build_execute_code_tool(public_tests, fn_name)
        return [tool]

    # ------------------------------------------------------------------ #
    #  Core rollout logic
    # ------------------------------------------------------------------ #

    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

        logger.info(
            f"[CodeRolloutProcessor] Starting rollout {rollout_id} "
            f"model={input.model_resource.model_name}"
        )

        start_time = time.time()
        tool_call_durations = []

        try:
            extra = input.rollout_extra or {}
            tools = self._build_tools(extra)
            trace_tool(tools)
            tool_node = ToolNode(tools)
            graph = self._build_graph()

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

            system_prompt = self._get_system_prompt(extra)
            messages = [SystemMessage(content=system_prompt)]
            messages.extend(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": tool_node,
                    "tool_call_durations": tool_call_durations,
                },
                metadata={
                    "rollout_id": rollout_id,
                    "model_name": input.model_resource.model_name,
                },
            )

            final_state = await graph.ainvoke({"messages": messages}, config=config)
            final_messages = final_state["messages"]

            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=0.0,
            )

            logger.info(
                f"[CodeRolloutProcessor] Done {rollout_id} | latency={latency}s | "
                f"tool_calls: {len(tool_call_durations)} | "
                f"messages={len(final_messages)}"
            )
            return RolloutOutput(
                agent_output=agent_output,
                status=TaskStatus.SUCCESS,
                error=None,
            )

        except Exception as e:
            logger.error(f"[CodeRolloutProcessor] 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:
        return await self._async_process(input)
