from __future__ import annotations

import json
from typing import Any, AsyncIterator, Generic, Iterator, TypeVar

import httpx

from ._sse import SSEDecoder
from .types.chat.completion import ChatCompletion
from .types.rungate_metadata import RungateMetadata

T = TypeVar("T")


# ─── Low-Level Stream ────────────────────────────────────


class RungateStream(Generic[T]):
    """Low-level streaming response. Yields parsed SSE data as dicts.

    Automatically captures the `event: rungate` metadata event silently.
    Must be used as a context manager to ensure cleanup.

    Usage:
        with client.chat.completions.create(stream=True, ...) as stream:
            for chunk in stream:
                print(chunk["choices"][0]["delta"]["content"], end="")
            print(stream.rungate.cost_usd)
    """

    def __init__(self, response: httpx.Response) -> None:
        self._response = response
        self._rungate_metadata: RungateMetadata | None = None
        self._decoder = SSEDecoder()

    @property
    def rungate(self) -> RungateMetadata:
        """Rungate governance metadata (cost, tokens, run_id, etc.).

        Available after stream iteration completes. Parsed from the
        final `event: rungate` SSE event.
        """
        if self._rungate_metadata is None:
            self._rungate_metadata = RungateMetadata()
        return self._rungate_metadata

    def __iter__(self) -> Iterator[dict[str, Any]]:
        for sse_event in self._decoder.iter_events(self._response):
            if sse_event.event == "rungate":
                # Capture metadata silently — don't yield to user
                try:
                    data = json.loads(sse_event.data)
                    tokens = data.get("tokens")
                    tokens_used = None
                    if isinstance(tokens, dict):
                        try:
                            tokens_used = int(tokens.get("input", 0)) + int(tokens.get("output", 0))
                        except (ValueError, TypeError):
                            tokens_used = None
                    self._rungate_metadata = RungateMetadata(
                        run_id=data.get("run_id"),
                        step_id=data.get("step_id"),
                        provider=data.get("provider"),
                        model=data.get("model"),
                        cost_usd=data.get("cost_usd"),
                        tokens_used=tokens_used,
                    )
                except Exception:
                    pass
                continue

            if sse_event.data == "[DONE]":
                # OpenAI stream terminator — stop iteration
                continue

            # Parse provider chunk
            try:
                chunk = json.loads(sse_event.data)
                yield chunk
            except json.JSONDecodeError:
                # Non-JSON data line — skip
                continue

    def close(self) -> None:
        """Close the underlying HTTP response."""
        self._response.close()

    def __enter__(self) -> RungateStream[T]:
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


class AsyncRungateStream(Generic[T]):
    """Async low-level streaming response."""

    def __init__(self, response: httpx.Response) -> None:
        self._response = response
        self._rungate_metadata: RungateMetadata | None = None
        self._decoder = SSEDecoder()

    @property
    def rungate(self) -> RungateMetadata:
        if self._rungate_metadata is None:
            self._rungate_metadata = RungateMetadata()
        return self._rungate_metadata

    async def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
        async for sse_event in self._decoder.aiter_events(self._response):
            if sse_event.event == "rungate":
                try:
                    data = json.loads(sse_event.data)
                    tokens = data.get("tokens")
                    tokens_used = None
                    if isinstance(tokens, dict):
                        try:
                            tokens_used = int(tokens.get("input", 0)) + int(tokens.get("output", 0))
                        except (ValueError, TypeError):
                            tokens_used = None
                    self._rungate_metadata = RungateMetadata(
                        run_id=data.get("run_id"),
                        step_id=data.get("step_id"),
                        provider=data.get("provider"),
                        model=data.get("model"),
                        cost_usd=data.get("cost_usd"),
                        tokens_used=tokens_used,
                    )
                except Exception:
                    pass
                continue

            if sse_event.data == "[DONE]":
                continue

            try:
                chunk = json.loads(sse_event.data)
                yield chunk
            except json.JSONDecodeError:
                continue

    async def close(self) -> None:
        await self._response.aclose()

    async def __aenter__(self) -> AsyncRungateStream[T]:
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()


# ─── High-Level Stream ────────────────────────────────────


