from __future__ import annotations

import json
import os
from collections.abc import AsyncGenerator, Sequence
from typing import Any, override

import openai
import pydantic

from .. import core


def _tools_to_openai(tools: Sequence[core.tools.ToolLike]) -> list[dict[str, Any]]:
    """Convert internal Tool objects to OpenAI tool schema format."""
    return [
        {
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.param_schema,
            },
        }
        for tool in tools
    ]


def _messages_to_openai(messages: list[core.messages.Message]) -> list[dict[str, Any]]:
    """Convert internal messages to OpenAI API format.

    Converts to the OpenAI wire format:

    - ``tool_calls`` on assistant messages
    - tool results as separate ``role: "tool"`` messages

    The Vercel AI Gateway preserves reasoning details across interactions,
    normalizing formats from different providers.

    See: https://vercel.com/docs/ai-gateway/openai-compat/advanced
    """
    result: list[dict[str, Any]] = []
    for msg in messages:
        if msg.role == "assistant":
            content = ""
            reasoning = ""
            tool_calls = []
            tool_results = []

            for part in msg.parts:
                if isinstance(part, core.messages.ReasoningPart):
                    reasoning += part.text
                elif isinstance(part, core.messages.TextPart):
                    content += part.text
                elif isinstance(part, core.messages.ToolPart):
                    tool_calls.append(
                        {
                            "id": part.tool_call_id,
                            "type": "function",
                            "function": {
                                "name": part.tool_name,
                                "arguments": part.tool_args,
                            },
                        }
                    )
                    if part.status in ("result", "error"):
                        tool_results.append(
                            {
                                "role": "tool",
                                "tool_call_id": part.tool_call_id,
                                "content": str(part.result)
                                if part.result is not None
                                else "",
                            }
                        )

            entry: dict[str, Any] = {"role": "assistant"}
            if content:
                entry["content"] = content
            if reasoning:
                entry["reasoning"] = reasoning
            if tool_calls:
                entry["tool_calls"] = tool_calls
            result.append(entry)

            # Emit tool results as separate messages (OpenAI API format)
            result.extend(tool_results)
        else:
            # User/system messages
            content = "".join(
                p.text for p in msg.parts if isinstance(p, core.messages.TextPart)
            )
            result.append({"role": msg.role, "content": content})
    return result


