from dataclasses import dataclass
from datetime import datetime

import asyncpg

CHANNEL_ID_THRESHOLD = -1000000000000

_UPSERT_MESSAGE = """
INSERT INTO messages
    (account_id, chat_id, message_id, date, sender_id, text, raw,
     has_media, is_self_destruct)
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9)
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
    sender_id = EXCLUDED.sender_id,
    text = EXCLUDED.text,
    raw = EXCLUDED.raw,
    has_media = EXCLUDED.has_media,
    is_self_destruct = EXCLUDED.is_self_destruct
"""

_TOUCH_EDITED = """
INSERT INTO messages
    (account_id, chat_id, message_id, date, sender_id, text, raw,
     has_media, is_self_destruct, edited_at)
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, CASE WHEN $10 THEN now() END)
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
    text = EXCLUDED.text,
    raw = EXCLUDED.raw,
    has_media = EXCLUDED.has_media,
    is_self_destruct = EXCLUDED.is_self_destruct,
    edited_at = coalesce(EXCLUDED.edited_at, messages.edited_at)
"""

_LOCK_MESSAGE = """
SELECT pg_advisory_xact_lock(
    hashtextextended(concat_ws(':', 'message', $1::int, $2::bigint, $3::bigint), 0)
)
"""

_FIND_DELETED = """
SELECT chat_id, message_id, date FROM messages
WHERE account_id = $1 AND message_id = ANY($2::bigint[]) AND deleted_at IS NULL
"""

_MARK_DELETED = """
UPDATE messages SET deleted_at = now()
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
"""

_MARK_DELETED_LATE = """
UPDATE messages SET deleted_at = now()
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
  AND deleted_at IS NULL
"""

_INSERT_VERSION = """
INSERT INTO message_versions
    (account_id, chat_id, message_id, observed_at, edit_date, text, raw)
VALUES ($1, $2, $3, clock_timestamp(), $4, $5, $6::jsonb)
ON CONFLICT DO NOTHING
"""

_SNAPSHOT_ORIGINAL = """
INSERT INTO message_versions
    (account_id, chat_id, message_id, observed_at, edit_date, text, raw)
SELECT account_id, chat_id, message_id, clock_timestamp(), NULL, text, raw
FROM messages m
WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3 AND m.date = $4
  AND NOT EXISTS (
    SELECT 1 FROM message_versions v
    WHERE v.account_id = m.account_id AND v.chat_id = m.chat_id
      AND v.message_id = m.message_id
  )
ON CONFLICT DO NOTHING
"""

_INSERT_PAST_VERSION = """
INSERT INTO message_versions
    (account_id, chat_id, message_id, observed_at, edit_date, text, raw)
SELECT $1::int, $2::bigint, $3::bigint,
    coalesce(
        (SELECT min(v.observed_at) - interval '1 microsecond'
         FROM message_versions v
         WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
           AND v.edit_date > $4::timestamptz),
        $4::timestamptz,
        $5::timestamptz
    ),
    $4::timestamptz, $6::text, $7::jsonb
WHERE NOT EXISTS (
    SELECT 1 FROM message_versions v
    WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
      AND v.edit_date IS NOT DISTINCT FROM $4::timestamptz
      AND v.text IS NOT DISTINCT FROM $6::text
)
ON CONFLICT DO NOTHING
"""

_CURRENT_CONTENT = """
SELECT
    m.text AS text,
    m.raw->>'edit_date' AS edit_date,
    (SELECT d.unique_id FROM media d
     WHERE d.account_id = m.account_id AND d.chat_id = m.chat_id
       AND d.message_id = m.message_id) AS media_unique_id
FROM messages m
WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3 AND m.date = $4
"""

_CURRENT_MEDIA = """
SELECT unique_id, storage_key, file_size, downloaded FROM media
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3
"""

_INSERT_MEDIA_VERSION = """
INSERT INTO media_versions
    (account_id, chat_id, message_id, observed_at, kind, storage_key,
     file_size, mime, ttl_seconds, unique_id)
VALUES ($1, $2, $3, clock_timestamp(), $4, $5, $6, $7, $8, $9)
ON CONFLICT (account_id, chat_id, message_id, storage_key) DO NOTHING
"""

