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

from datetime import datetime

import asyncpg

from utils.policy.models import FolderSpec
from utils.read.accounts import self_user_id
from utils.read.message_view import build_message_view, load_raw, media_ref_from
from utils.read.models import (
    ChatListItem,
    MediaRef,
    MessageVersionView,
    MessageView,
    Page,
)
from utils.read.read_receipts import read_times, read_up_to

_MESSAGE_COLS = (
    "chat_id, message_id, date, sender_id, text, has_media, is_self_destruct, "
    "edited_at, deleted_at, raw, raw->>'media_group_id' AS media_group_id"
)


async def _media_map(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> dict[tuple[int, int], asyncpg.Record]:
    message_ids = list({row["message_id"] for row in rows})
    if not message_ids:
        return {}
    media_rows = await pool.fetch(
        "SELECT id, chat_id, message_id, kind, downloaded, mime, file_size, "
        "ttl_seconds, extracted_text FROM media "
        "WHERE account_id = $1 AND message_id = ANY($2::bigint[])",
        account_id,
        message_ids,
    )
    return {(row["chat_id"], row["message_id"]): row for row in media_rows}


def _single_media(
    row: asyncpg.Record, raw: dict, media_by_key: dict[tuple[int, int], asyncpg.Record]
) -> list[MediaRef]:
    media_row = media_by_key.get((row["chat_id"], row["message_id"]))
    if not (row["has_media"] or media_row):
        return []
    ref = media_ref_from(row["message_id"], raw, media_row)
    return [ref] if ref else []


_ALL_IDS = """
    SELECT chat_id FROM chat_stats WHERE account_id = $1
    UNION
    SELECT chat_id FROM dialogs WHERE account_id = $1
    UNION
    SELECT scope_id AS chat_id FROM capture_policy
    WHERE account_id = $1 AND scope_type = 'chat' AND scope_id IS NOT NULL
"""

_ONE_ID = """
    SELECT chat_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2
    UNION
    SELECT chat_id FROM dialogs WHERE account_id = $1 AND chat_id = $2
    UNION
    SELECT scope_id AS chat_id FROM capture_policy
    WHERE account_id = $1 AND scope_type = 'chat' AND scope_id = $2
"""

_MANY_IDS = """
    SELECT chat_id FROM chat_stats
    WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
    UNION
    SELECT chat_id FROM dialogs
    WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
    UNION
    SELECT scope_id AS chat_id FROM capture_policy
    WHERE account_id = $1 AND scope_type = 'chat'
      AND scope_id = ANY($2::bigint[])
"""

_CHAT_ROWS = """
WITH ids AS ({ids})
SELECT ids.chat_id,
       COALESCE(cs.message_count, 0) AS message_count,
       cs.last_date, cs.last_text, cs.last_sender_id, cs.last_message_id,
       COALESCE(cs.last_deleted, false) AS last_deleted,
       COALESCE(named.title,
                NULLIF(trim(concat_ws(' ', p.first_name, p.last_name)), ''),
                p.username) AS title,
       COALESCE(typed.is_broadcast, false) AS is_broadcast,
       COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool,
                p.raw->>'type' = 'ChatType.BOT', false) AS is_bot,
       COALESCE((p.raw->>'is_contact')::bool, (p.raw->>'contact')::bool,
                false) AS is_contact,
       EXISTS (SELECT 1 FROM avatars a
               WHERE a.account_id = $1 AND a.owner_id = ids.chat_id) AS has_avatar
FROM ids
LEFT JOIN chat_stats cs ON cs.account_id = $1 AND cs.chat_id = ids.chat_id
LEFT JOIN peers p ON p.account_id = $1 AND p.peer_id = ids.chat_id
LEFT JOIN LATERAL (
    SELECT ch.title FROM chat_history ch
    WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id AND ch.title IS NOT NULL
    ORDER BY ch.ts DESC LIMIT 1
) named ON true
LEFT JOIN LATERAL (
    SELECT COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type')
           = 'ChatType.CHANNEL' AS is_broadcast
    FROM chat_history ch
    WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id
      AND COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type') IS NOT NULL
    ORDER BY ch.ts DESC LIMIT 1
) typed ON true
"""


def _folder_filter(base: int) -> str:
    return (
        f"NOT (chat.chat_id = ANY(${base + 1}::bigint[])) "
        f"AND (chat.chat_id = ANY(${base + 2}::bigint[]) "
        f"OR (NOT ${base + 3}::bool AND CASE "
        f"WHEN chat.is_broadcast THEN ${base + 4}::bool "
        f"WHEN chat.chat_id < 0 THEN ${base + 5}::bool "
        f"WHEN chat.is_bot THEN ${base + 6}::bool "
        f"WHEN chat.is_contact THEN ${base + 7}::bool "
        f"ELSE ${base + 8}::bool END))"
    )


def _folder_params(folder: FolderSpec) -> list[object]:
    return [
        sorted(folder.exclude_ids),
        sorted(folder.include_ids | folder.pinned_ids),
        folder.is_chatlist,
        folder.broadcasts,
        folder.groups,
        folder.bots,
        folder.contacts,
        folder.non_contacts,
    ]


def _chat_item(row: asyncpg.Record) -> ChatListItem:
    return ChatListItem(
        chat_id=row["chat_id"],
        title=row["title"],
        kind="private" if row["chat_id"] > 0 else "group",
        has_avatar=row["has_avatar"],
        is_bot=bool(row["is_bot"]),
        is_contact=bool(row["is_contact"]),
        is_broadcast=bool(row["is_broadcast"]),
        message_count=row["message_count"],
        last_date=row["last_date"],
        last_text=row["last_text"],
        last_sender_id=row["last_sender_id"],
        last_message_id=row["last_message_id"],
        last_deleted=row["last_deleted"],
    )


async def list_chats(
    pool: asyncpg.Pool,
    account_id: int,
    page: Page,
    *,
    folder: FolderSpec | None = None,
    search: str | None = None,
) -> list[ChatListItem]:
    params: list[object] = [account_id, page.capped_limit, page.offset]
    rows_sql = _CHAT_ROWS.format(ids=_ALL_IDS)
    clauses: list[str] = []
    if folder is not None:
        clauses.append(_folder_filter(len(params)))
        params.extend(_folder_params(folder))
    if search:
        params.append(f"%{search}%")
        clauses.append(f"chat.title ILIKE ${len(params)}")
    where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
    query = (
        f"SELECT chat.* FROM ({rows_sql}) chat{where} "  # noqa: S608
        "ORDER BY last_date DESC NULLS LAST, chat_id DESC LIMIT $2 OFFSET $3"
    )
    rows = await pool.fetch(query, *params)
    return [_chat_item(row) for row in rows]


async def folder_chat_ids(
    pool: asyncpg.Pool, account_id: int, folder: FolderSpec
) -> list[int]:
    rows_sql = _CHAT_ROWS.format(ids=_ALL_IDS)
    rows = await pool.fetch(
        f"SELECT chat.chat_id FROM ({rows_sql}) chat WHERE {_folder_filter(1)}",  # noqa: S608
        account_id,
        *_folder_params(folder),
    )
    return [row["chat_id"] for row in rows]


async def get_chat(
    pool: asyncpg.Pool, account_id: int, chat_id: int
) -> ChatListItem | None:
    row = await pool.fetchrow(_CHAT_ROWS.format(ids=_ONE_ID), account_id, chat_id)
    return _chat_item(row) if row is not None else None


async def get_chats(
    pool: asyncpg.Pool, account_id: int, ids: list[int]
) -> list[ChatListItem]:
    if not ids:
        return []
    rows = await pool.fetch(_CHAT_ROWS.format(ids=_MANY_IDS), account_id, ids)
    return [_chat_item(row) for row in rows]


async def get_chat_history(  # noqa: PLR0913
    pool: asyncpg.Pool,
    account_id: int,
    chat_id: int,
    page: Page,
    *,
    include_deleted: bool = True,
    before_id: int | None = None,
    after_id: int | None = None,
) -> list[MessageView]:
    where = "account_id = $1 AND chat_id = $2"
    if not include_deleted:
        where += " AND deleted_at IS NULL"
    params: list[object] = [account_id, chat_id]
    if after_id is not None:
        params.append(after_id)
        where += f" AND message_id > ${len(params)}"
        order = "date ASC, message_id ASC"
    elif before_id is not None:
        params.append(before_id)
        where += f" AND message_id < ${len(params)}"
        order = "date DESC, message_id DESC"
    else:
        order = "date DESC, message_id DESC"
    params.append(page.capped_limit)
    query = (
        f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} "  # noqa: S608
        f"ORDER BY {order} LIMIT ${len(params)}"
    )
    if before_id is None and after_id is None:
        params.append(page.offset)
        query += f" OFFSET ${len(params)}"
    rows = await pool.fetch(query, *params)
    media_by_key = await _media_map(pool, account_id, rows)
    parsed = [(row, load_raw(row["raw"])) for row in rows]
    views: list[MessageView] = []
    index = 0
    while index < len(parsed):
        group_id = parsed[index][0]["media_group_id"]
        end = index + 1
        if group_id is not None:
            while end < len(parsed) and parsed[end][0]["media_group_id"] == group_id:
                end += 1
        members = parsed[index:end]
        if len(members) == 1:
            row, raw = members[0]
            views.append(
                build_message_view(row, raw, _single_media(row, raw, media_by_key))
            )
        else:
            views.append(_build_album(members, media_by_key))
        index = end
    await _apply_read_status(pool, account_id, chat_id, views)
    await _apply_notes(pool, account_id, chat_id, views)
    return views