class OpenAIModel(core.llm.LanguageModel):
    """OpenAI adapter with reasoning/thinking support via Vercel AI Gateway.

    Supports reasoning for models like GPT 5.x, o-series, and Claude via gateway.
    Uses the Vercel AI Gateway's unified reasoning API format.

    See: https://vercel.com/docs/ai-gateway/openai-compat/advanced
    """

    def __init__(
        self,
        model: str = "gpt-4o",
        base_url: str | None = None,
        api_key: str | None = None,
        thinking: bool = False,
        budget_tokens: int | None = None,
        reasoning_effort: str | None = None,
    ) -> None:
        """Initialize OpenAI model adapter.

        Args:
            model: Model identifier
                (e.g., 'openai/gpt-5.2', 'anthropic/claude-sonnet-4.5')
            base_url: API base URL
                (e.g., 'https://ai-gateway.vercel.sh/v1')
            api_key: API key for authentication
            thinking: Enable reasoning/thinking output
            budget_tokens: Max tokens for reasoning
                (mutually exclusive with reasoning_effort)
            reasoning_effort: Effort level — 'none', 'minimal',
                'low', 'medium', 'high', 'xhigh'
                (mutually exclusive with budget_tokens)
        """
        self._model = model
        self._thinking = thinking
        self._budget_tokens = budget_tokens
        self._reasoning_effort = reasoning_effort
        resolved_key = api_key or os.environ.get("OPENAI_API_KEY") or ""
        self._client = openai.AsyncOpenAI(base_url=base_url, api_key=resolved_key)

    async def stream_events(
        self,
        messages: list[core.messages.Message],
        tools: Sequence[core.tools.ToolLike] | None = None,
        output_type: type[pydantic.BaseModel] | None = None,
    ) -> AsyncGenerator[core.llm.StreamEvent]:
        """Yield raw stream events from OpenAI API."""
        openai_messages = _messages_to_openai(messages)
        openai_tools = _tools_to_openai(tools) if tools else None

        kwargs: dict[str, Any] = {
            "model": self._model,
            "messages": openai_messages,
            "stream": True,
        }
        if openai_tools:
            kwargs["tools"] = openai_tools
        kwargs["stream_options"] = {"include_usage": True}

        if output_type is not None:
            from openai.lib._pydantic import to_strict_json_schema

            kwargs["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": output_type.__name__,
                    "schema": to_strict_json_schema(output_type),
                    "strict": True,
                },
            }

        # Enable reasoning/thinking via Vercel AI Gateway's unified format
        # See: https://vercel.com/docs/ai-gateway/openai-compat/advanced
        if self._thinking:
            reasoning_config: dict[str, Any] = {"enabled": True}
            # Use budget_tokens OR reasoning_effort (mutually exclusive per docs)
            if self._budget_tokens is not None:
                reasoning_config["max_tokens"] = self._budget_tokens
            elif self._reasoning_effort is not None:
                reasoning_config["effort"] = self._reasoning_effort
            kwargs["extra_body"] = {"reasoning": reasoning_config}

        stream = await self._client.chat.completions.create(**kwargs)

        # Track active blocks for Start/End events
        text_started = False
        reasoning_started = False
        tool_calls: dict[int, dict[str, Any]] = {}  # index -> {id, name, started}
        finish_reason: str | None = None
        usage: core.messages.Usage | None = None

        async for chunk in stream:
            # Extract usage from any chunk that carries it (typically the final
            # chunk when stream_options.include_usage is True).
            if chunk.usage is not None:
                raw = chunk.usage.model_dump(exclude_none=True)
                # Extract optional breakdowns
                reasoning_tokens: int | None = None
                cache_read: int | None = None
                completion_details = getattr(
                    chunk.usage, "completion_tokens_details", None
                )
                if completion_details:
                    reasoning_tokens = getattr(
                        completion_details, "reasoning_tokens", None
                    )
                prompt_details = getattr(chunk.usage, "prompt_tokens_details", None)
                if prompt_details:
                    cache_read = getattr(prompt_details, "cached_tokens", None)
                usage = core.messages.Usage(
                    input_tokens=chunk.usage.prompt_tokens or 0,
                    output_tokens=chunk.usage.completion_tokens or 0,
                    reasoning_tokens=reasoning_tokens,
                    cache_read_tokens=cache_read,
                    raw=raw,
                )

            if not chunk.choices:
                continue

            choice = chunk.choices[0]
            delta = choice.delta

            # Handle reasoning/thinking content via Vercel AI Gateway
            # The gateway may return reasoning in different ways:
            # 1. As a direct attribute (if SDK supports it)
            # 2. In model_extra (Pydantic v2 extra fields)
            reasoning_value = None
            if hasattr(delta, "reasoning") and delta.reasoning:
                reasoning_value = delta.reasoning
            elif hasattr(delta, "model_extra") and delta.model_extra:
                reasoning_value = delta.model_extra.get("reasoning")

            if reasoning_value:
                if not reasoning_started:
                    reasoning_started = True
                    yield core.llm.ReasoningStart(block_id="reasoning")
                yield core.llm.ReasoningDelta(
                    block_id="reasoning", delta=reasoning_value
                )

            if delta.content:
                # Close reasoning block when text starts (reasoning precedes text)
                if reasoning_started:
                    yield core.llm.ReasoningEnd(block_id="reasoning")
                    reasoning_started = False

                if not text_started:
                    text_started = True
                    yield core.llm.TextStart(block_id="text")
                yield core.llm.TextDelta(block_id="text", delta=delta.content)

            if delta.tool_calls:
                for tc in delta.tool_calls:
                    idx = tc.index
                    if idx not in tool_calls:
                        tool_calls[idx] = {"id": tc.id, "name": None, "started": False}
                    if tc.id:
                        tool_calls[idx]["id"] = tc.id
                    if tc.function:
                        if tc.function.name:
                            tool_calls[idx]["name"] = tc.function.name
                        if tc.function.arguments:
                            tool_id = tool_calls[idx]["id"]
                            tool_name = tool_calls[idx]["name"] or ""

                            # Emit start if not started
                            if not tool_calls[idx]["started"] and tool_id:
                                tool_calls[idx]["started"] = True
                                yield core.llm.ToolStart(
                                    tool_call_id=tool_id, tool_name=tool_name
                                )

                            if tool_id:
                                yield core.llm.ToolArgsDelta(
                                    tool_call_id=tool_id, delta=tc.function.arguments
                                )

            if choice.finish_reason is not None:
                finish_reason = choice.finish_reason
                # Close any open blocks
                if reasoning_started:
                    yield core.llm.ReasoningEnd(block_id="reasoning")
                if text_started:
                    yield core.llm.TextEnd(block_id="text")
                for tc in tool_calls.values():
                    if tc["started"] and tc["id"]:
                        yield core.llm.ToolEnd(tool_call_id=tc["id"])

                # Don't return yet — the usage chunk may arrive after
                # finish_reason. We'll emit MessageDone after the loop.

        yield core.llm.MessageDone(finish_reason=finish_reason, usage=usage)

    @override
    async def stream(
        self,
        messages: list[core.messages.Message],
        tools: Sequence[core.tools.ToolLike] | None = None,
        output_type: type[pydantic.BaseModel] | None = None,
    ) -> AsyncGenerator[core.messages.Message]:
        """Stream Messages (uses StreamHandler internally)."""
        handler = core.llm.StreamHandler()
        msg: core.messages.Message | None = None
        async for event in self.stream_events(messages, tools, output_type=output_type):
            msg = handler.handle_event(event)
            yield msg

        # After stream completes, validate and attach structured output part
        if output_type is not None and msg is not None and msg.text:
            data = json.loads(msg.text)
            output_type.model_validate(data)  # fail fast on bad data
            part = core.messages.StructuredOutputPart(
                data=data,
                output_type_name=f"{output_type.__module__}.{output_type.__qualname__}",
            )
            msg = msg.model_copy()
            msg.parts = [*msg.parts, part]
            yield msg
