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

import json
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from enum import StrEnum

import asyncpg

from utils.read import chats, peers
from utils.read.accounts import self_user_id
from utils.read.models import (
    DEFAULT_LIMIT,
    MAX_LIMIT,
    ChatListItem,
    FeedChat,
    FeedItem,
    FeedMessagesPage,
    FeedPage,
    FeedUnseen,
    MessageView,
    PeerView,
)

CURSOR_PATTERN = r"^\d+\.\d+$"
_EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
_MICROSECOND = timedelta(microseconds=1)

_ITEM_COLS = (
    "id, kind, chat_id, chat_kind, message_id, message_date, at, since, "
    "message_count, detail"
)

_VERSIONS = """
SELECT v.chat_id, v.message_id, v.observed_at, v.edit_date, v.text,
       p.text AS before, p.observed_at IS NOT NULL AS has_before
FROM unnest($2::bigint[], $3::bigint[], $4::timestamptz[])
    AS k(chat_id, message_id, observed_at)
JOIN message_versions v
  ON v.account_id = $1 AND v.chat_id = k.chat_id
 AND v.message_id = k.message_id AND v.observed_at = k.observed_at
LEFT JOIN LATERAL (
    SELECT text, observed_at FROM message_versions p
    WHERE p.account_id = $1 AND p.chat_id = v.chat_id
      AND p.message_id = v.message_id AND p.observed_at < v.observed_at
    ORDER BY p.observed_at DESC LIMIT 1
) p ON true
"""

_UNSEEN = """
SELECT count(*) FILTER (WHERE kind = 'deleted') AS deleted,
       count(*) FILTER (WHERE kind = 'edited') AS edited
FROM feed_items
WHERE account_id = $1 AND at > coalesce($2, '-infinity'::timestamptz)
"""

_MARK_SEEN = """
INSERT INTO feed_seen (account_id, seen_at) VALUES ($1, coalesce($2, now()))
ON CONFLICT (account_id) DO UPDATE
SET seen_at = greatest(feed_seen.seen_at, EXCLUDED.seen_at)
"""


class FeedKind(StrEnum):
    ALL = "all"
    MESSAGES = "messages"
    DELETED = "deleted"
    EDITED = "edited"
    PROFILE = "profile"


class FeedChatKind(StrEnum):
    DM = "dm"
    GROUP = "group"
    CHANNEL = "channel"


@dataclass(slots=True)
class FeedFilter:
    kind: FeedKind = FeedKind.ALL
    chat_id: int | None = None
    chat_kind: FeedChatKind | None = None
    since: datetime | None = None


def encode_cursor(at: datetime, key: int) -> str:
    return f"{(at - _EPOCH) // _MICROSECOND}.{key}"


def decode_cursor(cursor: str) -> tuple[datetime, int]:
    micros, _, key = cursor.partition(".")
    return _EPOCH + int(micros) * _MICROSECOND, int(key)


def _where(
    account_id: int, feed: FeedFilter, cursor: str | None
) -> tuple[str, list[object]]:
    params: list[object] = [account_id]
    clauses = ["account_id = $1"]
    if feed.kind is FeedKind.MESSAGES:
        clauses.append("kind <> 'profile'")
    exact = feed.kind not in {FeedKind.ALL, FeedKind.MESSAGES}
    conditions = (
        ("kind = {}", feed.kind.value if exact else None),
        ("chat_id = {}", feed.chat_id),
        ("chat_kind = {}", feed.chat_kind.value if feed.chat_kind else None),
        ("at >= {}", feed.since),
    )
    for template, value in conditions:
        if value is not None:
            params.append(value)
            clauses.append(template.format(f"${len(params)}"))
    if cursor is not None:
        params.extend(decode_cursor(cursor))
        clauses.append(f"(at, id) < (${len(params) - 1}, ${len(params)})")
    return " AND ".join(clauses), params


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