async def get_message(
    pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> MessageView | None:
    row = await pool.fetchrow(
        f"SELECT {_MESSAGE_COLS} FROM messages "  # noqa: S608
        "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
        account_id,
        chat_id,
        message_id,
    )
    if row is None:
        return None
    media_by_key = await _media_map(pool, account_id, [row])
    raw = load_raw(row["raw"])
    view = build_message_view(row, raw, _single_media(row, raw, media_by_key))
    await _apply_read_status(pool, account_id, chat_id, [view])
    await _apply_notes(pool, account_id, chat_id, [view])
    return view


async def _apply_read_status(
    pool: asyncpg.Pool, account_id: int, chat_id: int, views: list[MessageView]
) -> None:
    self_id = await self_user_id(pool, account_id)
    if self_id is None:
        return
    own = [view for view in views if view.sender_id == self_id]
    if not own:
        return
    marker = await read_up_to(pool, account_id, chat_id)
    times = await read_times(
        pool, account_id, chat_id, [(view.message_id, view.date) for view in own]
    )
    for view in own:
        view.read_at = times.get(view.message_id)
        view.read = view.read_at is not None or (
            marker is not None and view.message_id <= marker
        )


async def _apply_notes(
    pool: asyncpg.Pool, account_id: int, chat_id: int, views: list[MessageView]
) -> None:
    if not views:
        return
    rows = await pool.fetch(
        "SELECT DISTINCT message_id FROM annotations "
        "WHERE account_id = $1 AND chat_id = $2 AND message_id = ANY($3::bigint[])",
        account_id,
        chat_id,
        [view.message_id for view in views],
    )
    noted = {row["message_id"] for row in rows}
    for view in views:
        view.has_note = view.message_id in noted


def _build_album(
    members: list[tuple[asyncpg.Record, dict]],
    media_by_key: dict[tuple[int, int], asyncpg.Record],
) -> MessageView:
    ordered = sorted(members, key=lambda m: m[0]["message_id"])
    media: list[MediaRef] = []
    for row, raw in ordered:
        media_row = media_by_key.get((row["chat_id"], row["message_id"]))
        ref = media_ref_from(row["message_id"], raw, media_row)
        if ref:
            media.append(ref)
    primary_row, primary_raw = next(
        ((row, raw) for row, raw in ordered if row["text"]), ordered[0]
    )
    return build_message_view(primary_row, primary_raw, media)


async def get_deleted_messages(
    pool: asyncpg.Pool, account_id: int, page: Page, *, chat_id: int | None = None
) -> list[MessageView]:
    params: list[object] = [account_id]
    where = "account_id = $1 AND deleted_at IS NOT NULL"
    if chat_id is not None:
        params.append(chat_id)
        where += f" AND chat_id = ${len(params)}"
    params.append(page.capped_limit)
    params.append(page.offset)
    rows = await pool.fetch(
        f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} "  # noqa: S608
        f"ORDER BY deleted_at DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
        *params,
    )
    return await _single_views(pool, account_id, rows)


