from __future__ import annotations

import asyncio
import json
from collections.abc import Awaitable, Callable
from typing import Any

import httpx
from httpx_sse import aconnect_sse


class SSEListener:
    def __init__(
        self,
        server_url: str,
        group_id: str,
        peer_id: str,
        on_event: Callable[[dict[str, Any]], Awaitable[None]],
        groups_header: str = "",
    ) -> None:
        self.server_url = server_url.rstrip("/")
        self.group_id = group_id
        self.peer_id = peer_id
        self.on_event = on_event
        self.groups_header = groups_header
        self.connected: bool = False
        self._task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        self._task = asyncio.create_task(self._listen_loop())

    async def stop(self) -> None:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        self.connected = False

    async def _listen_loop(self) -> None:
        backoff = 1.0
        max_backoff = 30.0
        url = f"{self.server_url}/notifications/{self.group_id}/{self.peer_id}"
        headers: dict[str, str] = {}
        if self.groups_header:
            headers["X-Peer-Id"] = self.peer_id
            headers["X-Groups"] = self.groups_header

        while True:
            try:
                async with httpx.AsyncClient(timeout=None, headers=headers) as client:
                    async with aconnect_sse(client, "GET", url) as event_source:
                        self.connected = True
                        backoff = 1.0
                        async for event in event_source.aiter_sse():
                            if event.event != "notification" or not event.data:
                                continue
                            try:
                                parsed: dict[str, Any] = json.loads(event.data)
                            except (json.JSONDecodeError, TypeError):
                                continue
                            await self.on_event(parsed)
            except asyncio.CancelledError:
                self.connected = False
                return
            except (httpx.ConnectError, httpx.ReadError, httpx.RemoteProtocolError, httpx.ConnectTimeout):
                self.connected = False
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, max_backoff)
            except Exception:
                self.connected = False
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, max_backoff)
