"""ART rollout function with CWSandbox execution.

This module implements the core rollout function that:
1. Uses OpenAI Responses API with tool calling
2. Executes tools (execute_code, submit_solution) in CWSandboxes
3. Builds an ART Trajectory with the conversation history
4. Computes binary reward (1.0 if tests pass, 0.0 otherwise)

The rollout loop allows max 20 tool calls per problem, with errors
returned as tool responses so the agent can observe and retry.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, cast

from openai import AsyncOpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from openai.types.chat.chat_completion_message_function_tool_call import Function
from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
)
from openai.types.responses import Response, ResponseFunctionToolCall
from tools import (
    EXECUTE_CODE_NAME,
    ROLLOUT_TOOLS,
    ROLLOUT_TOOLS_RESPONSES,
    SUBMIT_SOLUTION_NAME,
    ExecuteCodeArgs,
    SubmitSolutionArgs,
)

import art

if TYPE_CHECKING:
    from cwsandbox import Sandbox

logger = logging.getLogger(__name__)

MAX_TOOL_CALLS = 20
MAX_OUTPUT_CHARS = 8192
EXECUTION_TIMEOUT_SECONDS = 60.0


@dataclass
class Problem:
    """A coding problem for the agent to solve.

    Attributes:
        task_id: Unique identifier for the problem
        prompt: Problem description/instructions
        test_code: Code that tests the solution (runs after solution code)
        test_imports: Optional imports needed for test code
    """

    task_id: str
    prompt: str
    test_code: str
    test_imports: str = ""


@dataclass
class RolloutConfig:
    """Configuration for rollout execution.

    Attributes:
        model: Model identifier (e.g., "gpt-4o", "gpt-4o-mini")
        base_url: OpenAI-compatible API base URL
        api_key: API key for authentication
        max_attempts: Maximum tool calls before giving up
    """

    model: str = "gpt-5.1-codex-mini"
    base_url: str | None = None
    api_key: str | None = None
    max_attempts: int = MAX_TOOL_CALLS
    execution_timeout: float = EXECUTION_TIMEOUT_SECONDS


_PROMPT_DIR = Path(__file__).parent / "prompts"
_SYSTEM_PROMPT_TEMPLATE: str | None = None


def _load_system_prompt_template() -> str:
    """Load the system prompt template from file (cached)."""
    global _SYSTEM_PROMPT_TEMPLATE
    if _SYSTEM_PROMPT_TEMPLATE is None:
        prompt_path = _PROMPT_DIR / "system_prompt.md"
        _SYSTEM_PROMPT_TEMPLATE = prompt_path.read_text()
    return _SYSTEM_PROMPT_TEMPLATE


def _extract_function_name(test_code: str) -> str | None:
    """Extract the expected function name from test assertions."""
    import re

    match = re.search(r"assert\s+(\w+)\s*\(", test_code)
    return match.group(1) if match else None


def build_system_message(problem: Problem, max_tool_calls: int = MAX_TOOL_CALLS) -> str:
    """Build the system message with problem context."""
    template = _load_system_prompt_template()
    function_name = _extract_function_name(problem.test_code)
    function_hint = f"\n\nREQUIRED FUNCTION NAME: `{function_name}`" if function_name else ""

    return template.format(
        max_tool_calls=max_tool_calls,
        function_hint=function_hint,
        problem_prompt=problem.prompt,
        test_code=problem.test_code,
    )


async def execute_code_in_sandbox(
    sandbox: Sandbox,
    code: str,
    timeout: float = EXECUTION_TIMEOUT_SECONDS,
) -> str:
    """Execute code in sandbox and return output or error message.

    Args:
        sandbox: CWSandbox instance
        code: Python code to execute
        timeout: Maximum execution time in seconds

    Returns:
        Output string: stdout on success, error message on failure
    """
    try:
        result = await sandbox.exec(["python", "-c", code], timeout_seconds=timeout)

        if result.returncode == 0:
            output = result.stdout.strip() if result.stdout else "(no output)"
        else:
            error = result.stderr.strip() if result.stderr else "(unknown error)"
            output = f"Error (exit code {result.returncode}):\n{error}"

    except Exception as e:
        logger.debug("Sandbox execution failed: %s: %s", type(e).__name__, e)
        output = "Error: execution failed"

    if len(output) > MAX_OUTPUT_CHARS:
        output = "[truncated]\n" + output[-MAX_OUTPUT_CHARS:]
    return output


async def run_tests_in_sandbox(
    sandbox: Sandbox,
    solution_code: str,
    test_code: str,
    test_imports: str = "",
    timeout: float = EXECUTION_TIMEOUT_SECONDS,
) -> tuple[bool, str]:
    """Run solution against test cases in sandbox.

    Args:
        sandbox: CWSandbox instance
        solution_code: The solution code to test
        test_code: Test code that validates the solution
        test_imports: Optional imports for test code
        timeout: Maximum execution time in seconds

    Returns:
        Tuple of (passed: bool, output: str)
    """
    full_code = f"{test_imports}\n\n{solution_code}\n\n{test_code}"

    try:
        result = await sandbox.exec(["python", "-c", full_code], timeout_seconds=timeout)

        if result.returncode == 0:
            output = result.stdout.strip() if result.stdout else "All tests passed"
            passed = True
        else:
            error = result.stderr.strip() if result.stderr else "(unknown error)"
            output = f"Tests failed (exit code {result.returncode}):\n{error}"
            passed = False

    except Exception as e:
        logger.debug("Sandbox test execution failed: %s: %s", type(e).__name__, e)
        output = "Error: execution failed"
        passed = False

    if len(output) > MAX_OUTPUT_CHARS:
        output = "[truncated]\n" + output[-MAX_OUTPUT_CHARS:]
    return passed, output


def _normalize_output_for_input(response: Response) -> list[dict[str, Any]]:
    """Convert response output items to valid input items for re-feeding.

    The Responses API output items need minor transformation to be valid input:
    - ResponseOutputMessage: extract content and convert to message format
    - ResponseFunctionToolCall: pass through (already valid as input)
    - ResponseReasoningItem: include for reasoning models (required before function_call)
    - Other items (refusal, error): convert to message format

    Args:
        response: The API response object

    Returns:
        List of input items ready to append to conversation
    """
    input_items: list[dict[str, Any]] = []

    for item in response.output:
        if item.type == "message":
            content_parts = []
            for part in item.content:
                if part.type == "output_text":
                    content_parts.append({"type": "text", "text": part.text})
                elif part.type == "refusal":
                    content_parts.append({"type": "text", "text": f"[Refusal: {part.refusal}]"})
            if content_parts:
                input_items.append(
                    {
                        "type": "message",
                        "role": "assistant",
                        "content": content_parts,
                    }
                )
        elif item.type == "reasoning":
            input_items.append(
                {
                    "type": "reasoning",
                    "id": item.id,
                    "summary": getattr(item, "summary", None),
                }
            )
        elif item.type == "function_call":
            input_items.append(
                {
                    "type": "function_call",
                    "id": item.id,
                    "call_id": item.call_id,
                    "name": item.name,
                    "arguments": item.arguments,
                }
            )

    return input_items


def _response_output_to_choice(response: Response) -> Choice:
    """Convert Responses API output to a Chat Completions Choice object for ART.

    ART expects Choice objects with ChatCompletionMessage for trajectory building.
    This function constructs the equivalent Choice from Responses API output.

    Args:
        response: The API response object

    Returns:
        Choice object compatible with ART Trajectory
    """
    text_content = response.output_text or ""

    tool_calls: list[ChatCompletionMessageToolCall] = []
    for item in response.output:
        if item.type == "function_call":
            tool_calls.append(
                ChatCompletionMessageToolCall(
                    id=item.call_id,
                    type="function",
                    function=Function(
                        name=item.name,
                        arguments=item.arguments,
                    ),
                )
            )

    message = ChatCompletionMessage(
        role="assistant",
        content=text_content if text_content else None,
        tool_calls=cast(Any, tool_calls) if tool_calls else None,
    )

    finish_reason: Literal["stop", "tool_calls"] = "tool_calls" if tool_calls else "stop"

    return Choice(
        index=0,
        message=message,
        finish_reason=finish_reason,
    )


async def handle_tool_call(
    tool_call: ResponseFunctionToolCall,
    sandbox: Sandbox,
    problem: Problem,
    config: RolloutConfig,
) -> tuple[str, bool, bool]:
    """Handle a single tool call.

    Args:
        tool_call: The tool call to handle
        sandbox: CWSandbox instance
        problem: The coding problem
        config: Rollout configuration

    Returns:
        Tuple of (result_content, is_submission, passed)
    """
    name = tool_call.name
    arguments = tool_call.arguments
    if not isinstance(arguments, str):
        return "Error: Tool arguments must be a string", False, False
    try:
        args = json.loads(arguments)
    except json.JSONDecodeError as e:
        return f"Error: Invalid JSON arguments: {e}", False, False

    if name == EXECUTE_CODE_NAME:
        typed_args: ExecuteCodeArgs = args
        try:
            code = typed_args["code"]
        except (KeyError, TypeError):
            return "Error: Missing required argument 'code'", False, False
        result = await execute_code_in_sandbox(sandbox, code, config.execution_timeout)
        return result, False, False

    elif name == SUBMIT_SOLUTION_NAME:
        typed_args_submit: SubmitSolutionArgs = args
        try:
            code = typed_args_submit["code"]
        except (KeyError, TypeError):
            return "Error: Missing required argument 'code'", False, False
        passed, result = await run_tests_in_sandbox(
            sandbox,
            code,
            problem.test_code,
            problem.test_imports,
            config.execution_timeout,
        )
        return result, True, passed

    else:
        return f"Error: Unknown tool '{name}'", False, False


async def rollout(
    problem: Problem,
    sandbox: Sandbox,
    config: RolloutConfig | None = None,
    client: AsyncOpenAI | None = None,
) -> art.Trajectory:
    """Execute a single rollout for a coding problem.

    This function runs an agent loop where the model can use tools
    to explore and test its solution before submitting.

    Args:
        problem: The coding problem to solve
        sandbox: CWSandbox instance (must be running)
        config: Optional rollout configuration
        client: Optional AsyncOpenAI client (reuses connections if provided)

    Returns:
        ART Trajectory with conversation history and reward
    """
    if config is None:
        config = RolloutConfig()

    if client is None:
        client = AsyncOpenAI(
            base_url=config.base_url,
            api_key=config.api_key,
        )

    system_msg = build_system_message(problem, max_tool_calls=config.max_attempts)

    input_items: list[dict[str, Any]] = [
        {"type": "message", "role": "system", "content": system_msg},
        {"type": "message", "role": "user", "content": "Please solve this problem."},
    ]

    messages_and_choices: list[dict[str, Any] | Choice] = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": "Please solve this problem."},
    ]
    tool_call_count = 0
    reward = 0.0
    submitted = False

    while tool_call_count < config.max_attempts and not submitted:
        response = await client.responses.create(
            model=config.model,
            input=cast(Any, input_items),
            tools=ROLLOUT_TOOLS_RESPONSES,
            tool_choice="auto",
            parallel_tool_calls=False,
        )

        if response.status != "completed":
            error_msg = response.error.message if response.error else "Unknown error"
            raise RuntimeError(f"Response not completed: {response.status} - {error_msg}")

        choice = _response_output_to_choice(response)
        messages_and_choices.append(choice)

        function_calls = [item for item in response.output if item.type == "function_call"]

        if function_calls:
            normalized_items = _normalize_output_for_input(response)
            input_items.extend(normalized_items)

            for tool_call in function_calls:
                tool_call_count += 1
                result_content, is_submission, passed = await handle_tool_call(
                    tool_call, sandbox, problem, config
                )

                tool_output_item: dict[str, Any] = {
                    "type": "function_call_output",
                    "call_id": tool_call.call_id,
                    "output": result_content,
                }
                input_items.append(tool_output_item)

                tool_response_msg: dict[str, Any] = {
                    "role": "tool",
                    "tool_call_id": tool_call.call_id,
                    "content": result_content,
                }
                messages_and_choices.append(tool_response_msg)

                if is_submission:
                    submitted = True
                    reward = 1.0 if passed else 0.0
                    break

                if tool_call_count >= config.max_attempts:
                    break
        else:
            break

    trajectory = art.Trajectory(
        messages_and_choices=messages_and_choices,
        tools=ROLLOUT_TOOLS,
        reward=reward,
        metadata={
            "task_id": problem.task_id,
            "tool_calls": tool_call_count,
            "submitted": submitted,
        },
    )

    return trajectory.finish()
