from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import queue
import threading
from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from dataclasses import dataclass
from typing import Any, TypeAlias, TypedDict, TypeVar, cast

from .exceptions import APIError

BodyData = bytes | Mapping[str, Any]

logger = logging.getLogger(__name__)

ContentCallback: TypeAlias = Callable[[bytes], int]


@dataclass
class SSEEvent:
    """A single parsed Server-Sent Event.

    Attributes:
        event: Event type (defaults to ``"message"`` when omitted by the server).
        data: Raw data payload, may be multi-line (lines joined with ``"\\n"``).
        id: Optional event ID set by the ``id:`` field.
        retry: Optional reconnection delay in milliseconds from the ``retry:`` field.
        headers: Response headers from the enclosing HTTP response (injected after parsing).
    """

    event: str
    data: str
    id: str | None = None
    retry: int | None = None
    headers: Mapping[str, str] | None = None


T = TypeVar("T")


class SSEDecoder:
    """Incremental SSE stream parser.

    Feed raw bytes chunks via :meth:`feed`; it returns a list of complete
    :class:`SSEEvent` objects found in the accumulated buffer.  Handles
    ``\\r\\n``, ``\\r``, and ``\\n`` line endings and resets the internal
    buffer if it exceeds 2 MB to guard against memory exhaustion.
    """

    __slots__ = ("_buf", "_data_lines", "_event", "_id", "_retry")

    def __init__(self) -> None:
        self._buf = bytearray()
        self._data_lines: list[str] = []
        self._event: str = "message"
        self._id: str | None = None
        self._retry: int | None = None

    def feed(self, chunk: bytes) -> list[SSEEvent]:
        self._buf.extend(chunk)
        out: list[SSEEvent] = []

        if len(self._buf) > 2_000_000:
            self._buf.clear()
            self._data_lines.clear()
            self._event = "message"
            self._id = None
            self._retry = None
            return out

        while True:
            nl = self._buf.find(b"\n")
            if nl < 0:
                break

            line = bytes(self._buf[:nl])
            del self._buf[: nl + 1]

            if line.endswith(b"\r"):
                line = line[:-1]

            if line == b"":
                if self._data_lines:
                    out.append(
                        SSEEvent(
                            event=self._event or "message",
                            data="\n".join(self._data_lines),
                            id=self._id,
                            retry=self._retry,
                        )
                    )
                    self._data_lines.clear()

                self._event = "message"
                self._retry = None
                self._id = None
                continue

            if line.startswith(b":"):
                continue

            if b":" in line:
                field_b, value_b = line.split(b":", 1)
                if value_b.startswith(b" "):
                    value_b = value_b[1:]
            else:
                field_b, value_b = line, b""

            if field_b == b"data":
                self._data_lines.append(value_b.decode("utf-8", errors="replace"))
            elif field_b == b"event":
                self._event = value_b.decode("utf-8", errors="replace")
            elif field_b == b"id":
                self._id = value_b.decode("utf-8", errors="replace")
            elif field_b == b"retry":
                with contextlib.suppress(ValueError):
                    self._retry = int(value_b.decode("utf-8", errors="replace"))

        return out


