"""AgentRunner implementation backed by the OpenAI Agents SDK.

Each registered MCP server is connected as a stdio MCP server so the OpenAI
agent can call IoT / FMSR / TSFM / utilities tools directly via MCP.

Usage::

    import anyio
    from agent.openai_agent import OpenAIAgentRunner

    runner = OpenAIAgentRunner(model="litellm_proxy/azure/gpt-5.4")
    result = anyio.run(runner.run, "What sensors are on Chiller 6?")
    print(result.answer)
"""

from __future__ import annotations

import datetime as _dt
import json
import logging
import os
import time
from contextlib import AsyncExitStack
from pathlib import Path

from openai import AsyncOpenAI

from agents import Agent, ModelProvider, OpenAIChatCompletionsModel, RunConfig, Runner, set_tracing_disabled
from agents.mcp import MCPServerStdio

from observability import agent_run_span, persist_trajectory

from .._litellm import LITELLM_PREFIX, resolve_model
from .._prompts import AGENT_SYSTEM_PROMPT
from ..models import AgentResult, ToolCall, Trajectory, TurnRecord
from ..runner import AgentRunner

_log = logging.getLogger(__name__)

_DEFAULT_MODEL = "litellm_proxy/azure/gpt-5.4"


def _build_run_config(model_id: str) -> RunConfig | None:
    """Build a RunConfig with a LiteLLM model provider when needed.

    When *model_id* starts with ``litellm_proxy/``, creates an
    :class:`AsyncOpenAI` client pointing at the LiteLLM proxy (using
    ``LITELLM_BASE_URL`` and ``LITELLM_API_KEY``) and wraps it in
    :class:`OpenAIChatCompletionsModel`.

    Returns ``None`` for direct OpenAI API usage.
    """
    if not model_id.startswith(LITELLM_PREFIX):
        return None

    base_url = os.environ.get("LITELLM_BASE_URL")
    api_key = os.environ.get("LITELLM_API_KEY")
    if not base_url or not api_key:
        raise ValueError(
            "LITELLM_BASE_URL and LITELLM_API_KEY must be set "
            f"when using {LITELLM_PREFIX!r} model prefix"
        )

    resolved = resolve_model(model_id)
    client = AsyncOpenAI(base_url=base_url, api_key=api_key)
    set_tracing_disabled(disabled=True)

    class _LiteLLMModelProvider(ModelProvider):
        def get_model(self, model_name: str | None):
            return OpenAIChatCompletionsModel(
                model=model_name or resolved,
                openai_client=client,
            )

    return RunConfig(model_provider=_LiteLLMModelProvider())


def _build_mcp_servers(
    server_paths: dict[str, Path | str],
) -> list[MCPServerStdio]:
    """Convert server_paths entries into MCPServerStdio instances.

    Entry-point names (str without path separators) become
    ``MCPServerStdio(command="uv", args=["run", name])``.
    Path objects become ``MCPServerStdio(command="uv", args=["run", str(path)])``.
    """
    servers: list[MCPServerStdio] = []
    for name, spec in server_paths.items():
        cmd_arg = str(spec) if isinstance(spec, Path) else spec
        servers.append(
            MCPServerStdio(
                name=name,
                params={
                    "command": "uv",
                    "args": ["run", cmd_arg],
                },
                cache_tools_list=True,
            )
        )
    return servers


def _build_trajectory(result) -> Trajectory:
    """Extract a Trajectory from a Runner.run result.

    Walks ``result.new_items`` to collect text messages, tool calls, and
    tool outputs.  Token usage is pulled from ``result.raw_responses``.
    """
    trajectory = Trajectory()
    turn_index = 0
    text_parts: list[str] = []
    tool_calls: list[ToolCall] = []

    def _flush() -> None:
        nonlocal text_parts, tool_calls, turn_index
        if not text_parts and not tool_calls:
            return
        trajectory.turns.append(
            TurnRecord(
                index=turn_index,
                text="".join(text_parts),
                tool_calls=list(tool_calls),
            )
        )
        turn_index += 1
        text_parts = []
        tool_calls = []

    for item in result.new_items:
        item_type = getattr(item, "type", "")
        if item_type == "message_output_item":
            # Flush any pending tool calls from previous turn
            _flush()
            raw = getattr(item, "raw_item", None)
            if raw:
                content = getattr(raw, "content", None) or []
                for part in content:
                    if hasattr(part, "text"):
                        text_parts.append(part.text)
        elif item_type == "tool_call_item":
            raw = getattr(item, "raw_item", None)
            if raw:
                tc_name = getattr(raw, "name", "") or ""
                tc_id = getattr(raw, "call_id", "") or getattr(raw, "id", "") or ""
                tc_args = getattr(raw, "arguments", "{}") or "{}"
                try:
                    tc_input = json.loads(tc_args) if isinstance(tc_args, str) else tc_args
                except (json.JSONDecodeError, TypeError):
                    tc_input = {"raw": tc_args}
                tool_calls.append(ToolCall(name=tc_name, input=tc_input, id=tc_id))
        elif item_type == "tool_call_output_item":
            output = getattr(item, "output", None)
            # Attach output to the last matching tool call
            if tool_calls:
                tool_calls[-1].output = output

    # Flush remaining
    _flush()

    # Distribute token usage from raw_responses across turns
    raw_responses = getattr(result, "raw_responses", []) or []
    for i, resp in enumerate(raw_responses):
        usage = getattr(resp, "usage", None)
        if usage and i < len(trajectory.turns):
            trajectory.turns[i].input_tokens = getattr(usage, "input_tokens", 0) or 0
            trajectory.turns[i].output_tokens = getattr(usage, "output_tokens", 0) or 0

    return trajectory