_INSERT_MEDIA = """
INSERT INTO media
    (account_id, chat_id, message_id, kind, storage_key, file_size, mime,
     ttl_seconds, downloaded, unique_id, evicted_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, CASE WHEN $11 THEN now() END)
ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
    kind = EXCLUDED.kind,
    storage_key = EXCLUDED.storage_key,
    file_size = EXCLUDED.file_size,
    mime = EXCLUDED.mime,
    ttl_seconds = EXCLUDED.ttl_seconds,
    downloaded = EXCLUDED.downloaded,
    unique_id = EXCLUDED.unique_id,
    evicted_at = CASE
        WHEN EXCLUDED.downloaded THEN NULL
        WHEN media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
            THEN EXCLUDED.evicted_at
        ELSE coalesce(media.evicted_at, EXCLUDED.evicted_at) END,
    evicted_key = CASE WHEN EXCLUDED.downloaded
        OR media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
        THEN NULL ELSE media.evicted_key END
"""

_IS_EVICTED = """
SELECT EXISTS (
    SELECT 1 FROM media WHERE unique_id = $1 AND evicted_at IS NOT NULL
)
"""

_UNEVICT = """
UPDATE media SET storage_key = $2, file_size = $3, downloaded = true,
    evicted_at = NULL, evicted_key = NULL
WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL)
"""


@dataclass(frozen=True, slots=True)
class MessageRow:
    chat_id: int
    message_id: int
    date: datetime
    sender_id: int | None
    text: str | None
    raw: str
    has_media: bool
    is_self_destruct: bool
    edit_date: datetime | None
    media_unique_id: str | None

    @property
    def columns(self) -> tuple:
        return (
            self.chat_id,
            self.message_id,
            self.date,
            self.sender_id,
            self.text,
            self.raw,
            self.has_media,
            self.is_self_destruct,
        )


async def _lock_current(
    conn: asyncpg.Connection, account_id: int, row: MessageRow
) -> asyncpg.Record | None:
    await conn.execute(_LOCK_MESSAGE, account_id, row.chat_id, row.message_id)
    return await conn.fetchrow(
        _CURRENT_CONTENT, account_id, row.chat_id, row.message_id, row.date
    )


def _is_stale(current: asyncpg.Record, row: MessageRow) -> bool:
    stored = current["edit_date"]
    if stored is None:
        return False
    if row.edit_date is None:
        return True
    try:
        return row.edit_date < datetime.fromisoformat(stored)
    except (TypeError, ValueError):
        return False


async def _record_past(
    conn: asyncpg.Connection, account_id: int, row: MessageRow
) -> None:
    await conn.execute(
        _INSERT_PAST_VERSION,
        account_id,
        row.chat_id,
        row.message_id,
        row.edit_date,
        row.date,
        row.text,
        row.raw,
    )


async def _record_edit(
    conn: asyncpg.Connection, account_id: int, row: MessageRow, *, mark_edited: bool
) -> None:
    await conn.execute(
        _SNAPSHOT_ORIGINAL, account_id, row.chat_id, row.message_id, row.date
    )
    await conn.execute(_TOUCH_EDITED, account_id, *row.columns, mark_edited)
    await conn.execute(
        _INSERT_VERSION,
        account_id,
        row.chat_id,
        row.message_id,
        row.edit_date,
        row.text,
        row.raw,
    )


async def upsert_message(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
    async with pool.acquire() as conn, conn.transaction():
        current = await _lock_current(conn, account_id, row)
        if current is None:
            await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
            return True
        stored_media = current["media_unique_id"]
        changed = current["text"] != row.text or (
            stored_media is not None and stored_media != row.media_unique_id
        )
        if _is_stale(current, row):
            if changed:
                await _record_past(conn, account_id, row)
            return False
        if changed:
            await _record_edit(
                conn, account_id, row, mark_edited=row.edit_date is not None
            )
        else:
            await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
    return True


async def max_message_id(
    pool: asyncpg.Pool, account_id: int, chat_id: int
) -> int | None:
    return await pool.fetchval(
        "SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
        account_id,
        chat_id,
    )


async def _mark_deleted(
    pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
) -> None:
    for row in rows:
        await pool.execute(
            _MARK_DELETED, account_id, row["chat_id"], row["message_id"], row["date"]
        )


async def mark_deleted_box(
    pool: asyncpg.Pool, account_id: int, message_ids: list[int]
) -> None:
    rows = await pool.fetch(
        f"{_FIND_DELETED} AND chat_id > $3",
        account_id,
        message_ids,
        CHANNEL_ID_THRESHOLD,
    )
    await _mark_deleted(pool, account_id, rows)


async def mark_deleted_channel(
    pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int]
) -> None:
    rows = await pool.fetch(
        f"{_FIND_DELETED} AND chat_id = $3", account_id, message_ids, chat_id
    )
    await _mark_deleted(pool, account_id, rows)


