devplacepy/devplacepy/services/telegram/bridge.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import html
import logging
import time
from typing import Any

from devplacepy.database import get_table
from devplacepy.services.audit import record as audit
from devplacepy.services.manager import service_manager
from devplacepy.utils import is_admin, is_primary_admin

from . import store
from .format import markdown_to_telegram_html, split_for_telegram

logger = logging.getLogger("telegram.bridge")

TURN_TIMEOUT_SECONDS = 300.0
TYPING_INTERVAL_SECONDS = 4.0
PROGRESS_EDIT_INTERVAL_SECONDS = 1.5
ATTEMPT_LIMIT = 5
ATTEMPT_WINDOW_SECONDS = 600.0
THINKING_PLACEHOLDER = "<i>Devii is thinking...</i>"

WELCOME = (
    "Welcome to Devii on Telegram. To connect your DevPlace account, open your profile, "
    "request a Telegram pairing code, and send me the four digit code here."
)
PAIRED_TEMPLATE = "Paired. Hello {username}. Send me a message and I will help."
BAD_CODE = (
    "That code is invalid or expired. Request a fresh code from your DevPlace profile and "
    "send it here."
)
TOO_MANY = "Too many attempts. Wait a few minutes, request a new code, and try again."


class TelegramConnection:
    def __init__(
        self,
        service: Any,
        session: Any,
        chat_id: int,
        message_id: int | None,
        bridge: "TelegramBridge | None" = None,
    ) -> None:
        self._service = service
        self._session = session
        self._chat_id = chat_id
        self._message_id = message_id
        self._bridge = bridge
        self.done = asyncio.Event()
        self._finalized = False
        self._last_status = ""
        self._last_edit = 0.0
        self._typing_task = asyncio.create_task(self._typing_loop())

    async def send_json(self, payload: dict[str, Any]) -> None:
        kind = payload.get("type")
        if kind == "reply":
            await self._finalize(markdown_to_telegram_html(payload.get("text", "")) or "(no response)")
        elif kind == "error":
            await self._finalize(html.escape("Error: " + str(payload.get("text", ""))))
        elif kind in ("status", "trace"):
            await self._progress(payload)
        elif kind == "interaction":
            await self._handle_interaction(payload)
        elif kind in ("avatar", "client"):
            query_id = payload.get("id")
            if query_id:
                self._session.resolve_query(
                    query_id, {"error": "No browser is attached (Telegram session)."}
                )

    async def _handle_interaction(self, payload: dict[str, Any]) -> None:
        request_id = str(payload.get("id") or "")
        args = payload.get("args") if isinstance(payload.get("args"), dict) else payload
        interaction_id = str(
            args.get("interaction_id") or payload.get("interaction_id") or ""
        )
        plain = str(args.get("plain") or args.get("markdown") or "")
        widgets = args.get("widgets") if isinstance(args.get("widgets"), list) else []
        markup = _telegram_keyboard(interaction_id, widgets, args)
        text = html.escape(plain) if plain else "Please choose:"
        if markup is not None:
            await self._service.send(
                self._chat_id, f"<pre>{text}</pre>", reply_markup=markup
            )
        else:
            await self._service.send(self._chat_id, f"<pre>{text}</pre>")
        if self._bridge is None:
            self._session.resolve_interaction(
                request_id,
                {
                    "status": "error",
                    "interaction_id": interaction_id,
                    "error": "Telegram bridge unavailable for interaction reply.",
                    "values": {},
                },
            )
            return
        try:
            result = await self._bridge.wait_interaction_reply(
                self._chat_id, interaction_id, args
            )
        except asyncio.TimeoutError:
            result = {
                "status": "timeout",
                "interaction_id": interaction_id,
                "values": {},
            }
        self._session.resolve_interaction(request_id, result)

    async def _typing_loop(self) -> None:
        try:
            while not self.done.is_set():
                await self._service.chat_action(self._chat_id, "typing")
                await asyncio.sleep(TYPING_INTERVAL_SECONDS)
        except asyncio.CancelledError:
            pass
        except Exception:  # noqa: BLE001 - typing is cosmetic, never fail a turn over it
            pass

    async def _progress(self, payload: dict[str, Any]) -> None:
        if self._finalized or self._message_id is None:
            return
        if payload.get("type") == "trace":
            if payload.get("event") != "call":
                return
            text = f"Working... ({payload.get('name', '')})"
        else:
            text = str(payload.get("text", ""))
        if not text or text == self._last_status:
            return
        now = time.monotonic()
        if now - self._last_edit < PROGRESS_EDIT_INTERVAL_SECONDS:
            return
        self._last_status = text
        self._last_edit = now
        await self._service.edit(
            self._chat_id, self._message_id, f"<i>{html.escape(text)}</i>"
        )

    async def _finalize(self, html_text: str) -> None:
        if self._finalized:
            return
        self._finalized = True
        self._stop_typing()
        chunks = split_for_telegram(html_text)
        first = chunks[0] if chunks else "(no response)"
        if self._message_id is not None:
            ok = await self._service.edit(self._chat_id, self._message_id, first)
            if not ok:
                await self._service.send(self._chat_id, first)
        else:
            await self._service.send(self._chat_id, first)
        for chunk in chunks[1:]:
            await self._service.send(self._chat_id, chunk)
        self.done.set()

    def _stop_typing(self) -> None:
        if self._typing_task is not None and not self._typing_task.done():
            self._typing_task.cancel()


