ForcePilot/backend/package/yuxi/channel/extensions/twitch/eventsub.py

import asyncio
import json
import logging
import time
import uuid

import websockets

from yuxi.channel.extensions.twitch.config import TwitchAccountConfig, normalize_twitch_channel
from yuxi.channel.extensions.twitch.token import denormalize_token, resolve_twitch_token
from yuxi.channel.extensions.twitch.types import TwitchEventSubNotification

logger = logging.getLogger(__name__)

EVENTSUB_WS_URL = "wss://eventsub.wss.twitch.tv/ws"
EVENTSUB_SUBSCRIBE_URL = "https://api.twitch.tv/helix/eventsub/subscriptions"

# Core EventSub subscription types we support
CORE_SUBSCRIPTION_TYPES = [
    "channel.follow",
    "channel.subscribe",
    "channel.subscription.gift",
    "channel.subscription.message",
    "channel.cheer",
    "stream.online",
    "stream.offline",
    "channel.raid",
    "channel.update",
    "channel.ban",
    "channel.unban",
]


class TwitchEventSubManager:
    def __init__(self, logger_=None):
        self._logger = logger_ or logger
        self._ws = None
        self._session_id: str | None = None
        self._subscriptions: set[str] = set()
        self._handlers: list[callable] = []
        self._abort = asyncio.Event()
        self._last_keepalive = time.time()
        self._reconnect_url: str | None = None
        self._account: TwitchAccountConfig | None = None
        self._account_id: str = "default"
        self._connected = False

    def on_notification(self, handler: callable):
        self._handlers.append(handler)

    async def start(self, account: TwitchAccountConfig, account_id: str = "default"):
        self._account = account
        self._account_id = account_id
        self._abort.clear()
        self._connected = False
        await self._connect()

    async def stop(self):
        self._abort.set()
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None
        self._connected = False
        self._session_id = None

    async def _connect(self):
        url = self._reconnect_url or EVENTSUB_WS_URL
        self._reconnect_url = None
        try:
            self._ws = await websockets.connect(url, ping_interval=None)
            self._logger.info("EventSub WebSocket connected: %s", url)
            await self._message_loop()
        except asyncio.CancelledError:
            raise
        except Exception as e:
            self._logger.warning("EventSub connection error: %s", e)

    async def _message_loop(self):
        while not self._abort.is_set() and self._ws:
            try:
                raw = await asyncio.wait_for(self._ws.recv(), timeout=60)
                await self._handle_message(raw)
            except TimeoutError:
                if time.time() - self._last_keepalive > 70:
                    self._logger.warning("EventSub keepalive timeout, reconnecting")
                    await self._reconnect()
                    break
            except websockets.exceptions.ConnectionClosed:
                self._logger.warning("EventSub connection closed")
                if not self._abort.is_set():
                    await self._reconnect()
                break
            except Exception as e:
                self._logger.warning("EventSub message loop error: %s", e)
                if not self._abort.is_set():
                    await asyncio.sleep(5)

    async def _handle_message(self, raw: str):
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            self._logger.warning("EventSub invalid JSON: %s", raw[:200])
            return

        msg_type = data.get("metadata", {}).get("message_type")
        payload = data.get("payload", {})

        if msg_type == "session_welcome":
            session = payload.get("session", {})
            self._session_id = session.get("id")
            self._connected = True
            self._last_keepalive = time.time()
            self._logger.info("EventSub session welcome: %s", self._session_id)
            await self._create_subscriptions()

        elif msg_type == "session_keepalive":
            self._last_keepalive = time.time()

        elif msg_type == "notification":
            self._last_keepalive = time.time()
            await self._dispatch_notification(payload)

        elif msg_type == "session_reconnect":
            session = payload.get("session", {})
            reconnect_url = session.get("reconnect_url")
            if reconnect_url:
                self._logger.info("EventSub reconnect requested: %s", reconnect_url)
                self._reconnect_url = reconnect_url
                await self._reconnect()

        elif msg_type == "revocation":
            sub_type = payload.get("subscription", {}).get("type")
            self._logger.warning("EventSub subscription revoked: %s", sub_type)

        else:
            self._logger.debug("EventSub unknown message type: %s", msg_type)

    async def _dispatch_notification(self, payload: dict):
        subscription = payload.get("subscription", {})
        event = payload.get("event", {})
        sub_type = subscription.get("type", "unknown")

        notification = TwitchEventSubNotification(
            id=str(uuid.uuid4()),
            subscription_type=sub_type,
            event=event,
            channel=event.get("broadcaster_user_login", "")
            or event.get("from_broadcaster_user_login", "")
            or event.get("broadcaster_user_name", ""),
        )

        for handler in self._handlers:
            try:
                await handler(notification)
            except Exception as e:
                self._logger.warning("EventSub handler error: %s", e)

    async def _create_subscriptions(self):
        if not self._session_id:
            return
        account = self._account
        if not account or not account.client_id or not account.client_secret:
            self._logger.warning("EventSub: missing client_id/client_secret, cannot create subscriptions")
            return

        token, _ = resolve_twitch_token(account)
        headers = {
            "Client-Id": account.client_id,
            "Authorization": f"Bearer {denormalize_token(token)}",
            "Content-Type": "application/json",
        }

        import aiohttp

        broadcaster_id = await self._resolve_broadcaster_id()
        if not broadcaster_id:
            self._logger.warning("EventSub: cannot resolve broadcaster_id")
            return

        for sub_type in CORE_SUBSCRIPTION_TYPES:
            if sub_type in self._subscriptions:
                continue
            condition = self._build_condition(sub_type, broadcaster_id)
            body = {
                "type": sub_type,
                "version": "1",
                "condition": condition,
                "transport": {
                    "method": "websocket",
                    "session_id": self._session_id,
                },
            }
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        EVENTSUB_SUBSCRIBE_URL,
                        headers=headers,
                        json=body,
                        timeout=aiohttp.ClientTimeout(total=15),
                    ) as resp:
                        if resp.status in (202, 409):
                            self._subscriptions.add(sub_type)
                            self._logger.debug("EventSub subscribed: %s", sub_type)
                        else:
                            text = await resp.text()
                            self._logger.warning("EventSub subscribe failed %s: %s - %s", sub_type, resp.status, text)
            except Exception as e:
                self._logger.warning("EventSub subscribe error %s: %s", sub_type, e)

    def _build_condition(self, sub_type: str, broadcaster_id: str) -> dict:
        if sub_type in (
            "channel.follow",
            "channel.subscribe",
            "channel.subscription.gift",
            "channel.subscription.message",
            "channel.cheer",
            "stream.online",
            "stream.offline",
            "channel.raid",
            "channel.update",
            "channel.ban",
            "channel.unban",
        ):
            return {"broadcaster_user_id": broadcaster_id}
        return {"broadcaster_user_id": broadcaster_id}

    async def _resolve_broadcaster_id(self) -> str | None:
        account = self._account
        if not account:
            return None
        try:
            import aiohttp

            token, _ = resolve_twitch_token(account)
            headers = {
                "Client-Id": account.client_id,
                "Authorization": f"Bearer {denormalize_token(token)}",
            }
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    "https://api.twitch.tv/helix/users",
                    headers=headers,
                    params={"login": normalize_twitch_channel(account.channel)},
                    timeout=aiohttp.ClientTimeout(total=10),
                ) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        users = data.get("data", [])
                        if users:
                            return users[0].get("id")
        except Exception as e:
            self._logger.warning("EventSub resolve broadcaster_id error: %s", e)
        return None

    async def _reconnect(self):
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None
        if not self._abort.is_set():
            await asyncio.sleep(2)
            await self._connect()