from __future__ import annotations

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

from ._types import ChatCompletionChunk

T = TypeVar("T")


class Stream:
    """Synchronous SSE stream wrapper for chat completion chunks."""

    def __init__(self, response: object, context_manager: object = None) -> None:
        self._response = response  # httpx.Response (streaming)
        self._cm = context_manager  # keep context manager alive
        self._iterator: Iterator[str] | None = None

    def __iter__(self) -> Iterator[ChatCompletionChunk]:
        return self._parse()

    def _parse(self) -> Iterator[ChatCompletionChunk]:
        for line in self._response.iter_lines():  # type: ignore[attr-defined]
            line = line.strip()
            if not line or not line.startswith("data: "):
                continue
            data = line[6:]
            if data == "[DONE]":
                break
            try:
                chunk = json.loads(data)
                yield ChatCompletionChunk(**chunk)
            except (json.JSONDecodeError, Exception):
                continue

    def close(self) -> None:
        self._response.close()  # type: ignore[attr-defined]
        if self._cm is not None:
            self._cm.__exit__(None, None, None)
            self._cm = None

    def __enter__(self) -> Stream:
        return self

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


class AsyncStream:
    """Asynchronous SSE stream wrapper for chat completion chunks."""

    def __init__(self, response: object, context_manager: object = None) -> None:
        self._response = response  # httpx.Response (async streaming)
        self._cm = context_manager  # keep context manager alive

    def __aiter__(self) -> AsyncIterator[ChatCompletionChunk]:
        return self._parse()

    async def _parse(self) -> AsyncIterator[ChatCompletionChunk]:
        async for line in self._response.aiter_lines():  # type: ignore[attr-defined]
            line = line.strip()
            if not line or not line.startswith("data: "):
                continue
            data = line[6:]
            if data == "[DONE]":
                break
            try:
                chunk = json.loads(data)
                yield ChatCompletionChunk(**chunk)
            except (json.JSONDecodeError, Exception):
                continue

    async def close(self) -> None:
        await self._response.aclose()  # type: ignore[attr-defined]
        if self._cm is not None:
            await self._cm.__aexit__(None, None, None)
            self._cm = None

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

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