class ChatCompletionStream:
    """High-level streaming with text extraction and accumulation.

    Usage:
        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()
            print(stream.rungate.cost_usd)
    """

    def __init__(self, response: httpx.Response) -> None:
        self._stream = RungateStream[dict[str, Any]](response)
        self._chunks: list[dict[str, Any]] = []
        self._consumed = False
        self._started = False
        self._text_gen: Iterator[str] | None = None

    @property
    def text_stream(self) -> Iterator[str]:
        """Yield just the text content from each chunk.

        Follows Anthropic SDK convention. Can only be iterated once.
        """
        if self._consumed or self._started:
            raise RuntimeError("Stream has already been consumed")
        self._started = True
        self._text_gen = self._iter_text()
        return self._text_gen

    def _iter_text(self) -> Iterator[str]:
        for chunk in self._stream:
            self._chunks.append(chunk)
            choices = chunk.get("choices", [])
            if choices:
                delta = choices[0].get("delta", {})
                content = delta.get("content")
                if content:
                    yield content
        self._consumed = True

    @property
    def rungate(self) -> RungateMetadata:
        return self._stream.rungate

    def get_final_completion(self) -> ChatCompletion:
        """Reconstruct a full ChatCompletion from accumulated chunks.

        If not yet consumed, consumes the remaining stream first.
        Safe to call after partial iteration (early break).

        Note: Rungate metadata (.rungate.cost_usd, etc.) is only available
        after the full stream has been consumed, since the metadata event
        arrives at the end of the SSE stream.
        """
        if not self._consumed:
            try:
                if self._text_gen is not None:
                    # Exhaust the existing text generator (resumes from where it paused)
                    for _ in self._text_gen:
                        pass
                else:
                    # Never started text_stream — consume stream directly
                    for chunk in self._stream:
                        self._chunks.append(chunk)
            finally:
                self._consumed = True
        return _accumulate_chunks(self._chunks, self._stream.rungate)

    def close(self) -> None:
        self._stream.close()

    def __enter__(self) -> ChatCompletionStream:
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


class AsyncChatCompletionStream:
    """Async high-level streaming."""

    def __init__(self, response: httpx.Response) -> None:
        self._stream = AsyncRungateStream[dict[str, Any]](response)
        self._chunks: list[dict[str, Any]] = []
        self._consumed = False
        self._started = False
        self._text_gen: AsyncIterator[str] | None = None

    @property
    def text_stream(self) -> AsyncIterator[str]:
        if self._consumed or self._started:
            raise RuntimeError("Stream has already been consumed")
        self._started = True
        self._text_gen = self._iter_text()
        return self._text_gen

    async def _iter_text(self) -> AsyncIterator[str]:
        async for chunk in self._stream:
            self._chunks.append(chunk)
            choices = chunk.get("choices", [])
            if choices:
                delta = choices[0].get("delta", {})
                content = delta.get("content")
                if content:
                    yield content
        self._consumed = True

    @property
    def rungate(self) -> RungateMetadata:
        return self._stream.rungate

    async def get_final_completion(self) -> ChatCompletion:
        """Reconstruct a full ChatCompletion. Safe after partial iteration.

        Note: Rungate metadata is only available after full stream consumption.
        """
        if not self._consumed:
            try:
                if self._text_gen is not None:
                    async for _ in self._text_gen:
                        pass
                else:
                    async for chunk in self._stream:
                        self._chunks.append(chunk)
            finally:
                self._consumed = True
        return _accumulate_chunks(self._chunks, self._stream.rungate)

    async def close(self) -> None:
        await self._stream.close()

    async def __aenter__(self) -> AsyncChatCompletionStream:
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()


# ─── Stream Managers ──────────────────────────────────────


class ChatCompletionStreamManager:
    """Context manager that defers the HTTP request until __enter__."""

    def __init__(
        self,
        client: Any,
        *,
        request_json: dict[str, Any],
        request_headers: dict[str, str] | None = None,
    ) -> None:
        self._client = client
        self._request_json = request_json
        self._request_headers = request_headers
        self._response: httpx.Response | None = None

    def __enter__(self) -> ChatCompletionStream:
        self._response = self._client._request_stream(
            "POST",
            "/v1/chat/completions",
            json=self._request_json,
            headers=self._request_headers,
        )
        return ChatCompletionStream(self._response)

    def __exit__(self, *args: Any) -> None:
        if self._response:
            self._response.close()


class AsyncChatCompletionStreamManager:
    """Async context manager that defers the HTTP request until __aenter__."""

    def __init__(
        self,
        client: Any,
        *,
        request_json: dict[str, Any],
        request_headers: dict[str, str] | None = None,
    ) -> None:
        self._client = client
        self._request_json = request_json
        self._request_headers = request_headers
        self._response: httpx.Response | None = None

    async def __aenter__(self) -> AsyncChatCompletionStream:
        self._response = await self._client._request_stream(
            "POST",
            "/v1/chat/completions",
            json=self._request_json,
            headers=self._request_headers,
        )
        return AsyncChatCompletionStream(self._response)

    async def __aexit__(self, *args: Any) -> None:
        if self._response:
            await self._response.aclose()


# ─── Chunk Accumulation ──────────────────────────────────


def _accumulate_chunks(chunks: list[dict[str, Any]], metadata: RungateMetadata) -> ChatCompletion:
    """Reconstruct a ChatCompletion from streaming chunks."""
    if not chunks:
        return ChatCompletion(id="", model="", choices=[], rungate=metadata)

    first = chunks[0]
    accumulated_content = ""
    finish_reason = None

    for chunk in chunks:
        choices = chunk.get("choices", [])
        if choices:
            delta = choices[0].get("delta", {})
            content = delta.get("content", "")
            accumulated_content += content
            if choices[0].get("finish_reason"):
                finish_reason = choices[0]["finish_reason"]

    return ChatCompletion(
        id=first.get("id", ""),
        model=first.get("model", ""),
        choices=[
            {
                "index": 0,
                "message": {"role": "assistant", "content": accumulated_content},
                "finish_reason": finish_reason,
            }
        ],
        usage=None,
        rungate=metadata,
    )