def stream_json_events(
    *,
    session_request: Callable[..., Any],
    method: str,
    url: str,
    headers: Mapping[str, str],
    params: Mapping[str, Any] | None,
    json_body: dict[str, Any] | None,
    data: Any | None,
    timeout: float | tuple[float, float],
    parser: Callable[[SSEEvent, dict[str, Any] | None], T],
    max_queue: int = 256,
) -> Iterator[T]:
    """Consume a streaming HTTP response as parsed SSE items (sync).

    Runs the HTTP request in a background daemon thread and forwards parsed
    items through a bounded queue.  Back-pressure stops the producer when
    the queue is full (``max_queue`` items).

    Args:
        session_request: Callable that performs the HTTP request (e.g. ``session.request``).
        method: HTTP method string (e.g. ``"POST"``).
        url: Fully-qualified request URL.
        headers: Request headers mapping.
        params: Optional URL query parameters.
        json_body: Optional JSON-serialisable request body.
        data: Optional raw request body (mutually exclusive with ``json_body``).
        timeout: Request timeout in seconds, or ``(connect_timeout, read_timeout)`` tuple.
        parser: Callable that converts an ``(SSEEvent, dict | None)`` pair to ``T``.
            Raised exceptions are logged and the event is skipped.
        max_queue: Maximum number of parsed items to buffer (default ``256``).

    Yields:
        Parsed items of type ``T`` as they arrive from the server.

    Raises:
        APIError: If the server returns a 4xx/5xx status code.
        Exception: Any other exception raised during streaming is re-raised in the caller.
    """
    q: queue.Queue[object] = queue.Queue(maxsize=max_queue)
    sentinel = object()
    stop = threading.Event()
    decoder = SSEDecoder()

    resp_ref: dict[str, Any] = {"resp": None}

    def _safe_put(x: object) -> None:
        try:
            q.put_nowait(x)
        except queue.Full:
            stop.set()

    def worker() -> None:
        resp = None
        try:
            resp = session_request(
                method,
                url,
                params=params,
                headers=headers,
                json=json_body,
                data=data,
                timeout=timeout,
                stream=True,
            )
            resp_ref["resp"] = resp

            resp_headers = {k.lower(): v for k, v in resp.headers.items()}

            if getattr(resp, "status_code", 0) >= 400:
                _safe_put(
                    APIError(
                        resp.status_code, b"", msg="HTTP error in streaming response"
                    )
                )
                return

            for chunk in resp.iter_content():
                if stop.is_set():
                    break

                for ev in decoder.feed(chunk):
                    ev.headers = resp_headers

                    obj: dict[str, Any] | None = None
                    payload = (ev.data or "").strip()
                    if payload:
                        if payload == "[DONE]":
                            continue
                        try:
                            obj_any: Any = json.loads(payload)
                        except Exception as e:
                            logger.warning(
                                "Invalid JSON chunk: event=%r data=%r (%s)",
                                ev.event,
                                payload,
                                e,
                            )
                            obj_any = None

                        if isinstance(obj_any, dict):
                            obj = cast("dict[str, Any]", obj_any)
                        elif obj_any is not None:
                            logger.warning(
                                "JSON chunk is not an object: event=%r data=%r",
                                ev.event,
                                obj_any,
                            )

                    try:
                        item = parser(ev, obj)
                    except Exception as e:
                        logger.warning("Parser error: event=%r obj=%r (%s)", ev, obj, e)
                        continue

                    try:
                        q.put_nowait(item)
                    except queue.Full:
                        stop.set()
                        if getattr(resp, "quit_now", None):
                            resp.quit_now.set()
                        return

        except Exception as e:
            if not stop.is_set():
                _safe_put(e)
        finally:
            try:
                if resp is not None and getattr(resp, "quit_now", None):
                    resp.quit_now.set()
                if resp is not None and getattr(resp, "close", None):
                    resp.close()
            except Exception:  # noqa: S110
                pass
            _safe_put(sentinel)

    t = threading.Thread(target=worker, name="seaart-stream", daemon=True)
    t.start()

    try:
        while True:
            try:
                x = q.get(timeout=0.5)
            except queue.Empty:
                if stop.is_set():
                    break
                continue

            if x is sentinel:
                break
            if isinstance(x, Exception):
                raise x
            yield cast("T", x)
    finally:
        stop.set()
        resp = resp_ref.get("resp")
        try:
            if resp is not None and getattr(resp, "quit_now", None):
                resp.quit_now.set()
        except Exception:  # noqa: S110
            pass
        t.join(timeout=1.0)


