devplacepy/devplacepy/services/telegram/worker.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import base64
import json
import os
import re
import sys
import time
from typing import Any, Awaitable, Callable

from .backend import HttpTelegramBackend, TelegramBackend

DEFAULT_POLL_TIMEOUT = 25
DEFAULT_IMAGE_MAX_BYTES = 2_000_000
DEFAULT_MIN_SEND_INTERVAL = 0.05
CONFLICT_BACKOFF_SECONDS = 3.0
ERROR_BACKOFF_SECONDS = 2.0
_TAG = re.compile(r"<[^>]+>")


def _guess_mime(file_path: str) -> str:
    lowered = file_path.lower()
    if lowered.endswith(".png"):
        return "image/png"
    if lowered.endswith(".webp"):
        return "image/webp"
    if lowered.endswith(".gif"):
        return "image/gif"
    return "image/jpeg"


def _strip_tags(text: str) -> str:
    import html as _html

    return _html.unescape(_TAG.sub("", text))


def _entity_error(resp: dict[str, Any]) -> bool:
    return "can't parse entities" in str(resp.get("description", "")).lower()


class TelegramWorker:
    def __init__(
        self,
        backend: TelegramBackend,
        on_emit: Callable[[dict[str, Any]], Awaitable[None]],
        *,
        poll_timeout: int = DEFAULT_POLL_TIMEOUT,
        image_max_bytes: int = DEFAULT_IMAGE_MAX_BYTES,
        min_send_interval: float = DEFAULT_MIN_SEND_INTERVAL,
    ) -> None:
        self._backend = backend
        self._on_emit = on_emit
        self._poll_timeout = poll_timeout
        self._image_max_bytes = image_max_bytes
        self._min_send_interval = min_send_interval
        self._offset = 0
        self._stop = asyncio.Event()
        self._gate = asyncio.Lock()
        self._last_call = 0.0

    async def _emit(self, frame: dict[str, Any]) -> None:
        await self._on_emit(frame)

    async def _log(self, line: str) -> None:
        await self._emit({"type": "log", "line": line})

    async def _throttle(self) -> None:
        async with self._gate:
            wait = self._min_send_interval - (time.monotonic() - self._last_call)
            if wait > 0:
                await asyncio.sleep(wait)
            self._last_call = time.monotonic()

    async def poll_once(self) -> list[dict[str, Any]]:
        resp = await self._backend.get_updates(self._offset, self._poll_timeout)
        if not resp.get("ok"):
            if resp.get("error_code") == 409:
                await self._log("getUpdates conflict (409); backing off")
                await asyncio.sleep(CONFLICT_BACKOFF_SECONDS)
            else:
                await self._log(f"getUpdates error: {resp.get('description', 'unknown')}")
                await asyncio.sleep(ERROR_BACKOFF_SECONDS)
            return []
        updates = resp.get("result", []) or []
        for update in updates:
            self._offset = max(self._offset, int(update.get("update_id", 0)) + 1)
        return updates

    async def process_update(self, update: dict[str, Any]) -> None:
        callback = update.get("callback_query")
        if isinstance(callback, dict):
            await self._process_callback(callback)
            return
        message = update.get("message")
        if not isinstance(message, dict):
            return
        chat = message.get("chat", {})
        if chat.get("type") != "private":
            return
        chat_id = chat.get("id")
        from_id = (message.get("from") or {}).get("id")
        if chat_id is None or from_id is None:
            return
        text = message.get("text") or message.get("caption") or ""
        images: list[str] = []
        photos = message.get("photo")
        if isinstance(photos, list) and photos:
            largest = max(photos, key=lambda p: p.get("file_size", 0) or 0)
            data_uri = await self._fetch_image(largest.get("file_id"))
            if data_uri:
                images.append(data_uri)
        document = message.get("document")
        if isinstance(document, dict) and str(document.get("mime_type", "")).startswith(
            "image/"
        ):
            data_uri = await self._fetch_image(
                document.get("file_id"), document.get("mime_type")
            )
            if data_uri:
                images.append(data_uri)
        await self._emit(
            {
                "type": "message",
                "chat_id": chat_id,
                "from_id": from_id,
                "text": text,
                "images": images,
            }
        )

    async def _process_callback(self, callback: dict[str, Any]) -> None:
        message = callback.get("message") or {}
        chat = message.get("chat") or {}
        if chat.get("type") != "private":
            return
        chat_id = chat.get("id")
        from_id = (callback.get("from") or {}).get("id")
        if chat_id is None or from_id is None:
            return
        await self._emit(
            {
                "type": "callback",
                "chat_id": chat_id,
                "from_id": from_id,
                "data": str(callback.get("data") or ""),
                "callback_query_id": str(callback.get("id") or ""),
                "message_id": message.get("message_id"),
            }
        )

    async def _fetch_image(self, file_id: str | None, mime: str | None = None) -> str:
        if not file_id:
            return ""
        info = await self._backend.get_file(file_id)
        if not info.get("ok"):
            return ""
        file_path = (info.get("result") or {}).get("file_path", "")
        if not file_path:
            return ""
        data = await self._backend.download(file_path)
        if not data:
            return ""
        if len(data) > self._image_max_bytes:
            await self._log(f"image skipped ({len(data)} bytes over cap)")
            return ""
        resolved_mime = mime or _guess_mime(file_path)
        encoded = base64.b64encode(data).decode("ascii")
        return f"data:{resolved_mime};base64,{encoded}"

    async def handle_command(self, command: dict[str, Any]) -> None:
        kind = command.get("cmd")
        try:
            if kind == "send":
                await self._send(command)
            elif kind == "edit":
                await self._edit(command)
            elif kind == "chat_action":
                await self._chat_action(command)
            elif kind == "answer_callback":
                await self._answer_callback(command)
        except Exception as exc:  # noqa: BLE001 - one bad command must not kill the worker
            await self._log(f"command {kind} failed: {exc}")
            req_id = command.get("req_id")
            if req_id is not None:
                await self._emit({"type": "result", "req_id": req_id, "ok": False})

    async def _send(self, command: dict[str, Any]) -> None:
        chat_id = command["chat_id"]
        text = str(command.get("text", ""))
        parse_mode = command.get("parse_mode")
        reply_markup = command.get("reply_markup")
        await self._throttle()
        resp = await self._backend.send_message(
            chat_id, text, parse_mode, reply_markup=reply_markup
        )
        resp = await self._retry_if_needed(
            resp,
            lambda mode: self._backend.send_message(
                chat_id, text, mode, reply_markup=reply_markup
            ),
            parse_mode,
        )
        await self._emit_result(command.get("req_id"), resp)

    async def _edit(self, command: dict[str, Any]) -> None:
        chat_id = command["chat_id"]
        message_id = command["message_id"]
        text = str(command.get("text", ""))
        parse_mode = command.get("parse_mode")
        reply_markup = command.get("reply_markup")
        await self._throttle()
        resp = await self._backend.edit_message_text(
            chat_id, message_id, text, parse_mode, reply_markup=reply_markup
        )
        resp = await self._retry_if_needed(
            resp,
            lambda mode: self._backend.edit_message_text(
                chat_id, message_id, text, mode, reply_markup=reply_markup
            ),
            parse_mode,
        )
        await self._emit_result(command.get("req_id"), resp)

    async def _answer_callback(self, command: dict[str, Any]) -> None:
        callback_query_id = str(command.get("callback_query_id") or "")
        text = command.get("text")
        await self._throttle()
        resp = await self._backend.answer_callback_query(callback_query_id, text)
        await self._emit_result(command.get("req_id"), resp)

    async def _chat_action(self, command: dict[str, Any]) -> None:
        await self._throttle()
        await self._backend.send_chat_action(command["chat_id"], command.get("action", "typing"))

    async def _retry_if_needed(
        self,
        resp: dict[str, Any],
        again: Callable[[str | None], Awaitable[dict[str, Any]]],
        parse_mode: str | None,
    ) -> dict[str, Any]:
        if resp.get("ok"):
            return resp
        if resp.get("error_code") == 429:
            retry_after = int((resp.get("parameters") or {}).get("retry_after", 1))
            await asyncio.sleep(min(retry_after, 30) + 0.25)
            await self._throttle()
            resp = await again(parse_mode)
            if resp.get("ok"):
                return resp
        if parse_mode and _entity_error(resp):
            await self._throttle()
            resp = await again(None)
        return resp

    async def _emit_result(self, req_id: Any, resp: dict[str, Any]) -> None:
        if req_id is None:
            if not resp.get("ok"):
                await self._log(f"send failed: {resp.get('description', 'unknown')}")
            return
        message_id = None
        if resp.get("ok"):
            message_id = (resp.get("result") or {}).get("message_id")
        await self._emit(
            {"type": "result", "req_id": req_id, "ok": bool(resp.get("ok")), "message_id": message_id}
        )

    async def _read_commands(self) -> None:
        loop = asyncio.get_event_loop()
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
        while not self._stop.is_set():
            line = await reader.readline()
            if not line:
                break
            stripped = line.strip()
            if not stripped:
                continue
            try:
                command = json.loads(stripped)
            except (ValueError, TypeError):
                continue
            asyncio.create_task(self.handle_command(command))
        self._stop.set()

    async def _poll_loop(self) -> None:
        while not self._stop.is_set():
            try:
                updates = await self.poll_once()
            except Exception as exc:  # noqa: BLE001 - keep polling through transient failures
                await self._log(f"poll loop error: {exc}")
                await asyncio.sleep(ERROR_BACKOFF_SECONDS)
                continue
            for update in updates:
                try:
                    await self.process_update(update)
                except Exception as exc:  # noqa: BLE001 - one bad update never stops the loop
                    await self._log(f"update processing error: {exc}")

    async def run(self) -> None:
        await self._emit({"type": "ready"})
        await self._log("Telegram worker started")
        reader_task = asyncio.create_task(self._read_commands())
        poll_task = asyncio.create_task(self._poll_loop())
        await self._stop.wait()
        poll_task.cancel()
        reader_task.cancel()
        await asyncio.gather(poll_task, reader_task, return_exceptions=True)


async def _stdout_emitter() -> Callable[[dict[str, Any]], Awaitable[None]]:
    lock = asyncio.Lock()

    async def emit(frame: dict[str, Any]) -> None:
        async with lock:
            sys.stdout.write(json.dumps(frame, ensure_ascii=False) + "\n")
            sys.stdout.flush()

    return emit


async def _main() -> None:
    token = os.environ.get("TELEGRAM_BOT_TOKEN", "").strip()
    emit = await _stdout_emitter()
    if not token:
        await emit({"type": "log", "line": "TELEGRAM_BOT_TOKEN is not set; exiting"})
        return
    poll_timeout = int(os.environ.get("TELEGRAM_POLL_TIMEOUT", DEFAULT_POLL_TIMEOUT))
    image_cap = int(os.environ.get("TELEGRAM_IMAGE_MAX_BYTES", DEFAULT_IMAGE_MAX_BYTES))
    backend = HttpTelegramBackend(token)
    worker = TelegramWorker(
        backend, emit, poll_timeout=poll_timeout, image_max_bytes=image_cap
    )
    await worker.run()


def main() -> None:
    try:
        asyncio.run(_main())
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()