beavergram/backend/src/utils/read/watches.py

import json
from datetime import UTC, datetime
from typing import Any

import asyncpg

from utils.read import chats as chats_read
from utils.read import peers as peers_read
from utils.read.models import (
    AlertItem,
    AlertMessage,
    AlertPage,
    AlertView,
    AlertWatch,
    ChatListItem,
    ChatRef,
    PeerView,
    PersonRef,
    WatchItem,
    WatchView,
)

WATCHES_CHANGED_CHANNEL = "watches_changed"
ALERTS_CHANGED_CHANNEL = "alerts_changed"
ALERT_PAGE = 50
MAX_ALERT_PAGE = 200

_WATCH_COLS = "id, account_id, kind, name, params, enabled, created_at, updated_at"
_ALERT_COLS = "id, account_id, watch_id, ts, payload, seen, created_at"
_LEGACY_EVENTS = {"keyword": "message", "peer_online": "online"}

_INSERT_ALERT = f"""
WITH alert AS (
    INSERT INTO alerts (account_id, watch_id, ts, payload, seen, dedup_key)
    VALUES ($1, $2, $3, $4::jsonb, false, $5)
    ON CONFLICT (watch_id, dedup_key) WHERE dedup_key IS NOT NULL DO NOTHING
    RETURNING {_ALERT_COLS}
), hooks AS (
    INSERT INTO webhook_deliveries (webhook_id, alert_id, event)
    SELECT w.id, alert.id, 'alert' FROM alert
    JOIN webhooks w ON w.account_id = alert.account_id AND w.enabled
        AND (w.watch_ids IS NULL OR alert.watch_id = ANY(w.watch_ids))
)
SELECT * FROM alert
"""  # noqa: S608

_WATCH_STATS = """
SELECT watch_id, count(*) AS count, max(ts) AS last FROM alerts
WHERE watch_id = ANY($1::int[]) GROUP BY watch_id
"""

_ALERT_ROWS = """
SELECT a.id, a.account_id, a.ts, a.seen, a.payload,
       w.id AS watch_id, w.kind AS watch_kind, w.name AS watch_name,
       w.params AS watch_params,
       m.text AS stored_text,
       m.sender_id AS stored_sender, m.date AS stored_date,
       m.deleted_at IS NOT NULL AS stored_deleted,
       m.edited_at IS NOT NULL AS stored_edited
FROM alerts a
JOIN watches w ON w.id = a.watch_id
LEFT JOIN LATERAL (
    SELECT text, sender_id, date, deleted_at, edited_at
    FROM messages
    WHERE account_id = a.account_id
      AND chat_id = (a.payload->>'chat_id')::bigint
      AND message_id = (a.payload->>'message_id')::bigint
    ORDER BY date DESC LIMIT 1
) m ON true
"""


def _to_watch(row: asyncpg.Record) -> WatchView:
    data = dict(row)
    data["params"] = json.loads(data["params"])
    return WatchView(**data)


def _to_alert(row: asyncpg.Record) -> AlertView:
    data = dict(row)
    data["payload"] = json.loads(data["payload"])
    return AlertView(**data)


async def _notify(pool: asyncpg.Pool, channel: str, payload: dict[str, Any]) -> None:
    await pool.execute("SELECT pg_notify($1, $2)", channel, json.dumps(payload))


async def list_watches(pool: asyncpg.Pool, account_id: int) -> list[WatchView]:
    rows = await pool.fetch(
        f"SELECT {_WATCH_COLS} FROM watches WHERE account_id = $1 "  # noqa: S608
        "ORDER BY id DESC",
        account_id,
    )
    return [_to_watch(row) for row in rows]


async def get_watch(pool: asyncpg.Pool, watch_id: int) -> WatchView | None:
    row = await pool.fetchrow(
        f"SELECT {_WATCH_COLS} FROM watches WHERE id = $1",  # noqa: S608
        watch_id,
    )
    return _to_watch(row) if row else None


async def create_watch(  # noqa: PLR0913
    pool: asyncpg.Pool,
    account_id: int,
    kind: str,
    params: dict[str, Any],
    *,
    name: str | None = None,
    enabled: bool = True,
) -> WatchView:
    row = await pool.fetchrow(
        "INSERT INTO watches (account_id, kind, name, params, enabled) "  # noqa: S608
        f"VALUES ($1, $2, $3, $4::jsonb, $5) RETURNING {_WATCH_COLS}",
        account_id,
        kind,
        name,
        json.dumps(params),
        enabled,
    )
    await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
    return _to_watch(row)


