"""
SSE 事件流解析器。

将 httpx-sse 推送的原始 SSE 事件转换为 CMDC Event dataclass。
"""

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Iterator
from typing import Any

import httpx
from httpx_sse import EventSource, aconnect_sse, connect_sse

from cmdc.types import Event


def _parse_data(raw_data: str) -> dict[str, Any]:
    """解析 SSE data 字段（JSON 或空字符串）。"""
    if not raw_data or raw_data.strip() == "{}":
        return {}
    try:
        parsed = json.loads(raw_data)
        return parsed if isinstance(parsed, dict) else {}
    except json.JSONDecodeError:
        return {}


class SSEIterator:
    """同步 SSE 事件迭代器。

    用法::

        with SSEIterator(client, url, headers) as it:
            for event in it:
                print(event.type, event.data)
    """

    def __init__(
        self,
        http_client: httpx.Client,
        url: str,
        headers: dict[str, str],
    ) -> None:
        self._client = http_client
        self._url = url
        self._headers = headers
        self._ctx: Any = None
        self._source: EventSource | None = None

    def __enter__(self) -> "SSEIterator":
        self._ctx = connect_sse(
            self._client,
            "GET",
            self._url,
            headers=self._headers,
            timeout=None,
        )
        self._source = self._ctx.__enter__()
        return self

    def __exit__(self, *args: Any) -> None:
        if self._ctx is not None:
            self._ctx.__exit__(*args)

    def __iter__(self) -> Iterator[Event]:
        if self._source is None:
            raise RuntimeError("SSEIterator must be used as a context manager")
        for sse in self._source.iter_sse():
            # 跳过心跳注释行（httpx-sse 不会暴露注释行，这里预防性处理）
            if not sse.event:
                continue
            event = Event(type=sse.event, data=_parse_data(sse.data))
            yield event
            # 终止事件：关闭 SSE 连接
            if sse.event in ("agent_end", "agent_abort"):
                break


class AsyncSSEIterator:
    """异步 SSE 事件迭代器。

    用法::

        async with AsyncSSEIterator(client, url, headers) as it:
            async for event in it:
                print(event.type, event.data)
    """

    def __init__(
        self,
        http_client: httpx.AsyncClient,
        url: str,
        headers: dict[str, str],
    ) -> None:
        self._client = http_client
        self._url = url
        self._headers = headers
        self._ctx: Any = None
        self._source: EventSource | None = None

    async def __aenter__(self) -> "AsyncSSEIterator":
        self._ctx = aconnect_sse(
            self._client,
            "GET",
            self._url,
            headers=self._headers,
            timeout=None,
        )
        self._source = await self._ctx.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._ctx is not None:
            await self._ctx.__aexit__(*args)

    async def __aiter__(self) -> AsyncIterator[Event]:
        if self._source is None:
            raise RuntimeError("AsyncSSEIterator must be used as an async context manager")
        async for sse in self._source.aiter_sse():
            if not sse.event:
                continue
            event = Event(type=sse.event, data=_parse_data(sse.data))
            yield event
            if sse.event in ("agent_end", "agent_abort"):
                break