async def _single_views(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> list[MessageView]:
    media_by_key = await _media_map(pool, account_id, rows)
    views: list[MessageView] = []
    for row in rows:
        raw = load_raw(row["raw"])
        views.append(
            build_message_view(row, raw, _single_media(row, raw, media_by_key))
        )
    return views


async def _rows_at(
    pool: asyncpg.Pool, account_id: int, keys: list[tuple[int, int, datetime]]
) -> list[asyncpg.Record]:
    if not keys:
        return []
    params: list[object] = [account_id]
    clauses: list[str] = []
    for key in keys:
        params.extend(key)
        n = len(params)
        clauses.append(
            f"(chat_id = ${n - 2} AND message_id = ${n - 1} AND date = ${n})"
        )
    return await pool.fetch(
        f"SELECT {_MESSAGE_COLS} FROM messages "  # noqa: S608
        f"WHERE account_id = $1 AND ({' OR '.join(clauses)})",
        *params,
    )


async def get_messages_at(
    pool: asyncpg.Pool, account_id: int, keys: list[tuple[int, int, datetime | None]]
) -> dict[tuple[int, int], MessageView]:
    dated = list(
        dict.fromkeys((chat, message, date) for chat, message, date in keys if date)
    )
    rows = await _rows_at(pool, account_id, dated)
    views = {
        (view.chat_id, view.message_id): view
        for view in await _single_views(pool, account_id, rows)
    }
    for chat_id, message_id, _ in keys:
        if (chat_id, message_id) not in views:
            view = await get_message(pool, account_id, chat_id, message_id)
            if view is not None:
                views[chat_id, message_id] = view
    return views


async def get_deleted_between(  # noqa: PLR0913
    pool: asyncpg.Pool,
    account_id: int,
    chat_id: int,
    since: datetime,
    until: datetime,
    limit: int,
    before: tuple[datetime, int] | None = None,
) -> list[MessageView]:
    params: list[object] = [account_id, chat_id, since, until, limit]
    where = "account_id = $1 AND chat_id = $2 AND deleted_at BETWEEN $3 AND $4"
    if before is not None:
        params.extend(before)
        where += " AND (date, message_id) < ($6, $7)"
    rows = await pool.fetch(
        f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} "  # noqa: S608
        "ORDER BY date DESC, message_id DESC LIMIT $5",
        *params,
    )
    return await _single_views(pool, account_id, rows)


async def get_message_versions(
    pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> list[MessageVersionView]:
    rows = await pool.fetch(
        "SELECT observed_at, edit_date, text FROM message_versions "
        "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
        "ORDER BY observed_at LIMIT 500",
        account_id,
        chat_id,
        message_id,
    )
    return [MessageVersionView(**dict(row)) for row in rows]