database.py

"""Module for database operations"""
import sqlite3
import json
from datetime import datetime
from typing import Optional, List, Dict, Any
from contextlib import contextmanager
import config


class Database:
    def __init__(self, db_path: str = config.DB_PATH):
        self.db_path = db_path
        self.init_db()
    
    @contextmanager
    def get_connection(self):
        """Context manager for database connection"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()
    
    def init_db(self):
        """Initialize database schema"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS chats (
                    chat_id INTEGER PRIMARY KEY,
                    chat_type TEXT,
                    title TEXT,
                    username TEXT,
                    first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_activity TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    message_id INTEGER,
                    chat_id INTEGER,
                    sender_id INTEGER,
                    sender_name TEXT,
                    text TEXT,
                    date TIMESTAMP,
                    media_type TEXT,
                    media_path TEXT,
                    is_edited BOOLEAN DEFAULT 0,
                    is_deleted BOOLEAN DEFAULT 0,
                    deleted_at TIMESTAMP,
                    reply_to_msg_id INTEGER,
                    raw_data TEXT,
                    FOREIGN KEY (chat_id) REFERENCES chats(chat_id),
                    UNIQUE(message_id, chat_id)
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS message_edits (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    message_id INTEGER,
                    chat_id INTEGER,
                    old_text TEXT,
                    new_text TEXT,
                    edited_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (message_id, chat_id) REFERENCES messages(message_id, chat_id)
                )
            ''')
            
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_messages_chat ON messages(chat_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_messages_date ON messages(date)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_messages_deleted ON messages(is_deleted)')
    
    def save_chat(self, chat_id: int, chat_type: str, title: str, username: Optional[str] = None):
        """Save chat information"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO chats (chat_id, chat_type, title, username, last_activity)
                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT(chat_id) DO UPDATE SET
                    title = excluded.title,
                    username = excluded.username,
                    last_activity = CURRENT_TIMESTAMP
            ''', (chat_id, chat_type, title, username))
    
    def save_message(self, message_data: Dict[str, Any], mark_edited: bool = False) -> int:
        """Save message"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            if mark_edited:
                cursor.execute('''
                    INSERT INTO messages (
                        message_id, chat_id, sender_id, sender_name, text, 
                        date, media_type, media_path, reply_to_msg_id, raw_data
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(message_id, chat_id) DO UPDATE SET
                        text = excluded.text,
                        media_type = excluded.media_type,
                        media_path = excluded.media_path,
                        is_edited = 1
                ''', (
                    message_data['message_id'],
                    message_data['chat_id'],
                    message_data['sender_id'],
                    message_data['sender_name'],
                    message_data['text'],
                    message_data['date'],
                    message_data.get('media_type'),
                    message_data.get('media_path'),
                    message_data.get('reply_to_msg_id'),
                    json.dumps(message_data.get('raw_data', {}))
                ))
            else:
                cursor.execute('''
                    INSERT INTO messages (
                        message_id, chat_id, sender_id, sender_name, text, 
                        date, media_type, media_path, reply_to_msg_id, raw_data
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(message_id, chat_id) DO NOTHING
                ''', (
                message_data['message_id'],
                message_data['chat_id'],
                message_data['sender_id'],
                message_data['sender_name'],
                message_data['text'],
                message_data['date'],
                message_data.get('media_type'),
                message_data.get('media_path'),
                message_data.get('reply_to_msg_id'),
                json.dumps(message_data.get('raw_data', {}))
                ))
            return cursor.lastrowid
    
    def mark_message_deleted(self, message_id: int, chat_id: int):
        """Mark message as deleted"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE messages 
                SET is_deleted = 1, deleted_at = CURRENT_TIMESTAMP
                WHERE message_id = ? AND chat_id = ?
            ''', (message_id, chat_id))
            return cursor.rowcount > 0
    
    def find_message_chat_id(self, message_id: int) -> Optional[int]:
        """Find chat_id by message_id"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT chat_id FROM messages
                WHERE message_id = ?
                LIMIT 1
            ''', (message_id,))
            result = cursor.fetchone()
            return result[0] if result else None
    
    def save_message_edit(self, message_id: int, chat_id: int, old_text: str, new_text: str):
        """Save message edit history"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO message_edits (message_id, chat_id, old_text, new_text)
                VALUES (?, ?, ?, ?)
            ''', (message_id, chat_id, old_text, new_text))
    
    def get_chats(self, limit: int = 100) -> List[Dict]:
        """Get list of chats"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT c.*, COUNT(m.id) as message_count
                FROM chats c
                LEFT JOIN messages m ON c.chat_id = m.chat_id
                GROUP BY c.chat_id
                ORDER BY c.last_activity DESC
                LIMIT ?
            ''', (limit,))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_messages(self, chat_id: int, limit: int = 100, offset: int = 0, 
                     include_deleted: bool = True) -> List[Dict]:
        """Get messages from chat"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            query = '''
                SELECT * FROM messages 
                WHERE chat_id = ?
            '''
            params = [chat_id]
            
            if not include_deleted:
                query += ' AND is_deleted = 0'
            
            query += ' ORDER BY date DESC LIMIT ? OFFSET ?'
            params.extend([limit, offset])
            
            cursor.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]
    
    def search_messages(self, query: str, chat_id: Optional[int] = None, 
                       limit: int = 100) -> List[Dict]:
        """Search messages by text"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            sql = '''
                SELECT m.*, c.title as chat_title
                FROM messages m
                JOIN chats c ON m.chat_id = c.chat_id
                WHERE m.text LIKE ?
            '''
            params = [f'%{query}%']
            
            if chat_id:
                sql += ' AND m.chat_id = ?'
                params.append(chat_id)
            
            sql += ' ORDER BY m.date DESC LIMIT ?'
            params.append(limit)
            
            cursor.execute(sql, params)
            return [dict(row) for row in cursor.fetchall()]
    
    def get_message_edits(self, message_id: int, chat_id: int) -> List[Dict]:
        """Get message edit history"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM message_edits
                WHERE message_id = ? AND chat_id = ?
                ORDER BY edited_at DESC
            ''', (message_id, chat_id))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_stats(self) -> Dict[str, Any]:
        """Database statistics"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            cursor.execute('SELECT COUNT(*) FROM chats')
            total_chats = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages')
            total_messages = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE is_deleted = 1')
            deleted_messages = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM messages WHERE is_edited = 1')
            edited_messages = cursor.fetchone()[0]
            
            return {
                'total_chats': total_chats,
                'total_messages': total_messages,
                'deleted_messages': deleted_messages,
                'edited_messages': edited_messages
            }