428th-exchange-bot/bot/db/repository.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import aiosqlite

SCHEMA = """
CREATE TABLE IF NOT EXISTS clients (
    telegram_id INTEGER PRIMARY KEY,
    username TEXT,
    first_name TEXT,
    last_name TEXT,
    topic_id INTEGER UNIQUE,
    language TEXT,
    ack_sent INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_clients_topic_id ON clients(topic_id);

CREATE TABLE IF NOT EXISTS message_links (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    client_telegram_id INTEGER NOT NULL,
    client_message_id INTEGER NOT NULL,
    manager_message_id INTEGER NOT NULL,
    created_at TEXT NOT NULL,
    UNIQUE(client_telegram_id, client_message_id),
    UNIQUE(manager_message_id)
);

CREATE INDEX IF NOT EXISTS idx_message_links_manager
    ON message_links(manager_message_id);
"""


@dataclass
class Client:
    telegram_id: int
    username: Optional[str]
    first_name: Optional[str]
    last_name: Optional[str]
    topic_id: Optional[int]
    language: Optional[str]
    ack_sent: bool
    created_at: str
    updated_at: str


@dataclass
class MessageLink:
    client_telegram_id: int
    client_message_id: int
    manager_message_id: int


def _utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()


class ClientRepository:
    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path
        self._conn: Optional[aiosqlite.Connection] = None

    async def connect(self) -> None:
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = await aiosqlite.connect(self._db_path)
        self._conn.row_factory = aiosqlite.Row
        await self._conn.execute("PRAGMA journal_mode=WAL")
        await self._conn.execute("PRAGMA busy_timeout=5000")
        await self._conn.executescript(SCHEMA)
        await self._ensure_columns()
        await self._conn.commit()

    async def _ensure_columns(self) -> None:
        conn = self._require_conn()
        async with conn.execute("PRAGMA table_info(clients)") as cursor:
            rows = await cursor.fetchall()
        columns = {row[1] for row in rows}
        if "language" not in columns:
            await conn.execute("ALTER TABLE clients ADD COLUMN language TEXT")
        if "ack_sent" not in columns:
            await conn.execute(
                "ALTER TABLE clients ADD COLUMN ack_sent INTEGER NOT NULL DEFAULT 0"
            )
            await conn.execute(
                """
                UPDATE clients
                SET ack_sent = 1
                WHERE telegram_id IN (
                    SELECT DISTINCT client_telegram_id FROM message_links
                )
                """
            )

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None

    def _require_conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("Database is not connected")
        return self._conn

    @staticmethod
    def _row_to_client(row: aiosqlite.Row) -> Client:
        return Client(
            telegram_id=row["telegram_id"],
            username=row["username"],
            first_name=row["first_name"],
            last_name=row["last_name"],
            topic_id=row["topic_id"],
            language=row["language"],
            ack_sent=bool(row["ack_sent"]),
            created_at=row["created_at"],
            updated_at=row["updated_at"],
        )

    async def upsert_client(
        self,
        telegram_id: int,
        username: Optional[str],
        first_name: Optional[str],
        last_name: Optional[str],
    ) -> Client:
        conn = self._require_conn()
        now = _utcnow()
        await conn.execute(
            """
            INSERT INTO clients (
                telegram_id, username, first_name, last_name, topic_id, language,
                ack_sent, created_at, updated_at
            ) VALUES (?, ?, ?, ?, NULL, NULL, 0, ?, ?)
            ON CONFLICT(telegram_id) DO UPDATE SET
                username = excluded.username,
                first_name = excluded.first_name,
                last_name = excluded.last_name,
                updated_at = excluded.updated_at
            """,
            (telegram_id, username, first_name, last_name, now, now),
        )
        await conn.commit()
        client = await self.get_by_telegram_id(telegram_id)
        if client is None:
            raise RuntimeError(f"Failed to upsert client {telegram_id}")
        return client

    async def get_by_telegram_id(self, telegram_id: int) -> Optional[Client]:
        conn = self._require_conn()
        async with conn.execute(
            "SELECT * FROM clients WHERE telegram_id = ?",
            (telegram_id,),
        ) as cursor:
            row = await cursor.fetchone()
        return self._row_to_client(row) if row else None

    async def get_by_topic_id(self, topic_id: int) -> Optional[Client]:
        conn = self._require_conn()
        async with conn.execute(
            "SELECT * FROM clients WHERE topic_id = ?",
            (topic_id,),
        ) as cursor:
            row = await cursor.fetchone()
        return self._row_to_client(row) if row else None

    async def set_topic_id(self, telegram_id: int, topic_id: int) -> None:
        conn = self._require_conn()
        await conn.execute(
            """
            UPDATE clients
            SET topic_id = ?, updated_at = ?
            WHERE telegram_id = ?
            """,
            (topic_id, _utcnow(), telegram_id),
        )
        await conn.commit()

    async def clear_topic_id(self, telegram_id: int) -> None:
        conn = self._require_conn()
        await conn.execute(
            """
            UPDATE clients
            SET topic_id = NULL, updated_at = ?
            WHERE telegram_id = ?
            """,
            (_utcnow(), telegram_id),
        )
        await conn.commit()

    async def mark_ack_sent(self, telegram_id: int) -> None:
        conn = self._require_conn()
        await conn.execute(
            """
            UPDATE clients
            SET ack_sent = 1, updated_at = ?
            WHERE telegram_id = ?
            """,
            (_utcnow(), telegram_id),
        )
        await conn.commit()

    async def set_language(self, telegram_id: int, language: str) -> None:
        conn = self._require_conn()
        await conn.execute(
            """
            UPDATE clients
            SET language = ?, updated_at = ?
            WHERE telegram_id = ?
            """,
            (language, _utcnow(), telegram_id),
        )
        await conn.commit()

    async def save_message_link(
        self,
        client_telegram_id: int,
        client_message_id: int,
        manager_message_id: int,
    ) -> None:
        conn = self._require_conn()
        await conn.execute(
            "DELETE FROM message_links WHERE manager_message_id = ?",
            (manager_message_id,),
        )
        await conn.execute(
            """
            DELETE FROM message_links
            WHERE client_telegram_id = ? AND client_message_id = ?
            """,
            (client_telegram_id, client_message_id),
        )
        await conn.execute(
            """
            INSERT INTO message_links (
                client_telegram_id, client_message_id, manager_message_id, created_at
            ) VALUES (?, ?, ?, ?)
            """,
            (client_telegram_id, client_message_id, manager_message_id, _utcnow()),
        )
        await conn.commit()

    async def get_link_by_client_message(
        self,
        client_telegram_id: int,
        client_message_id: int,
    ) -> Optional[MessageLink]:
        conn = self._require_conn()
        async with conn.execute(
            """
            SELECT client_telegram_id, client_message_id, manager_message_id
            FROM message_links
            WHERE client_telegram_id = ? AND client_message_id = ?
            """,
            (client_telegram_id, client_message_id),
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            return None
        return MessageLink(
            client_telegram_id=row["client_telegram_id"],
            client_message_id=row["client_message_id"],
            manager_message_id=row["manager_message_id"],
        )

    async def get_link_by_manager_message(
        self,
        manager_message_id: int,
    ) -> Optional[MessageLink]:
        conn = self._require_conn()
        async with conn.execute(
            """
            SELECT client_telegram_id, client_message_id, manager_message_id
            FROM message_links
            WHERE manager_message_id = ?
            """,
            (manager_message_id,),
        ) as cursor:
            row = await cursor.fetchone()
        if row is None:
            return None
        return MessageLink(
            client_telegram_id=row["client_telegram_id"],
            client_message_id=row["client_message_id"],
            manager_message_id=row["manager_message_id"],
        )

    async def delete_message_link(
        self,
        *,
        client_telegram_id: int,
        client_message_id: int,
        manager_message_id: int,
    ) -> None:
        conn = self._require_conn()
        await conn.execute(
            """
            DELETE FROM message_links
            WHERE client_telegram_id = ?
              AND client_message_id = ?
              AND manager_message_id = ?
            """,
            (client_telegram_id, client_message_id, manager_message_id),
        )
        await conn.commit()