async def _versions(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> dict[tuple[int, int, datetime], asyncpg.Record]:
    edited = [row for row in rows if row["kind"] == FeedKind.EDITED]
    if not edited:
        return {}
    found = await pool.fetch(
        _VERSIONS,
        account_id,
        [row["chat_id"] for row in edited],
        [row["message_id"] for row in edited],
        [row["at"] for row in edited],
    )
    return {
        (row["chat_id"], row["message_id"], row["observed_at"]): row for row in found
    }


async def _sender_names(
    pool: asyncpg.Pool,
    account_id: int,
    messages: dict[tuple[int, int], MessageView],
    chat_views: dict[int, ChatListItem],
) -> dict[int, str | None]:
    ids = {m.sender_id for m in messages.values() if m.sender_id is not None}
    found = await peers.get_peers(pool, account_id, list(ids))
    names: dict[int, str | None] = {
        chat_id: chat.title for chat_id, chat in chat_views.items()
    }
    names.update({peer.peer_id: _peer_name(peer) for peer in found})
    return names


async def _build_items(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> list[FeedItem]:
    if not rows:
        return []
    messages = await chats.get_messages_at(
        pool,
        account_id,
        [
            (row["chat_id"], row["message_id"], row["message_date"])
            for row in rows
            if row["kind"] != FeedKind.PROFILE
        ],
    )
    versions = await _versions(pool, account_id, rows)
    chat_views = {
        chat.chat_id: chat
        for chat in await chats.get_chats(
            pool, account_id, list({row["chat_id"] for row in rows})
        )
    }
    names = await _sender_names(pool, account_id, messages, chat_views)
    self_id = await self_user_id(pool, account_id)
    items: list[FeedItem] = []
    for row in rows:
        message = messages.get((row["chat_id"], row["message_id"]))
        chat = chat_views.get(row["chat_id"])
        version = versions.get((row["chat_id"], row["message_id"], row["at"]))
        sender_id = message.sender_id if message else None
        items.append(
            FeedItem(
                id=row["id"],
                kind=row["kind"],
                at=row["at"],
                since=row["since"],
                message_count=row["message_count"],
                chat=FeedChat(
                    chat_id=row["chat_id"],
                    title=chat.title if chat else None,
                    kind=row["chat_kind"],
                    has_avatar=chat.has_avatar if chat else False,
                ),
                message_id=row["message_id"],
                message=message,
                sender_id=sender_id,
                sender_name=names.get(sender_id) if sender_id is not None else None,
                outgoing=self_id is not None and sender_id == self_id,
                before=version["before"] if version else None,
                after=version["text"] if version else None,
                has_before=bool(version and version["has_before"]),
                edit_date=version["edit_date"] if version else None,
                detail=json.loads(row["detail"]) if row["detail"] else None,
            )
        )
    return items


async def list_feed(
    pool: asyncpg.Pool,
    account_id: int,
    feed: FeedFilter,
    *,
    limit: int = DEFAULT_LIMIT,
    cursor: str | None = None,
) -> FeedPage:
    size = max(1, min(limit, MAX_LIMIT))
    where, params = _where(account_id, feed, cursor)
    params.append(size + 1)
    rows = await pool.fetch(
        f"SELECT {_ITEM_COLS} FROM feed_items WHERE {where} "  # noqa: S608
        f"ORDER BY at DESC, id DESC LIMIT ${len(params)}",
        *params,
    )
    page = rows[:size]
    more = len(rows) > size
    return FeedPage(
        items=await _build_items(pool, account_id, page),
        next_cursor=encode_cursor(page[-1]["at"], page[-1]["id"]) if more else None,
    )


async def unseen(pool: asyncpg.Pool, account_id: int) -> FeedUnseen:
    seen_at = await pool.fetchval(
        "SELECT seen_at FROM feed_seen WHERE account_id = $1", account_id
    )
    counts = await pool.fetchrow(_UNSEEN, account_id, seen_at)
    latest = await list_feed(
        pool, account_id, FeedFilter(kind=FeedKind.MESSAGES), limit=1
    )
    return FeedUnseen(
        count=counts["deleted"] + counts["edited"],
        deleted=counts["deleted"],
        edited=counts["edited"],
        seen_at=seen_at,
        latest=latest.items[0] if latest.items else None,
    )


async def mark_seen(
    pool: asyncpg.Pool, account_id: int, at: datetime | None
) -> FeedUnseen:
    await pool.execute(_MARK_SEEN, account_id, at)
    return await unseen(pool, account_id)


async def item_messages(
    pool: asyncpg.Pool,
    account_id: int,
    item_id: int,
    *,
    limit: int = DEFAULT_LIMIT,
    cursor: str | None = None,
) -> FeedMessagesPage | None:
    row = await pool.fetchrow(
        "SELECT chat_id, since, at FROM feed_items "
        "WHERE id = $1 AND account_id = $2 AND kind = 'deleted'",
        item_id,
        account_id,
    )
    if row is None:
        return None
    size = max(1, min(limit, MAX_LIMIT))
    views = await chats.get_deleted_between(
        pool,
        account_id,
        row["chat_id"],
        row["since"],
        row["at"],
        size + 1,
        decode_cursor(cursor) if cursor else None,
    )
    page = views[:size]
    more = len(views) > size
    return FeedMessagesPage(
        messages=page,
        next_cursor=encode_cursor(page[-1].date, page[-1].message_id) if more else None,
    )