async def update_watch(  # noqa: PLR0913
    pool: asyncpg.Pool,
    watch_id: int,
    kind: str,
    params: dict[str, Any],
    *,
    name: str | None,
    enabled: bool,
) -> WatchView | None:
    row = await pool.fetchrow(
        "UPDATE watches SET kind = $2, name = $3, params = $4::jsonb, "  # noqa: S608
        "enabled = $5, updated_at = now() "
        f"WHERE id = $1 RETURNING {_WATCH_COLS}",
        watch_id,
        kind,
        name,
        json.dumps(params),
        enabled,
    )
    if row is None:
        return None
    await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
    return _to_watch(row)


async def delete_watch(pool: asyncpg.Pool, watch_id: int) -> bool:
    account_id = await pool.fetchval(
        "DELETE FROM watches WHERE id = $1 RETURNING account_id", watch_id
    )
    if account_id is None:
        return False
    await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
    await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id})
    return True


def _chat_ref(chat_id: int, chat: ChatListItem | None, name: str | None) -> ChatRef:
    if chat is None:
        kind = "private" if chat_id > 0 else "group"
        return ChatRef(id=chat_id, title=name, kind=kind, has_avatar=False)
    kind = "channel" if chat.is_broadcast else chat.kind
    return ChatRef(
        id=chat_id, title=chat.title or name, kind=kind, has_avatar=chat.has_avatar
    )


def _peer_name(peer: PeerView) -> str | None:
    full = " ".join(part for part in (peer.first_name, peer.last_name) if part)
    return full or peer.username


class Refs:
    def __init__(self, chats: list[ChatListItem], people: list[PeerView]) -> None:
        self._chats = {chat.chat_id: chat for chat in chats}
        self._people = {peer.peer_id: peer for peer in people}

    def chat(self, chat_id: int | None) -> ChatRef | None:
        if chat_id is None:
            return None
        peer = self._people.get(chat_id)
        name = _peer_name(peer) if peer else None
        return _chat_ref(chat_id, self._chats.get(chat_id), name)

    def person(self, person_id: int | None) -> PersonRef | None:
        if person_id is None:
            return None
        peer = self._people.get(person_id)
        chat = self._chats.get(person_id)
        name = (_peer_name(peer) if peer else None) or (chat.title if chat else None)
        has_avatar = (peer.has_avatar if peer else False) or bool(
            chat and chat.has_avatar
        )
        return PersonRef(id=person_id, name=name, has_avatar=has_avatar)


async def load_refs(
    pool: asyncpg.Pool, account_id: int, chat_ids: set[int], person_ids: set[int]
) -> Refs:
    chats = await chats_read.get_chats(pool, account_id, sorted(chat_ids | person_ids))
    people = await peers_read.get_peers(
        pool, account_id, sorted(person_ids | {i for i in chat_ids if i > 0})
    )
    return Refs(chats, people)


def _int(value: object) -> int | None:
    return value if isinstance(value, int) and not isinstance(value, bool) else None


def _watch_person(params: dict[str, Any]) -> int | None:
    return _int(params.get("sender_id")) or _int(params.get("peer_id"))


async def describe_watches(
    pool: asyncpg.Pool, account_id: int, watches: list[WatchView]
) -> list[WatchItem]:
    targets = [
        (_int(watch.params.get("chat_id")), _watch_person(watch.params))
        for watch in watches
    ]
    refs = await load_refs(
        pool,
        account_id,
        {chat for chat, _ in targets if chat is not None},
        {person for _, person in targets if person is not None},
    )
    stats = {
        row["watch_id"]: row
        for row in await pool.fetch(_WATCH_STATS, [watch.id for watch in watches])
    }
    return [
        WatchItem(
            **watch.model_dump(),
            chat=refs.chat(chat),
            person=refs.person(person),
            alert_count=stats[watch.id]["count"] if watch.id in stats else 0,
            last_alert_at=stats[watch.id]["last"] if watch.id in stats else None,
        )
        for watch, (chat, person) in zip(watches, targets, strict=True)
    ]


def _path(chat_id: int | None, message_id: int | None, peer_id: int | None) -> str:
    if chat_id is not None and message_id is not None:
        return f"/app/{chat_id}?message={message_id}"
    if peer_id is not None:
        return f"/app/{peer_id}"
    return "/app"


