PTSRaiAgent/app/channels/telegram.py

"""Telegram channel (#11). Prototype on synthetic data only; prod channel is MAX / Gosuslugi.

A thin layer over the dialog engine: update → `get_engine().handle(token, event)` → render `Reply` list.
- `/start <code>` (deep link from the doctor's QR) → `integrations.identity.verify(code)` → bind chat to token.
- Text / button presses of a bound chat → engine events; unbound chats get a fixed "ask your doctor" text.
- `TelegramOutbox` delivers system-initiated replies (scheduler pushes) to the bound chat.

The channel is the registration boundary, so it is allowed to use the identity DB (chat id ↔ token, guardrail #5).
Logs carry only event types and short token prefixes — never chat ids or message text (guardrail #8).
Without `TG_BOT_TOKEN` `start()` is a no-op.
"""

import asyncio
import contextlib
import logging
from functools import lru_cache
from pathlib import Path
from typing import Protocol

import yaml
from aiogram import Bot, Dispatcher, F, Router
from aiogram.filters import CommandObject, CommandStart
from aiogram.types import CallbackQuery, InlineKeyboardButton, InlineKeyboardMarkup, Message
from pydantic import BaseModel

from app.config import get_settings
from app.db import identity_repo
from app.db.session import get_identity_session
from app.dialog.contracts import Event, Reply, get_engine
from app.integrations.base import Outbox, integrations

log = logging.getLogger("app.channels.telegram")

CALLBACK_DATA_MAX = 64  # Telegram limit for callback_data, bytes
MESSAGE_MAX = 4096  # Telegram limit for message text, characters


class TelegramTexts(BaseModel):
    not_connected: str
    code_invalid: str
    error: str
    unsupported: str


@lru_cache
def load_texts(path: Path | None = None) -> TelegramTexts:
    path = path or get_settings().config_dir / "texts" / "telegram.yaml"
    return TelegramTexts.model_validate(yaml.safe_load(path.read_text(encoding="utf-8")))


class Sender(Protocol):
    """The part of `aiogram.Bot` the channel uses; tests pass a recorder."""

    async def send_message(
        self, chat_id: int, text: str, parse_mode: str | None = None, reply_markup: InlineKeyboardMarkup | None = None
    ) -> object: ...


def _short(token: str) -> str:
    return token[:6]


# ----------------------------------------------------------------- rendering


def render_markup(reply: Reply) -> InlineKeyboardMarkup | None:
    """Abstract button rows → inline keyboard. Over-long `data` is dropped (Telegram rejects the whole message)."""
    rows: list[list[InlineKeyboardButton]] = []
    for row in reply.buttons:
        buttons = []
        for button in row:
            if len(button.data.encode()) > CALLBACK_DATA_MAX:
                log.warning("button data too long, dropped len=%d", len(button.data.encode()))
                continue
            buttons.append(InlineKeyboardButton(text=button.label, callback_data=button.data))
        if buttons:
            rows.append(buttons)
    return InlineKeyboardMarkup(inline_keyboard=rows) if rows else None


def split_text(text: str, limit: int = MESSAGE_MAX) -> list[str]:
    """Split by paragraphs/lines so long replies (safety plan) fit Telegram's limit; hard cut as a last resort."""
    chunks: list[str] = []
    rest = text
    while len(rest) > limit:
        cut = rest.rfind("\n", 0, limit)
        if cut <= 0:
            cut = limit
        chunks.append(rest[:cut].rstrip())
        rest = rest[cut:].lstrip("\n")
    chunks.append(rest)
    return [c for c in chunks if c] or [""]


async def send_replies(sender: Sender, chat_id: int, replies: list[Reply]) -> None:
    """Send replies in order; text is sent verbatim (no parse_mode: crisis scripts must not be altered)."""
    for reply in replies:
        parts = split_text(reply.text)
        markup = render_markup(reply)
        for i, part in enumerate(parts):
            # Buttons go under the last chunk so they stay next to the question.
            await sender.send_message(
                chat_id, part, parse_mode=None, reply_markup=markup if i == len(parts) - 1 else None
            )


# ----------------------------------------------------------------- chat ↔ token


class ChatDirectory:
    """chat id ↔ token on top of the identity DB, with an in-process cache (a chat is bound once)."""

    def __init__(self) -> None:
        self._token_by_chat: dict[int, str] = {}
        self._chat_by_token: dict[str, int] = {}

    def _remember(self, chat_id: int, token: str) -> None:
        self._token_by_chat[chat_id] = token
        self._chat_by_token[token] = chat_id

    async def token_for(self, chat_id: int) -> str | None:
        if chat_id in self._token_by_chat:
            return self._token_by_chat[chat_id]
        async with get_identity_session() as s:
            patient = await identity_repo.get_patient_by_chat(s, chat_id)
        if patient is None:
            return None
        self._remember(chat_id, patient.token)
        return patient.token

    async def chat_for(self, token: str) -> int | None:
        if token in self._chat_by_token:
            return self._chat_by_token[token]
        async with get_identity_session() as s:
            patient = await identity_repo.get_patient(s, token)
        if patient is None or patient.tg_chat_id is None:
            return None
        self._remember(patient.tg_chat_id, token)
        return patient.tg_chat_id

    async def bind(self, chat_id: int, token: str) -> bool:
        async with get_identity_session() as s:
            # A chat id is unique in the identity DB: rebinding to another patient frees it first.
            previous = await identity_repo.get_patient_by_chat(s, chat_id)
            if previous is not None and previous.token != token:
                previous.tg_chat_id = None
                await s.flush()
            patient = await identity_repo.bind_chat(s, token=token, tg_chat_id=chat_id)
        if patient is None:
            return False
        stale = self._token_by_chat.pop(chat_id, None)
        if stale is not None:
            self._chat_by_token.pop(stale, None)
        self._remember(chat_id, token)
        return True


