ForcePilot/backend/package/yuxi/channel/extensions/twitter/gateway.py

from __future__ import annotations

import asyncio
import json
import logging

import tweepy

from yuxi.channel.extensions.twitter.auth import create_tweepy_client
from yuxi.channel.extensions.twitter.config import TwitterConfigAdapter
from yuxi.channel.extensions.twitter.dedupe import get_deduplicator
from yuxi.channel.extensions.twitter.monitor import (
    convert_dm_event_to_unified,
    convert_group_event_to_unified,
)
from yuxi.channel.extensions.twitter.security import TwitterSecurity

logger = logging.getLogger(__name__)


class TwitterGateway:
    MODE_WEBHOOK = "webhook"
    MODE_POLLING = "polling"

    def __init__(self):
        self._config_adapter = TwitterConfigAdapter()
        self._security = TwitterSecurity()
        self._tasks: list[asyncio.Task] = []
        self._running = False
        self._queue: asyncio.Queue | None = None
        self._abort_event: asyncio.Event | None = None
        self._account: dict = {}
        self._mode: str | None = None
        self._client: tweepy.Client | None = None

    async def start(self, ctx) -> object:
        account = await self._resolve_account(ctx)
        self._account = account
        if not account.get("api_key") or not account.get("access_token"):
            logger.warning(
                "twitter account %s not configured, skipping start",
                account.get("account_id"),
            )
            return {"running": False, "reason": "not-configured"}

        self._queue = asyncio.Queue(maxsize=1000)
        self._abort_event = asyncio.Event()
        self._running = True

        self._client = create_tweepy_client(account)

        connection_mode = account.get("connection_mode", "auto")
        webhook_url = account.get("webhook_url", "")

        if connection_mode == "webhook" or (connection_mode == "auto" and webhook_url):
            self._mode = self.MODE_WEBHOOK
            task = asyncio.create_task(
                self._start_webhook_mode(account, self._queue, self._abort_event)
            )
        else:
            self._mode = self.MODE_POLLING
            task = asyncio.create_task(
                self._poll_loop(account, self._queue, self._abort_event)
            )

        self._tasks.append(task)
        logger.info(
            "twitter gateway started for account %s (mode=%s)",
            account.get("account_id"),
            self._mode,
        )
        return {
            "running": True,
            "account_id": account.get("account_id"),
            "queue": self._queue,
            "mode": self._mode,
        }

    async def stop(self, ctx) -> None:
        self._running = False
        if self._abort_event:
            self._abort_event.set()

        for task in self._tasks:
            task.cancel()
        self._tasks.clear()

        if self._mode == self.MODE_WEBHOOK and self._account.get("webhook_url"):
            await self._cleanup_webhook(self._account)

        self._queue = None
        self._abort_event = None
        self._client = None
        self._account = {}
        self._mode = None
        logger.info("twitter gateway stopped")

    async def _poll_loop(
        self, account: dict, queue: asyncio.Queue, abort_event: asyncio.Event
    ) -> None:
        deduplicator = get_deduplicator(account.get("account_id", "default"))
        self_user_id = account.get("user_id", "")

        try:
            me = await asyncio.to_thread(self._client.get_me)
            me_data = me.get("data", {})
            if me_data.get("id"):
                self_user_id = str(me_data["id"])
        except Exception:
            logger.warning(
                "Twitter: cannot get self_user_id from API, falling back to config"
            )

        since_id: str | None = None
        interval = account.get("polling_interval_sec", 180)

        while self._running and not abort_event.is_set():
            try:
                if not self._client:
                    self._client = create_tweepy_client(account)

                kwargs: dict = {
                    "max_results": 50,
                    "dm_event_fields": [
                        "id",
                        "text",
                        "event_type",
                        "created_at",
                        "sender_id",
                        "participant_ids",
                        "dm_conversation_id",
                        "attachments",
                        "referenced_tweets",
                    ],
                    "expansions": ["sender_id", "attachments.media_key"],
                }
                if since_id:
                    kwargs["since_id"] = since_id

                response = await asyncio.to_thread(
                    self._client.get_dm_events,
                    **kwargs,
                )

                events = response.get("data", [])
                if events:
                    for event in events:
                        event_type = event.get("event_type", "")
                        if event_type not in (
                            "MessageCreate",
                            "ParticipantsJoin",
                            "ParticipantsLeave",
                        ):
                            continue

                        if event_type in ("ParticipantsJoin", "ParticipantsLeave"):
                            unified = convert_group_event_to_unified(
                                event, account.get("account_id", "default")
                            )
                        else:
                            sender_id = str(event.get("sender_id", ""))
                            if sender_id == self_user_id:
                                continue
                            unified = convert_dm_event_to_unified(
                                event, account.get("account_id", "default")
                            )

                        if unified:
                            event_id = str(event.get("id", ""))
                            if deduplicator.is_duplicate(event_id):
                                continue
                            try:
                                queue.put_nowait(unified)
                            except asyncio.QueueFull:
                                logger.warning(
                                    "Twitter message queue full, dropping message"
                                )

                    since_id = str(events[-1].get("id", ""))

            except tweepy.TooManyRequests:
                logger.warning("Twitter polling rate limited, waiting %ds", interval)
                await asyncio.sleep(interval)
            except Exception:
                logger.exception(
                    "Twitter polling error for account %s", account.get("account_id")
                )
                await asyncio.sleep(min(interval, 60))

            try:
                await asyncio.wait_for(abort_event.wait(), timeout=interval)
                break
            except TimeoutError:
                pass

    async def _start_webhook_mode(
        self, account: dict, queue: asyncio.Queue, abort_event: asyncio.Event
    ) -> None:
        from yuxi.channel.extensions.twitter.webhook import (
            register_webhook,
            subscribe_account_activity,
        )

        try:
            await register_webhook(account)
            await subscribe_account_activity(account)
            logger.info(
                "Twitter webhook registered and subscribed for account %s",
                account.get("account_id"),
            )
        except Exception:
            logger.exception("Twitter webhook setup failed, falling back to polling")
            self._mode = self.MODE_POLLING
            task = asyncio.create_task(self._poll_loop(account, queue, abort_event))
            self._tasks.append(task)
            return

        while not abort_event.is_set():
            await asyncio.sleep(1)

    async def _cleanup_webhook(self, account: dict) -> None:
        from yuxi.channel.extensions.twitter.webhook import unsubscribe_account_activity

        try:
            await unsubscribe_account_activity(account)
        except Exception:
            logger.exception("Twitter webhook cleanup failed")

    async def handle_webhook_callback(self, body: bytes, signature: str) -> dict:
        from yuxi.channel.extensions.twitter.webhook import verify_webhook_signature
        from yuxi.channel.extensions.twitter.monitor import (
            convert_webhook_event_to_unified,
        )

        account = self._account
        if not verify_webhook_signature(signature, body, account.get("api_secret", "")):
            logger.warning("Twitter webhook: invalid signature")
            return {"status": "invalid_signature"}

        data = json.loads(body)
        for dm_event in data.get("direct_message_events", []):
            unified = convert_webhook_event_to_unified(
                dm_event,
                account.get("user_id", ""),
                account.get("account_id", "default"),
            )
            if unified and self._queue:
                try:
                    self._queue.put_nowait(unified)
                except asyncio.QueueFull:
                    logger.warning("Twitter webhook: queue full, dropping message")

        return {"status": "ok"}

    async def _resolve_account(self, ctx) -> dict:
        config = getattr(ctx, "config", {}) if ctx else {}
        account_id = getattr(ctx, "account_id", "default")
        self._config_adapter.list_account_ids(config)
        return await self._config_adapter.resolve_account(account_id)