def _alert_item(row: asyncpg.Record, refs: Refs) -> AlertItem:
    payload = json.loads(row["payload"])
    event = payload.get("event") or _LEGACY_EVENTS.get(row["watch_kind"], "message")
    chat_id = _int(payload.get("chat_id"))
    message_id = _int(payload.get("message_id"))
    peer_id = _int(payload.get("peer_id"))
    message = None
    if message_id is not None:
        date = row["stored_date"] or (
            datetime.fromisoformat(payload["date"]) if payload.get("date") else None
        )
        message = AlertMessage(
            id=message_id,
            text=payload.get("text") or row["stored_text"],
            sender=refs.person(_int(payload.get("sender_id")) or row["stored_sender"]),
            date=date,
            deleted=bool(row["stored_deleted"]) or event == "delete",
            edited=bool(row["stored_edited"]) or event == "edit",
        )
    return AlertItem(
        id=row["id"],
        account_id=row["account_id"],
        ts=row["ts"],
        seen=row["seen"],
        event=event,
        watch=AlertWatch(
            id=row["watch_id"],
            kind=row["watch_kind"],
            name=row["watch_name"],
            params=json.loads(row["watch_params"]),
        ),
        chat=refs.chat(chat_id),
        message=message,
        peer=refs.person(peer_id),
        away_since=payload.get("since"),
        path=_path(chat_id, message_id, peer_id),
    )


async def _alert_items(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> list[AlertItem]:
    chat_ids: set[int] = set()
    person_ids: set[int] = set()
    for row in rows:
        payload = json.loads(row["payload"])
        for key, bucket in (
            ("chat_id", chat_ids),
            ("sender_id", person_ids),
            ("peer_id", person_ids),
        ):
            value = _int(payload.get(key))
            if value is not None:
                bucket.add(value)
        if row["stored_sender"] is not None:
            person_ids.add(row["stored_sender"])
    refs = await load_refs(pool, account_id, chat_ids, person_ids)
    return [_alert_item(row, refs) for row in rows]


async def list_alerts(
    pool: asyncpg.Pool,
    account_id: int,
    *,
    seen: bool | None = None,
    before: int | None = None,
    limit: int = ALERT_PAGE,
) -> AlertPage:
    size = max(1, min(limit, MAX_ALERT_PAGE))
    params: list[object] = [account_id]
    where = "a.account_id = $1"
    if seen is not None:
        params.append(seen)
        where += f" AND a.seen = ${len(params)}"
    if before is not None:
        params.append(before)
        where += f" AND a.id < ${len(params)}"
    params.append(size + 1)
    rows = await pool.fetch(
        f"{_ALERT_ROWS} WHERE {where} ORDER BY a.id DESC LIMIT ${len(params)}", *params
    )
    items = await _alert_items(pool, account_id, rows[:size])
    return AlertPage(
        items=items,
        next_before=items[-1].id if len(rows) > size else None,
        unseen=await unseen_count(pool, account_id),
    )


async def get_alert(pool: asyncpg.Pool, alert_id: int) -> AlertItem | None:
    row = await pool.fetchrow(f"{_ALERT_ROWS} WHERE a.id = $1", alert_id)
    if row is None:
        return None
    return (await _alert_items(pool, row["account_id"], [row]))[0]


async def unseen_count(pool: asyncpg.Pool, account_id: int) -> int:
    return await pool.fetchval(
        "SELECT count(*) FROM alerts WHERE account_id = $1 AND NOT seen", account_id
    )


async def insert_alert(
    pool: asyncpg.Pool,
    account_id: int,
    watch_id: int,
    payload: dict[str, Any],
    *,
    dedup_key: str | None = None,
) -> AlertView | None:
    row = await pool.fetchrow(
        _INSERT_ALERT,
        account_id,
        watch_id,
        datetime.now(UTC),
        json.dumps(payload),
        dedup_key,
    )
    if row is None:
        return None
    await _notify(
        pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id, "alert_id": row["id"]}
    )
    return _to_alert(row)


async def mark_alert_seen(pool: asyncpg.Pool, alert_id: int) -> bool:
    account_id = await pool.fetchval(
        "UPDATE alerts SET seen = true WHERE id = $1 RETURNING account_id", alert_id
    )
    if account_id is None:
        return False
    await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id})
    return True


async def mark_all_seen(
    pool: asyncpg.Pool, account_id: int, up_to: int | None = None
) -> int:
    result = await pool.execute(
        "UPDATE alerts SET seen = true WHERE account_id = $1 AND NOT seen "
        "AND ($2::int IS NULL OR id <= $2)",
        account_id,
        up_to,
    )
    marked = int(result.split()[-1])
    if marked:
        await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id})
    return marked