# ----------------------------------------------------------------- channel logic (no aiogram types)


class TelegramChannel:
    """Channel logic decoupled from aiogram updates, so it is testable with a fake sender."""

    def __init__(self, sender: Sender, directory: ChatDirectory | None = None) -> None:
        self.sender = sender
        self.directory = directory or ChatDirectory()

    async def _say(self, chat_id: int, text: str) -> None:
        await send_replies(self.sender, chat_id, [Reply(text=text, kind="info")])

    async def _dispatch(self, chat_id: int, token: str, event: Event) -> None:
        try:
            replies = await get_engine().handle(token, event)
        except Exception:
            # The engine already logs details; here only the fact, without text or chat id.
            log.exception("engine failed tok=%s event=%s", _short(token), event.type)
            replies = [Reply(text=load_texts().error, kind="info")]
        await send_replies(self.sender, chat_id, replies)

    async def on_start(self, chat_id: int, code: str | None) -> None:
        texts = load_texts()
        if code:
            person = await integrations.identity.verify(code)
            if person is None:
                log.info("start: code rejected")
                await self._say(chat_id, texts.code_invalid)
                return
            if not await self.directory.bind(chat_id, person.token):
                log.warning("start: verified token has no identity record tok=%s", _short(person.token))
                await self._say(chat_id, texts.code_invalid)
                return
            log.info("start: chat bound tok=%s", _short(person.token))
            await self._dispatch(chat_id, person.token, Event(type="start", params={"code": code}))
            return
        token = await self.directory.token_for(chat_id)
        if token is None:
            await self._say(chat_id, texts.not_connected)
            return
        await self._dispatch(chat_id, token, Event(type="start"))

    async def on_text(self, chat_id: int, text: str | None) -> None:
        texts = load_texts()
        token = await self.directory.token_for(chat_id)
        if token is None:
            await self._say(chat_id, texts.not_connected)
            return
        if not text:
            await self._say(chat_id, texts.unsupported)
            return
        await self._dispatch(chat_id, token, Event(type="text", text=text))

    async def on_button(self, chat_id: int, data: str) -> None:
        token = await self.directory.token_for(chat_id)
        if token is None:
            await self._say(chat_id, load_texts().not_connected)
            return
        await self._dispatch(chat_id, token, Event(type="button", button=data))


class TelegramOutbox:
    """Outbox for scheduler pushes: bound chats get Telegram messages, the rest go to the wrapped outbox
    (the web demo console polls it), so a patient connected via the console keeps working."""

    def __init__(self, sender: Sender, directory: ChatDirectory, fallback: Outbox) -> None:
        self.sender = sender
        self.directory = directory
        self.fallback = fallback

    async def push(self, token: str, replies: list[Reply]) -> None:
        chat_id = await self.directory.chat_for(token)
        if chat_id is None:
            await self.fallback.push(token, replies)
            return
        await send_replies(self.sender, chat_id, replies)


# ----------------------------------------------------------------- aiogram wiring


def build_router(channel: TelegramChannel) -> Router:
    router = Router(name="patient")

    @router.message(CommandStart())
    async def _start(message: Message, command: CommandObject) -> None:
        await channel.on_start(message.chat.id, (command.args or "").strip() or None)

    @router.message(F.chat.type == "private")
    async def _text(message: Message) -> None:
        if message.bot is not None and message.text:
            await message.bot.send_chat_action(message.chat.id, "typing")
        await channel.on_text(message.chat.id, message.text)

    @router.callback_query(F.data)
    async def _button(query: CallbackQuery) -> None:
        await query.answer()
        msg = query.message
        if not isinstance(msg, Message):  # inaccessible (too old) message
            return
        # Remove the keyboard so the same answer cannot be submitted twice.
        try:
            await msg.edit_reply_markup(reply_markup=None)
        except Exception:
            log.debug("could not remove keyboard")
        await channel.on_button(msg.chat.id, query.data or "")

    return router


_bot: Bot | None = None
_dp: Dispatcher | None = None
_task: asyncio.Task[None] | None = None
_prev_outbox: Outbox | None = None


async def start() -> None:
    """Start polling if TG_BOT_TOKEN is set; no-op otherwise."""
    global _bot, _dp, _task, _prev_outbox
    token = get_settings().tg_bot_token
    if not token or _task is not None:
        if not token:
            log.info("telegram disabled: TG_BOT_TOKEN is empty")
        return
    bot = Bot(token)
    channel = TelegramChannel(bot)
    dp = Dispatcher()
    dp.include_router(build_router(channel))
    _prev_outbox = integrations.outbox
    integrations.outbox = TelegramOutbox(bot, channel.directory, fallback=_prev_outbox)
    _bot, _dp = bot, dp
    _task = asyncio.create_task(dp.start_polling(bot, handle_signals=False), name="telegram-polling")
    log.info("telegram polling started")


async def stop() -> None:
    """Stop polling and restore the previous outbox."""
    global _bot, _dp, _task, _prev_outbox
    if _task is None:
        return
    if _dp is not None:
        with contextlib.suppress(RuntimeError):  # polling never got to start
            await _dp.stop_polling()
    _task.cancel()
    with contextlib.suppress(asyncio.CancelledError, Exception):
        await _task
    if _bot is not None:
        await _bot.session.close()
    if _prev_outbox is not None and isinstance(integrations.outbox, TelegramOutbox):
        integrations.outbox = _prev_outbox
    _bot = _dp = _task = _prev_outbox = None
    log.info("telegram polling stopped")