"""Helpers for OpenAI-compatible tool calling.

Many providers (OpenAI, Groq, Together, Mistral, Ollama, OpenAI-compatible
endpoints, and Azure OpenAI for the chat payload) use broadly the same
`chat/completions` shapes for tool calling:
- `tool_calls` in the assistant message
- tool calls contain `{id, function: {name, arguments}}`
- tool results are supplied as `{role: "tool", tool_call_id, content}`
"""

from __future__ import annotations

import json
from typing import Any

from mcpbridge.adapters.base import ToolCall, ToolResult


def extract_tool_calls_openai(response: Any) -> list[ToolCall]:
    """Extract ToolCall objects from OpenAI-compatible responses."""
    msg = _extract_message(response)
    tool_calls = msg.get("tool_calls") or []
    out: list[ToolCall] = []
    for call in tool_calls:
        if not isinstance(call, dict):
            continue
        call_id = str(call.get("id") or "")
        func = call.get("function") or {}
        name = str(func.get("name") or "")
        args = func.get("arguments") or {}
        if isinstance(args, str):
            try:
                args = json.loads(args) if args else {}
            except Exception:
                args = {}
        if not isinstance(args, dict):
            args = {}
        if call_id and name:
            out.append(ToolCall(id=call_id, name=name, arguments=args))
    return out


def extract_text_openai(response: Any) -> str:
    """Extract final assistant text from OpenAI-compatible responses."""
    msg = _extract_message(response)
    content = msg.get("content")
    if content is None:
        return ""
    return str(content)


def is_done_openai(response: Any) -> bool:
    """Done when there are no tool calls."""
    msg = _extract_message(response)
    tool_calls = msg.get("tool_calls") or []
    return len(tool_calls) == 0


def append_tool_results_openai(
    messages: list[dict[str, Any]],
    response: Any,
    tool_results: list[ToolResult],
) -> list[dict[str, Any]]:
    """Append assistant tool-call message and tool result messages."""
    msg = _extract_message(response)
    tool_calls = msg.get("tool_calls") or []
    content = msg.get("content")

    # Append the assistant message that triggered the tool calls.
    messages.append(
        {
            "role": "assistant",
            "content": "" if content is None else content,
            "tool_calls": tool_calls,
        }
    )

    # Append tool outputs.
    for tr in tool_results:
        messages.append({"role": "tool", "tool_call_id": tr.tool_call_id, "content": tr.content})
    return messages


def _extract_message(response: Any) -> dict[str, Any]:
    # Standard OpenAI: {choices:[{message:{...}}]}
    if isinstance(response, dict):
        if "choices" in response and isinstance(response["choices"], list) and response["choices"]:
            choice0 = response["choices"][0]
            if isinstance(choice0, dict):
                message = choice0.get("message") or {}
                if isinstance(message, dict):
                    return message
        # Ollama/others: {message:{...}} or {output:{message:{...}}}
        if "message" in response and isinstance(response["message"], dict):
            return response["message"]
        if "output" in response and isinstance(response["output"], dict) and "message" in response["output"]:
            out_msg = response["output"].get("message")
            if isinstance(out_msg, dict):
                return out_msg
    raise ValueError("Unsupported OpenAI-compatible response shape")

