ForcePilot/backend/package/yuxi/channel/extensions/telegram/sticker.py

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

TELEGRAM_API_BASE = "https://api.telegram.org"

_CACHE_TTL_SECONDS = 3600


@dataclass
class StickerInfo:
    file_id: str
    file_unique_id: str
    width: int = 0
    height: int = 0
    is_animated: bool = False
    is_video: bool = False
    emoji: str = ""
    set_name: str = ""
    thumbnail_file_id: str = ""
    file_size: int = 0


@dataclass
class StickerSet:
    name: str = ""
    title: str = ""
    sticker_type: str = "regular"
    is_animated: bool = False
    is_video: bool = False
    stickers: list[StickerInfo] = field(default_factory=list)
    cached_at: float = 0.0

    @property
    def size(self) -> int:
        return len(self.stickers)

    @property
    def is_expired(self) -> bool:
        return time.monotonic() - self.cached_at > _CACHE_TTL_SECONDS


class TelegramStickerCache:
    def __init__(self):
        self._sets: dict[str, StickerSet] = {}
        self._emoji_index: dict[str, list[StickerInfo]] = {}

    def get_set(self, name: str) -> StickerSet | None:
        entry = self._sets.get(name)
        if entry and not entry.is_expired:
            return entry
        if entry:
            self.remove_set(name)
        return None

    def put_set(self, sticker_set: StickerSet) -> None:
        self._sets[sticker_set.name] = sticker_set
        self._index_emojis(sticker_set)

    def remove_set(self, name: str) -> None:
        entry = self._sets.pop(name, None)
        if entry:
            self._remove_emoji_index(entry)

    def search_by_emoji(self, emoji: str) -> list[StickerInfo]:
        return self._emoji_index.get(emoji, [])

    def clear(self) -> None:
        self._sets.clear()
        self._emoji_index.clear()

    def list_cached_sets(self) -> list[str]:
        return list(self._sets)

    def _index_emojis(self, sticker_set: StickerSet) -> None:
        for sticker in sticker_set.stickers:
            if sticker.emoji:
                self._emoji_index.setdefault(sticker.emoji, []).append(sticker)

    def _remove_emoji_index(self, sticker_set: StickerSet) -> None:
        to_remove: set[str] = set()
        for emoji, stickers in self._emoji_index.items():
            self._emoji_index[emoji] = [s for s in stickers if s.set_name != sticker_set.name]
            if not self._emoji_index[emoji]:
                to_remove.add(emoji)
        for emoji in to_remove:
            self._emoji_index.pop(emoji, None)


_sticker_cache = TelegramStickerCache()


def _parse_sticker(raw: dict) -> StickerInfo:
    thumb = raw.get("thumbnail", {}) or {}
    return StickerInfo(
        file_id=raw.get("file_id", ""),
        file_unique_id=raw.get("file_unique_id", ""),
        width=raw.get("width", 0),
        height=raw.get("height", 0),
        is_animated=raw.get("is_animated", False),
        is_video=raw.get("is_video", False),
        emoji=raw.get("emoji", ""),
        set_name=raw.get("set_name", ""),
        thumbnail_file_id=thumb.get("file_id", ""),
        file_size=raw.get("file_size", 0),
    )


def extract_sticker_from_message(message: dict) -> StickerInfo | None:
    sticker_raw = message.get("sticker")
    if not sticker_raw:
        return None
    return _parse_sticker(sticker_raw)


async def get_sticker_set(token: str, name: str) -> StickerSet | None:
    cached = _sticker_cache.get_set(name)
    if cached:
        return cached

    import httpx

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
            resp = await client.get(
                f"{TELEGRAM_API_BASE}/bot{token}/getStickerSet",
                params={"name": name},
            )
            data = resp.json() if resp.content else {}
            if not data.get("ok"):
                logger.warning("getStickerSet failed: %s", data.get("description", ""))
                return None

            result = data["result"]
            sticker_set = StickerSet(
                name=result.get("name", name),
                title=result.get("title", ""),
                sticker_type=result.get("sticker_type", "regular"),
                is_animated=result.get("is_animated", False),
                is_video=result.get("is_video", False),
                stickers=[_parse_sticker(s) for s in result.get("stickers", [])],
                cached_at=time.monotonic(),
            )
            _sticker_cache.put_set(sticker_set)
            logger.info(
                "Sticker set cached: %s (%d stickers)", sticker_set.name, sticker_set.size,
            )
            return sticker_set
    except Exception:
        logger.exception("getStickerSet network error for %s", name)
        return None


async def send_sticker(
    token: str,
    target_id: str,
    file_id: str,
    *,
    reply_to_id: str | None = None,
    thread_id: str | None = None,
    emoji: str | None = None,
) -> dict | None:
    import httpx

    payload: dict = {"chat_id": target_id, "sticker": file_id}
    if reply_to_id:
        payload["reply_parameters"] = {"message_id": int(reply_to_id)}
    if thread_id:
        payload["message_thread_id"] = int(thread_id)
    if emoji:
        payload["emoji"] = emoji

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
            resp = await client.post(
                f"{TELEGRAM_API_BASE}/bot{token}/sendSticker",
                json=payload,
            )
            data = resp.json() if resp.content else {}
            if data.get("ok"):
                return data.get("result")
            logger.warning("sendSticker failed: %s", data.get("description", ""))
            return None
    except Exception:
        logger.exception("sendSticker network error")
        return None


async def resolve_sticker_by_emoji(
    token: str,
    emoji: str,
    *,
    set_name: str | None = None,
) -> StickerInfo | None:
    if set_name:
        sticker_set = await get_sticker_set(token, set_name)
        if sticker_set:
            for sticker in sticker_set.stickers:
                if sticker.emoji == emoji:
                    return sticker
        return None

    cached = _sticker_cache.search_by_emoji(emoji)
    if cached:
        return cached[0]
    return None


def get_cache() -> TelegramStickerCache:
    return _sticker_cache