muzovkantv2/utils/database.py

import aiosqlite
import logging

logger = logging.getLogger(__name__)


class FunchosaDatabase:
    def __init__(self, db_path='data/funchosa.db'):
        self.db_path = db_path
        self._conn: aiosqlite.Connection | None = None

    async def init_db(self):
        self._conn = await aiosqlite.connect(self.db_path)
        self._conn.row_factory = aiosqlite.Row

        await self._conn.executescript('''
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                message_id BIGINT UNIQUE NOT NULL,
                channel_id BIGINT NOT NULL,
                author_id BIGINT NOT NULL,
                author_name TEXT NOT NULL,
                content TEXT,
                timestamp TIMESTAMP NOT NULL,
                message_url TEXT NOT NULL,
                has_attachments BOOLEAN DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS attachments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                message_id INTEGER,
                url TEXT UNIQUE NOT NULL,
                filename TEXT,
                FOREIGN KEY (message_id) REFERENCES messages (id)
            );

            CREATE TABLE IF NOT EXISTS parsing_status (
                id INTEGER PRIMARY KEY CHECK (id = 1),
                first_parse_done BOOLEAN DEFAULT 0,
                last_parsed_message_id BIGINT,
                last_parse_time TIMESTAMP
            );

            CREATE INDEX IF NOT EXISTS idx_message_id ON messages(message_id);
            CREATE INDEX IF NOT EXISTS idx_author_id ON messages(author_id);
            CREATE INDEX IF NOT EXISTS idx_timestamp ON messages(timestamp);
        ''')
        await self._conn.commit()
        logger.info("Database initialized")

    async def close(self):
        if self._conn:
            await self._conn.close()

    def _parse_message_row(self, row) -> dict:
        message = dict(row)
        if message.get('attachment_urls_list'):
            urls = message['attachment_urls_list'].split(',')
            filenames = (message['attachment_filenames'] or '').split(',')
            message['attachments'] = [
                {'url': url, 'filename': filename}
                for url, filename in zip(urls, filenames)
            ]
        else:
            message['attachments'] = []
        return message

    async def get_parsing_status(self) -> dict:
        cursor = await self._conn.execute(
            'SELECT first_parse_done, last_parsed_message_id FROM parsing_status WHERE id = 1'
        )
        result = await cursor.fetchone()
        if result:
            return {'first_parse_done': bool(result[0]), 'last_parsed_message_id': result[1]}

        await self._conn.execute(
            'INSERT INTO parsing_status (id, first_parse_done, last_parsed_message_id) VALUES (1, 0, NULL)'
        )
        await self._conn.commit()
        return {'first_parse_done': False, 'last_parsed_message_id': None}

    async def update_parsing_status(self, first_parse_done=False, last_parsed_message_id=None):
        await self._conn.execute('''
            UPDATE parsing_status
            SET first_parse_done = ?,
                last_parsed_message_id = ?,
                last_parse_time = CURRENT_TIMESTAMP
            WHERE id = 1
        ''', (first_parse_done, last_parsed_message_id))
        await self._conn.commit()

    async def get_last_message_in_db(self):
        cursor = await self._conn.execute(
            'SELECT message_id FROM messages ORDER BY message_id DESC LIMIT 1'
        )
        result = await cursor.fetchone()
        return result[0] if result else None

    async def save_message(self, message_data: dict) -> bool:
        cursor = await self._conn.execute('''
            INSERT OR IGNORE INTO messages
            (message_id, channel_id, author_id, author_name, content,
             timestamp, message_url, has_attachments)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            message_data['message_id'],
            message_data['channel_id'],
            message_data['author_id'],
            message_data['author_name'],
            message_data['content'],
            message_data['timestamp'],
            message_data['message_url'],
            message_data['has_attachments'],
        ))

        if cursor.rowcount > 0:
            message_db_id = cursor.lastrowid
            for attachment in message_data.get('attachments', []):
                await self._conn.execute('''
                    INSERT OR IGNORE INTO attachments (message_id, url, filename)
                    VALUES (?, ?, ?)
                ''', (message_db_id, attachment['url'], attachment['filename']))
            await self._conn.commit()
            return True
        return False

    async def message_exists(self, message_id: int) -> bool:
        cursor = await self._conn.execute(
            'SELECT 1 FROM messages WHERE message_id = ? LIMIT 1', (message_id,)
        )
        return await cursor.fetchone() is not None

    async def get_random_message(self) -> dict | None:
        cursor = await self._conn.execute('''
            SELECT m.*,
                   GROUP_CONCAT(a.url) as attachment_urls_list,
                   GROUP_CONCAT(a.filename) as attachment_filenames
            FROM messages m
            LEFT JOIN attachments a ON m.id = a.message_id
            GROUP BY m.id
            ORDER BY RANDOM()
            LIMIT 1
        ''')
        row = await cursor.fetchone()
        return self._parse_message_row(row) if row else None

    async def get_message_by_number(self, number: int) -> dict | None:
        cursor = await self._conn.execute('''
            SELECT m.*,
                   GROUP_CONCAT(a.url) as attachment_urls_list,
                   GROUP_CONCAT(a.filename) as attachment_filenames
            FROM messages m
            LEFT JOIN attachments a ON m.id = a.message_id
            WHERE m.id = ?
            GROUP BY m.id
        ''', (number,))
        row = await cursor.fetchone()
        return self._parse_message_row(row) if row else None

    async def get_total_count(self) -> int:
        cursor = await self._conn.execute('SELECT COUNT(*) FROM messages')
        result = await cursor.fetchone()
        return result[0] if result else 0