class TelegramBridge:
    def __init__(self, service: Any, max_concurrent: int = 8) -> None:
        self._service = service
        self._semaphore = asyncio.Semaphore(max(1, max_concurrent))
        self._chat_locks: dict[int, asyncio.Lock] = {}
        self._attempts: dict[int, list[float]] = {}
        self._pending_interactions: dict[int, dict[str, Any]] = {}

    def _chat_lock(self, chat_id: int) -> asyncio.Lock:
        lock = self._chat_locks.get(chat_id)
        if lock is None:
            if len(self._chat_locks) > 1024:
                for key in [k for k, v in self._chat_locks.items() if not v.locked()]:
                    del self._chat_locks[key]
            lock = asyncio.Lock()
            self._chat_locks[chat_id] = lock
        return lock

    async def wait_interaction_reply(
        self, chat_id: int, interaction_id: str, args: dict[str, Any]
    ) -> dict[str, Any]:
        loop = asyncio.get_event_loop()
        future: asyncio.Future = loop.create_future()
        timeout = float(args.get("timeout_sec") or 0) or 300.0
        self._pending_interactions[chat_id] = {
            "interaction_id": interaction_id,
            "future": future,
            "args": args,
        }
        try:
            return await asyncio.wait_for(future, timeout=timeout)
        finally:
            current = self._pending_interactions.get(chat_id)
            if current and current.get("future") is future:
                self._pending_interactions.pop(chat_id, None)

    def _resolve_pending(
        self, chat_id: int, result: dict[str, Any]
    ) -> bool:
        pending = self._pending_interactions.get(chat_id)
        if pending is None:
            return False
        future = pending.get("future")
        if future is not None and not future.done():
            future.set_result(result)
            return True
        return False

    async def handle_inbound(self, event: dict[str, Any]) -> None:
        kind = event.get("type") or "message"
        chat_id = int(event["chat_id"])
        from_id = int(event["from_id"])
        if kind == "callback":
            await self._handle_callback(event)
            return
        text = str(event.get("text", "")).strip()
        images = [img for img in event.get("images", []) if img]
        if chat_id in self._pending_interactions and text:
            from devplacepy.services.devii.interaction.parse import parse_plain_reply
            from devplacepy.services.devii.interaction.schema import InteractionRequest

            pending = self._pending_interactions[chat_id]
            try:
                request = InteractionRequest.model_validate(
                    (pending.get("args") or {}).get("request")
                    or {
                        "title": "?",
                        "widgets": (pending.get("args") or {}).get("widgets") or [],
                    }
                )
                status, values, error = parse_plain_reply(request, text)
                self._resolve_pending(
                    chat_id,
                    {
                        "status": status,
                        "interaction_id": pending.get("interaction_id"),
                        "values": values,
                        "error": error,
                        "meta": {"adapter": "telegram", "degraded": True},
                    },
                )
            except Exception:
                self._resolve_pending(
                    chat_id,
                    {
                        "status": "error",
                        "interaction_id": pending.get("interaction_id"),
                        "values": {},
                        "error": "Could not parse reply.",
                        "meta": {"adapter": "telegram", "degraded": True},
                    },
                )
            return
        link = store.user_for_chat(chat_id)
        if link is None:
            await self._handle_unpaired(chat_id, from_id, text)
            return
        if int(link.get("from_id", 0)) != from_id:
            return
        user = get_table("users").find_one(uid=link["user_uid"])
        if not user:
            await self._service.send(
                chat_id, "Your DevPlace account was not found. Please re-pair from your profile."
            )
            return
        async with self._chat_lock(chat_id):
            async with self._semaphore:
                await self._run_turn(chat_id, user, text, images)

    async def _handle_callback(self, event: dict[str, Any]) -> None:
        chat_id = int(event["chat_id"])
        data = str(event.get("data") or "")
        callback_id = str(event.get("callback_query_id") or "")
        if callback_id:
            try:
                await self._service.answer_callback(callback_id)
            except Exception:
                pass
        pending = self._pending_interactions.get(chat_id)
        if pending is None:
            return
        interaction_id = str(pending.get("interaction_id") or "")
        parsed = _parse_callback_data(data)
        if not parsed:
            return
        if parsed.get("i") and parsed["i"] not in interaction_id:
            if not interaction_id.endswith(parsed["i"]) and parsed["i"] not in interaction_id:
                pass
        status = parsed.get("s") or "submitted"
        if status == "cancel":
            self._resolve_pending(
                chat_id,
                {
                    "status": "cancelled",
                    "interaction_id": interaction_id,
                    "values": {},
                    "meta": {"adapter": "telegram", "degraded": True},
                },
            )
            return
        name = parsed.get("k")
        value = parsed.get("v")
        values: dict[str, Any] = {}
        if name is not None:
            if value in ("true", "false"):
                values[name] = value == "true"
            else:
                values[name] = value
        self._resolve_pending(
            chat_id,
            {
                "status": "submitted",
                "interaction_id": interaction_id,
                "values": values,
                "meta": {"adapter": "telegram", "degraded": True},
            },
        )

    async def _handle_unpaired(self, chat_id: int, from_id: int, text: str) -> None:
        if text.isdigit() and len(text) == 4:
            if self._too_many_attempts(chat_id):
                await self._service.send(chat_id, TOO_MANY)
                return
            user = store.verify_code(text, chat_id, from_id)
            if user:
                self._attempts.pop(chat_id, None)
                await self._service.send(
                    chat_id, PAIRED_TEMPLATE.format(username=user.get("username", "there"))
                )
                audit.record_system(
                    "telegram.pair.success",
                    actor_kind="user",
                    actor_uid=user["uid"],
                    actor_username=user.get("username", ""),
                    summary=f"Telegram paired for {user.get('username', user['uid'])}",
                    metadata={"chat_id": chat_id},
                )
                await self._publish_pairing(user["uid"], True)
                return
            self._register_attempt(chat_id)
            await self._service.send(chat_id, BAD_CODE)
            audit.record_system(
                "telegram.pair.failure",
                actor_kind="guest",
                result="failure",
                summary="Telegram pairing code rejected",
                metadata={"chat_id": chat_id},
            )
            return
        await self._service.send(chat_id, WELCOME)

    def _too_many_attempts(self, chat_id: int) -> bool:
        now = time.monotonic()
        recent = [t for t in self._attempts.get(chat_id, []) if now - t < ATTEMPT_WINDOW_SECONDS]
        if recent:
            self._attempts[chat_id] = recent
        else:
            self._attempts.pop(chat_id, None)
        return len(recent) >= ATTEMPT_LIMIT

    def _register_attempt(self, chat_id: int) -> None:
        self._attempts.setdefault(chat_id, []).append(time.monotonic())

    async def _publish_pairing(self, user_uid: str, paired: bool) -> None:
        try:
            from devplacepy.services.pubsub import publish

            await publish(f"user.{user_uid}.telegram", {"paired": paired})
        except Exception:  # noqa: BLE001 - live pairing update is best-effort
            pass

    async def _run_turn(
        self, chat_id: int, user: dict[str, Any], text: str, images: list[str]
    ) -> None:
        devii = service_manager.get_service("devii")
        if devii is None or not devii.is_enabled():
            await self._service.send(chat_id, "Devii is currently unavailable.")
            return
        owner_id = user["uid"]
        owner_is_admin = is_admin(user)
        owner_is_primary_admin = is_primary_admin(user)
        if devii.quota_exceeded("user", owner_id, owner_is_admin):
            limit = devii.daily_limit_for("user", owner_is_admin)
            audit.record_system(
                "ai.quota.exceeded",
                actor_kind="user",
                actor_uid=owner_id,
                actor_username=user.get("username", ""),
                actor_role="admin" if owner_is_admin else "member",
                origin="telegram",
                via_agent=1,
                result="denied",
                summary=f"Telegram AI request by {user.get('username', owner_id)} blocked - 24h quota reached",
                metadata={"limit_usd": limit},
            )
            await self._service.send(
                chat_id, "Your daily AI quota is reached (100%). Please try again later."
            )
            return
        session = devii.hub().get_or_create(
            "user",
            owner_id,
            user.get("username", ""),
            user.get("api_key", ""),
            devii.instance_base_url(),
            is_admin=owner_is_admin,
            is_primary_admin=owner_is_primary_admin,
            channel="telegram",
        )
        session.set_timezone(user.get("timezone") or "")
        placeholder_id = await self._service.send(chat_id, THINKING_PLACEHOLDER)
        connection = TelegramConnection(
            self._service, session, chat_id, placeholder_id, bridge=self
        )
        session.interaction_broker.bind(present_telegram=None)
        session.attach(connection)
        content, audit_text = _build_content(text, images)
        started = time.monotonic()
        try:
            session.spawn_turn(content, audit_text=audit_text)
            await asyncio.wait_for(connection.done.wait(), timeout=TURN_TIMEOUT_SECONDS)
        except asyncio.TimeoutError:
            await connection._finalize(
                html.escape("Devii took too long to respond. Please try again.")
            )
        except Exception:  # noqa: BLE001 - never let one turn crash the bridge
            logger.exception("Telegram turn failed for %s", owner_id)
            await connection._finalize(html.escape("Something went wrong handling your message."))
        finally:
            session.detach(connection)
            self._service.record_latency(time.monotonic() - started)


