from __future__ import annotations

from typing import Any, Iterator, AsyncIterator, Union, overload

import httpx

from .._types import ChatCompletion, ChatCompletionChunk
from .._streaming import Stream, AsyncStream


class Completions:
    """Sync chat completions: client.chat.completions.create(...)"""

    def __init__(self, client: object) -> None:
        self._client = client  # RightNow (sync)

    @overload
    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = ...,
        **kwargs: Any,
    ) -> ChatCompletion: ...

    @overload
    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = ...,
        **kwargs: Any,
    ) -> Stream: ...

    def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = False,
        **kwargs: Any,
    ) -> ChatCompletion | Stream:
        payload: dict[str, Any] = {"model": model, "messages": messages, "stream": stream, **kwargs}

        if stream:
            cm = self._client._http.stream(  # type: ignore[attr-defined]
                "POST",
                self._client._build_url("/v1/chat/completions"),  # type: ignore[attr-defined]
                json=payload,
                headers=self._client._headers,  # type: ignore[attr-defined]
                timeout=self._client.timeout,  # type: ignore[attr-defined]
            )
            raw = cm.__enter__()
            self._client._raise_for_status(raw)  # type: ignore[attr-defined]
            return Stream(raw, context_manager=cm)

        resp = self._client._http.post(  # type: ignore[attr-defined]
            self._client._build_url("/v1/chat/completions"),  # type: ignore[attr-defined]
            json=payload,
            headers=self._client._headers,  # type: ignore[attr-defined]
            timeout=self._client.timeout,  # type: ignore[attr-defined]
        )
        self._client._raise_for_status(resp)  # type: ignore[attr-defined]
        return ChatCompletion(**resp.json())


class AsyncCompletions:
    """Async chat completions: client.chat.completions.create(...)"""

    def __init__(self, client: object) -> None:
        self._client = client  # AsyncRightNow

    async def create(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        stream: bool = False,
        **kwargs: Any,
    ) -> ChatCompletion | AsyncStream:
        payload: dict[str, Any] = {"model": model, "messages": messages, "stream": stream, **kwargs}

        if stream:
            cm = self._client._http.stream(  # type: ignore[attr-defined]
                "POST",
                self._client._build_url("/v1/chat/completions"),  # type: ignore[attr-defined]
                json=payload,
                headers=self._client._headers,  # type: ignore[attr-defined]
                timeout=self._client.timeout,  # type: ignore[attr-defined]
            )
            response = await cm.__aenter__()
            self._client._raise_for_status(response)  # type: ignore[attr-defined]
            return AsyncStream(response, context_manager=cm)

        resp = await self._client._http.post(  # type: ignore[attr-defined]
            self._client._build_url("/v1/chat/completions"),  # type: ignore[attr-defined]
            json=payload,
            headers=self._client._headers,  # type: ignore[attr-defined]
            timeout=self._client.timeout,  # type: ignore[attr-defined]
        )
        self._client._raise_for_status(resp)  # type: ignore[attr-defined]
        return ChatCompletion(**resp.json())


class Chat:
    """Sync chat resource: client.chat.completions"""

    def __init__(self, client: object) -> None:
        self.completions = Completions(client)


class AsyncChat:
    """Async chat resource: client.chat.completions"""

    def __init__(self, client: object) -> None:
        self.completions = AsyncCompletions(client)
