from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal, overload

from ..._constants import HEADER_RUN_ID
from ..._streaming import (
    AsyncChatCompletionStreamManager,
    AsyncRungateStream,
    ChatCompletionStreamManager,
    RungateStream,
)
from ...types.chat.completion import ChatCompletion

if TYPE_CHECKING:
    from ..._client import AsyncRungate, Rungate


def _build_body(
    *,
    model: str,
    messages: list[dict[str, Any]],
    temperature: float | None = None,
    max_tokens: int | None = None,
    top_p: float | None = None,
    tools: list[dict[str, Any]] | None = None,
    tool_choice: str | dict[str, Any] | None = None,
    stop: str | list[str] | None = None,
    stream: bool = False,
    extra_body: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build the request body for a chat completion."""
    body: dict[str, Any] = {"model": model, "messages": messages}
    if temperature is not None:
        body["temperature"] = temperature
    if max_tokens is not None:
        body["max_tokens"] = max_tokens
    if top_p is not None:
        body["top_p"] = top_p
    if tools is not None:
        body["tools"] = tools
    if tool_choice is not None:
        body["tool_choice"] = tool_choice
    if stop is not None:
        body["stop"] = stop
    if stream:
        body["stream"] = True
    if extra_body:
        body.update(extra_body)
    return body


def _build_headers(*, run_id: str | None = None) -> dict[str, str] | None:
    """Build extra headers for a chat completion."""
    if run_id:
        return {HEADER_RUN_ID: run_id}
    return None


class Completions:
    """Chat completions resource — proxies requests through Rungate."""

    def __init__(self, client: Rungate) -> None:
        self._client = client

    @overload
    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: Literal[False] = ...,
        temperature: float | None = ...,
        max_tokens: int | None = ...,
        top_p: float | None = ...,
        tools: list[dict[str, Any]] | None = ...,
        tool_choice: str | dict[str, Any] | None = ...,
        stop: str | list[str] | None = ...,
        run_id: str | None = ...,
        extra_body: dict[str, Any] | None = ...,
    ) -> ChatCompletion: ...

    @overload
    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: Literal[True],
        temperature: float | None = ...,
        max_tokens: int | None = ...,
        top_p: float | None = ...,
        tools: list[dict[str, Any]] | None = ...,
        tool_choice: str | dict[str, Any] | None = ...,
        stop: str | list[str] | None = ...,
        run_id: str | None = ...,
        extra_body: dict[str, Any] | None = ...,
    ) -> RungateStream[dict[str, Any]]: ...

    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        stop: str | list[str] | None = None,
        stream: bool = False,
        run_id: str | None = None,
        extra_body: dict[str, Any] | None = None,
    ) -> ChatCompletion | RungateStream[dict[str, Any]]:
        """Create a chat completion via the Rungate proxy.

        Args:
            model: The model to use (e.g., "gpt-4o", "claude-sonnet-4-20250514").
            messages: The conversation messages.
            stream: Whether to stream the response.
            temperature: Sampling temperature (0-2).
            max_tokens: Maximum tokens to generate.
            top_p: Nucleus sampling parameter.
            tools: Tool definitions for function calling.
            tool_choice: How the model should use tools.
            stop: Stop sequences.
            run_id: Explicit Rungate run ID.
            extra_body: Additional fields to include in the request body.

        Returns:
            ChatCompletion when stream=False, RungateStream when stream=True.
        """
        body = _build_body(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            tools=tools,
            tool_choice=tool_choice,
            stop=stop,
            stream=stream,
            extra_body=extra_body,
        )
        headers = _build_headers(run_id=run_id)

        if stream:
            response = self._client._request_stream("POST", "/v1/chat/completions", json=body, headers=headers)
            return RungateStream(response)

        response = self._client._request("POST", "/v1/chat/completions", json=body, headers=headers)
        data = response.body
        return ChatCompletion(
            id=data.get("id", ""),
            object=data.get("object", "chat.completion"),
            model=data.get("model", model),
            choices=data.get("choices", []),
            usage=data.get("usage"),
            created=data.get("created"),
            system_fingerprint=data.get("system_fingerprint"),
            rungate=response.metadata,
        )

    def stream(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        stop: str | list[str] | None = None,
        run_id: str | None = None,
        extra_body: dict[str, Any] | None = None,
    ) -> ChatCompletionStreamManager:
        """High-level streaming with text extraction and accumulation.

        Must be used as a context manager:
            with client.chat.completions.stream(model="gpt-4o", messages=[...]) as stream:
                for text in stream.text_stream:
                    print(text, end="")
                completion = stream.get_final_completion()
        """
        body = _build_body(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            tools=tools,
            tool_choice=tool_choice,
            stop=stop,
            stream=True,
            extra_body=extra_body,
        )
        headers = _build_headers(run_id=run_id)
        return ChatCompletionStreamManager(self._client, request_json=body, request_headers=headers)


class AsyncCompletions:
    """Async chat completions resource."""

    def __init__(self, client: AsyncRungate) -> None:
        self._client = client

    @overload
    async def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: Literal[False] = ...,
        temperature: float | None = ...,
        max_tokens: int | None = ...,
        top_p: float | None = ...,
        tools: list[dict[str, Any]] | None = ...,
        tool_choice: str | dict[str, Any] | None = ...,
        stop: str | list[str] | None = ...,
        run_id: str | None = ...,
        extra_body: dict[str, Any] | None = ...,
    ) -> ChatCompletion: ...

    @overload
    async def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: Literal[True],
        temperature: float | None = ...,
        max_tokens: int | None = ...,
        top_p: float | None = ...,
        tools: list[dict[str, Any]] | None = ...,
        tool_choice: str | dict[str, Any] | None = ...,
        stop: str | list[str] | None = ...,
        run_id: str | None = ...,
        extra_body: dict[str, Any] | None = ...,
    ) -> AsyncRungateStream[dict[str, Any]]: ...

    async def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        stop: str | list[str] | None = None,
        stream: bool = False,
        run_id: str | None = None,
        extra_body: dict[str, Any] | None = None,
    ) -> ChatCompletion | AsyncRungateStream[dict[str, Any]]:
        """Create a chat completion via the Rungate proxy (async)."""
        body = _build_body(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            tools=tools,
            tool_choice=tool_choice,
            stop=stop,
            stream=stream,
            extra_body=extra_body,
        )
        headers = _build_headers(run_id=run_id)

        if stream:
            response = await self._client._request_stream("POST", "/v1/chat/completions", json=body, headers=headers)
            return AsyncRungateStream(response)

        response = await self._client._request("POST", "/v1/chat/completions", json=body, headers=headers)
        data = response.body
        return ChatCompletion(
            id=data.get("id", ""),
            object=data.get("object", "chat.completion"),
            model=data.get("model", model),
            choices=data.get("choices", []),
            usage=data.get("usage"),
            created=data.get("created"),
            system_fingerprint=data.get("system_fingerprint"),
            rungate=response.metadata,
        )

    def stream(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        top_p: float | None = None,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        stop: str | list[str] | None = None,
        run_id: str | None = None,
        extra_body: dict[str, Any] | None = None,
    ) -> AsyncChatCompletionStreamManager:
        """High-level async streaming."""
        body = _build_body(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            tools=tools,
            tool_choice=tool_choice,
            stop=stop,
            stream=True,
            extra_body=extra_body,
        )
        headers = _build_headers(run_id=run_id)
        return AsyncChatCompletionStreamManager(self._client, request_json=body, request_headers=headers)