def _telegram_keyboard(
    interaction_id: str, widgets: list[Any], args: dict[str, Any]
) -> dict[str, Any] | None:
    short = interaction_id.replace("ix-", "")[:8]
    rows: list[list[dict[str, str]]] = []
    inputs = [
        w
        for w in widgets
        if isinstance(w, dict) and w.get("type") not in ("//", "group", None)
    ]
    if len(inputs) == 1 and inputs[0].get("type") == "confirm":
        name = str(inputs[0].get("name") or "confirm")
        submit = str(args.get("submit_label") or "Confirm")[:20]
        cancel = str(args.get("cancel_label") or "Cancel")[:20]
        rows.append(
            [
                {
                    "text": submit,
                    "callback_data": f"i={short};k={name};v=true"[:64],
                },
                {
                    "text": cancel,
                    "callback_data": f"i={short};s=cancel"[:64],
                },
            ]
        )
        return {"inline_keyboard": rows}
    if len(inputs) == 1 and inputs[0].get("type") in ("choice", "select"):
        name = str(inputs[0].get("name") or "choice")
        options = inputs[0].get("options") or []
        for option in options[:12]:
            if not isinstance(option, dict):
                continue
            value = str(option.get("value") or "")[:20]
            label = str(option.get("label") or value)[:40]
            rows.append(
                [
                    {
                        "text": label,
                        "callback_data": f"i={short};k={name};v={value}"[:64],
                    }
                ]
            )
        if args.get("cancelable", True):
            rows.append(
                [{"text": "Cancel", "callback_data": f"i={short};s=cancel"[:64]}]
            )
        return {"inline_keyboard": rows} if rows else None
    return None


def _parse_callback_data(data: str) -> dict[str, str]:
    out: dict[str, str] = {}
    for part in str(data or "").split(";"):
        if "=" not in part:
            continue
        key, value = part.split("=", 1)
        out[key.strip()] = value.strip()
    return out


def _build_content(text: str, images: list[str]) -> tuple[Any, str]:
    if not images:
        return text, text
    prompt = text or "Look at the attached image and respond."
    content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
    for uri in images:
        content.append({"type": "image_url", "image_url": {"url": uri}})
    audit_text = f"{text} [{len(images)} image(s)]".strip()
    return content, audit_text