async def add_version(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
    async with pool.acquire() as conn, conn.transaction():
        current = await _lock_current(conn, account_id, row)
        if current is not None:
            changed = (
                current["text"] != row.text
                or current["media_unique_id"] != row.media_unique_id
            )
            if not changed:
                return False
            if _is_stale(current, row):
                await _record_past(conn, account_id, row)
                return False
        await _record_edit(conn, account_id, row, mark_edited=True)
    return True


async def mark_deleted_late(
    pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int, date: datetime
) -> None:
    await pool.execute(_MARK_DELETED_LATE, account_id, chat_id, message_id, date)


async def current_media(
    pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> asyncpg.Record | None:
    return await pool.fetchrow(_CURRENT_MEDIA, account_id, chat_id, message_id)


async def insert_media(  # noqa: PLR0913
    pool: asyncpg.Pool,
    account_id: int,
    chat_id: int,
    message_id: int,
    kind: str,
    storage_key: str | None,
    file_size: int | None,
    mime: str | None,
    ttl_seconds: int | None,
    unique_id: str | None,
    *,
    downloaded: bool,
    evicted: bool = False,
) -> None:
    async with pool.acquire() as conn, conn.transaction():
        await conn.execute(
            _INSERT_MEDIA,
            account_id,
            chat_id,
            message_id,
            kind,
            storage_key,
            file_size,
            mime,
            ttl_seconds,
            downloaded,
            unique_id,
            evicted,
        )
        if storage_key is not None:
            await conn.execute(
                _INSERT_MEDIA_VERSION,
                account_id,
                chat_id,
                message_id,
                kind,
                storage_key,
                file_size,
                mime,
                ttl_seconds,
                unique_id,
            )


async def is_evicted(pool: asyncpg.Pool, unique_id: str | None) -> bool:
    return unique_id is not None and await pool.fetchval(_IS_EVICTED, unique_id)


async def unevict(
    pool: asyncpg.Pool, unique_id: str | None, storage_key: str, file_size: int | None
) -> bool:
    async with pool.acquire() as conn, conn.transaction():
        await conn.execute(_UNEVICT, unique_id, storage_key, file_size)
        dropped = await conn.fetchval(
            "DELETE FROM media_trash WHERE storage_key = $1 RETURNING storage_key",
            storage_key,
        )
    return dropped is not None


async def insert_callbacks(
    pool: asyncpg.Pool,
    account_id: int,
    chat_id: int,
    message_id: int,
    buttons: list[tuple[int, str | None, bytes | None]],
) -> None:
    async with pool.acquire() as conn, conn.transaction():
        await conn.execute(
            "DELETE FROM callbacks "
            "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
            account_id,
            chat_id,
            message_id,
        )
        await conn.executemany(
            "INSERT INTO callbacks "
            "(account_id, chat_id, message_id, position, label, data) "
            "VALUES ($1, $2, $3, $4, $5, $6)",
            [
                (account_id, chat_id, message_id, position, label, data)
                for position, label, data in buttons
            ],
        )


async def sync_reactions(
    pool: asyncpg.Pool,
    account_id: int,
    chat_id: int,
    message_id: int,
    current: list[tuple[int, str]],
) -> None:
    async with pool.acquire() as conn, conn.transaction():
        rows = await conn.fetch(
            "SELECT id, peer_id, reaction FROM reactions "
            "WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
            "AND removed_at IS NULL",
            account_id,
            chat_id,
            message_id,
        )
        existing = {(row["peer_id"], row["reaction"]): row["id"] for row in rows}
        current_set = set(current)
        stale = [pid for key, pid in existing.items() if key not in current_set]
        if stale:
            await conn.execute(
                "UPDATE reactions SET removed_at = now() WHERE id = ANY($1::bigint[])",
                stale,
            )
        fresh = [key for key in current_set if key not in existing]
        if fresh:
            await conn.executemany(
                "INSERT INTO reactions "
                "(account_id, chat_id, message_id, peer_id, reaction) "
                "VALUES ($1, $2, $3, $4, $5)",
                [
                    (account_id, chat_id, message_id, peer_id, reaction)
                    for peer_id, reaction in fresh
                ],
            )