class OpenAIAgentRunner(AgentRunner):
    """Agent runner that delegates to the OpenAI Agents SDK agentic loop.

    The SDK handles tool discovery, invocation, and multi-turn conversation
    against the registered MCP servers.

    Routes all requests through a LiteLLM proxy via the ``litellm_proxy/``
    model ID prefix (requires ``LITELLM_BASE_URL`` and ``LITELLM_API_KEY``).

    Args:
        llm: Unused — OpenAIAgentRunner uses the OpenAI Agents SDK directly.
             Accepted for interface compatibility with ``AgentRunner``.
        server_paths: MCP server specs identical to ``PlanExecuteRunner``.
                      Defaults to all registered servers.
        model: LiteLLM model string with ``litellm_proxy/`` prefix
               (default: ``litellm_proxy/azure/gpt-5.4``).
        max_turns: Maximum agentic loop turns (default: 30).
    """

    def __init__(
        self,
        llm=None,
        server_paths: dict[str, Path | str] | None = None,
        model: str = _DEFAULT_MODEL,
        max_turns: int = 30,
    ) -> None:
        super().__init__(llm, server_paths)
        self._model_id = model
        self._model = resolve_model(model)
        self._run_config = _build_run_config(model)
        self._max_turns = max_turns

    async def run(self, question: str) -> AgentResult:
        """Run the OpenAI Agents SDK loop for *question*.

        Args:
            question: Natural-language question to answer.

        Returns:
            AgentResult with the final answer and full execution trajectory.
        """
        with agent_run_span(
            "openai-agent", model=self._model_id, question=question
        ) as span:
            run_started = time.perf_counter()
            started_at = _dt.datetime.now(_dt.UTC).isoformat()
            mcp_servers = _build_mcp_servers(self._server_paths)

            # AsyncExitStack enters every server and closes them in LIFO order
            # on exit (success or exception).
            async with AsyncExitStack() as stack:
                active_servers = [
                    await stack.enter_async_context(s) for s in mcp_servers
                ]
                agent = Agent(
                    name="AssetOps Assistant",
                    instructions=AGENT_SYSTEM_PROMPT,
                    mcp_servers=active_servers,
                    model=self._model,
                )

                _log.info(
                    "OpenAIAgentRunner: starting query (model=%s, servers=%d)",
                    self._model,
                    len(active_servers),
                )

                run_kwargs: dict = dict(max_turns=self._max_turns)
                if self._run_config is not None:
                    run_kwargs["run_config"] = self._run_config

                result = await Runner.run(
                    agent,
                    question,
                    **run_kwargs,
                )

                answer = result.final_output or ""
                trajectory = _build_trajectory(result)
                trajectory.started_at = started_at

                _log.info(
                    "OpenAIAgentRunner: done (turns=%d, input_tokens=%d, "
                    "output_tokens=%d)",
                    len(trajectory.turns),
                    trajectory.total_input_tokens,
                    trajectory.total_output_tokens,
                )

                span.set_attribute("agent.answer.length", len(answer))
                span.set_attribute("gen_ai.usage.input_tokens", trajectory.total_input_tokens)
                span.set_attribute("gen_ai.usage.output_tokens", trajectory.total_output_tokens)
                span.set_attribute("agent.turns", len(trajectory.turns))
                span.set_attribute("agent.tool_calls", len(trajectory.all_tool_calls))
                span.set_attribute(
                    "agent.duration_ms", (time.perf_counter() - run_started) * 1000
                )
                persist_trajectory(
                    runner_name="openai-agent",
                    model=self._model_id,
                    question=question,
                    answer=answer,
                    trajectory=trajectory,
                )
                return AgentResult(
                    question=question,
                    answer=answer,
                    trajectory=trajectory,
                )