async def astream_json_events(
    *,
    session_request: Callable[..., Any],
    method: str,
    url: str,
    headers: Mapping[str, str],
    params: Mapping[str, Any] | None,
    json_body: dict[str, Any] | None,
    data: Any | None,
    timeout: float | tuple[float, float],
    parser: Callable[[SSEEvent, dict[str, Any] | None], T],
    max_queue: int = 256,
) -> AsyncIterator[T]:
    """Consume a streaming HTTP response as parsed SSE items (async).

    Runs the HTTP request in a background ``asyncio`` task and forwards
    parsed items through a bounded ``asyncio.Queue``.  The task is cancelled
    when the caller breaks out of the iteration or an error occurs.

    Args:
        session_request: Async callable that performs the HTTP request.
        method: HTTP method string (e.g. ``"POST"``).
        url: Fully-qualified request URL.
        headers: Request headers mapping.
        params: Optional URL query parameters.
        json_body: Optional JSON-serialisable request body.
        data: Optional raw request body (mutually exclusive with ``json_body``).
        timeout: Request timeout in seconds, or ``(connect_timeout, read_timeout)`` tuple.
        parser: Callable that converts an ``(SSEEvent, dict | None)`` pair to ``T``.
            Raised exceptions are logged and the event is skipped.
        max_queue: Maximum number of parsed items to buffer (default ``256``).

    Yields:
        Parsed items of type ``T`` as they arrive from the server.

    Raises:
        APIError: If the server returns a 4xx/5xx status code.
        Exception: Any other exception raised during streaming is re-raised in the caller.
    """
    q: asyncio.Queue[object] = asyncio.Queue(maxsize=max_queue)
    sentinel = object()
    decoder = SSEDecoder()

    async def runner() -> None:
        resp = None
        try:
            resp = await session_request(
                method,
                url,
                params=params,
                headers=headers,
                json=json_body,
                data=data,
                timeout=timeout,
                stream=True,
            )

            resp_headers = {k.lower(): v for k, v in resp.headers.items()}

            if getattr(resp, "status_code", 0) >= 400:
                await q.put(
                    APIError(
                        resp.status_code, b"", msg="HTTP error in streaming response"
                    )
                )
                return

            async for chunk in resp.aiter_content():
                for ev in decoder.feed(chunk):
                    ev.headers = resp_headers

                    obj: dict[str, Any] | None = None
                    payload = (ev.data or "").strip()
                    if payload:
                        if payload == "[DONE]":
                            continue
                        try:
                            obj_any: Any = json.loads(payload)
                        except Exception as e:
                            logger.warning(
                                "Invalid JSON chunk: event=%r data=%r (%s)",
                                ev.event,
                                payload,
                                e,
                            )
                            obj_any = None

                        if isinstance(obj_any, dict):
                            obj = cast("dict[str, Any]", obj_any)
                        elif obj_any is not None:
                            logger.warning(
                                "JSON chunk is not an object: event=%r data=%r",
                                ev.event,
                                obj_any,
                            )

                    try:
                        item = parser(ev, obj)
                    except Exception as e:
                        logger.warning("Parser error: event=%r obj=%r (%s)", ev, obj, e)
                        continue

                    try:
                        q.put_nowait(item)
                    except asyncio.QueueFull:
                        if getattr(resp, "quit_now", None):
                            resp.quit_now.set()
                        return

        except Exception as e:
            await q.put(e)
        finally:
            try:
                if resp is not None and getattr(resp, "quit_now", None):
                    resp.quit_now.set()
                if resp is not None:
                    await resp.aclose()
            except Exception:  # noqa: S110
                pass
            await q.put(sentinel)

    task = asyncio.create_task(runner(), name="seaart-astream")

    try:
        while True:
            x = await q.get()
            if x is sentinel:
                break
            if isinstance(x, Exception):
                raise x
            yield cast("T", x)
    finally:
        task.cancel()
        with contextlib.suppress(Exception):
            await task


class RequestStreamKwargs(TypedDict, total=False):
    params: Mapping[str, Any] | None
    headers: Mapping[str, str] | None
    json: dict[str, Any] | None
    data: BodyData | None
    parser: Callable[[SSEEvent, dict[str, Any] | None], Any]
