database.py

"""Database module for storing bookmarks and scan results"""
import json
from typing import Any, Dict, List, Optional

import aiosqlite
from loguru import logger

import config


class Database:
    """Handle database operations for bookmarks and scan history"""
    
    def __init__(self, db_path: str = None):
        self.db_path = db_path or config.DATABASE_PATH
        
    async def initialize(self):
        """Create database tables if they don't exist"""
        async with aiosqlite.connect(self.db_path) as db:
            # Bookmarks table
            await db.execute('''
                CREATE TABLE IF NOT EXISTS bookmarks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    url TEXT NOT NULL,
                    title TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_scan TIMESTAMP,
                    last_status TEXT,
                    scan_count INTEGER DEFAULT 0,
                    UNIQUE(user_id, url)
                )
            ''')
            
            # Scan history table
            await db.execute('''
                CREATE TABLE IF NOT EXISTS scan_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    url TEXT NOT NULL,
                    scan_data TEXT NOT NULL,
                    scan_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # User preferences table
            await db.execute('''
                CREATE TABLE IF NOT EXISTS user_preferences (
                    user_id INTEGER PRIMARY KEY,
                    show_detailed BOOLEAN DEFAULT 1,
                    auto_bookmark BOOLEAN DEFAULT 0,
                    notification_enabled BOOLEAN DEFAULT 1,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Users table - track all users who have used the bot
            await db.execute('''
                CREATE TABLE IF NOT EXISTS users (
                    user_id INTEGER PRIMARY KEY,
                    username TEXT,
                    first_name TEXT,
                    last_name TEXT,
                    first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            await db.commit()
            logger.info("Database initialized successfully")
    
    async def close(self):
        """Close database connection - placeholder for cleanup"""
        # aiosqlite uses context managers, no persistent connection to close
        logger.info("Database cleanup complete")
    
    async def add_bookmark(
        self, 
        user_id: int, 
        url: str, 
        title: Optional[str] = None
    ) -> bool:
        """Add a bookmark for a user"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    '''INSERT OR IGNORE INTO bookmarks 
                       (user_id, url, title) VALUES (?, ?, ?)''',
                    (user_id, url, title or url)
                )
                await db.commit()
                logger.info(f"Bookmark added: {url} for user {user_id}")
                return True
        except Exception as e:
            logger.error(f"Error adding bookmark: {e}")
            return False
    
    async def remove_bookmark(self, user_id: int, url: str) -> bool:
        """Remove a bookmark"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    'DELETE FROM bookmarks WHERE user_id = ? AND url = ?',
                    (user_id, url)
                )
                await db.commit()
                logger.info(f"Bookmark removed: {url} for user {user_id}")
                return True
        except Exception as e:
            logger.error(f"Error removing bookmark: {e}")
            return False
    
    async def get_bookmarks(self, user_id: int) -> List[Dict[str, Any]]:
        """Get all bookmarks for a user"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                db.row_factory = aiosqlite.Row
                async with db.execute(
                    '''SELECT * FROM bookmarks 
                       WHERE user_id = ? 
                       ORDER BY created_at DESC''',
                    (user_id,)
                ) as cursor:
                    rows = await cursor.fetchall()
                    return [dict(row) for row in rows]
        except Exception as e:
            logger.error(f"Error getting bookmarks: {e}")
            return []
    
    async def is_bookmarked(self, user_id: int, url: str) -> bool:
        """Check if URL is bookmarked"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                async with db.execute(
                    'SELECT COUNT(*) FROM bookmarks WHERE user_id = ? AND url = ?',
                    (user_id, url)
                ) as cursor:
                    count = await cursor.fetchone()
                    return count[0] > 0
        except Exception as e:
            logger.error(f"Error checking bookmark: {e}")
            return False
    
    async def update_bookmark_status(
        self, 
        user_id: int, 
        url: str, 
        status: str
    ) -> bool:
        """Update bookmark status after scan"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    '''UPDATE bookmarks 
                       SET last_scan = CURRENT_TIMESTAMP,
                           last_status = ?,
                           scan_count = scan_count + 1
                       WHERE user_id = ? AND url = ?''',
                    (status, user_id, url)
                )
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Error updating bookmark: {e}")
            return False
    
    async def save_scan_history(
        self, 
        user_id: int, 
        url: str, 
        scan_data: Dict[str, Any]
    ) -> bool:
        """Save scan results to history"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    '''INSERT INTO scan_history 
                       (user_id, url, scan_data) VALUES (?, ?, ?)''',
                    (user_id, url, json.dumps(scan_data))
                )
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Error saving scan history: {e}")
            return False
    
    async def get_scan_history(
        self, 
        user_id: int, 
        url: Optional[str] = None,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """Get scan history for user or specific URL"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                db.row_factory = aiosqlite.Row
                
                if url:
                    query = '''SELECT * FROM scan_history 
                              WHERE user_id = ? AND url = ?
                              ORDER BY scan_date DESC LIMIT ?'''
                    params = (user_id, url, limit)
                else:
                    query = '''SELECT * FROM scan_history 
                              WHERE user_id = ?
                              ORDER BY scan_date DESC LIMIT ?'''
                    params = (user_id, limit)
                
                async with db.execute(query, params) as cursor:
                    rows = await cursor.fetchall()
                    results = []
                    for row in rows:
                        result = dict(row)
                        result['scan_data'] = json.loads(result['scan_data'])
                        results.append(result)
                    return results
        except Exception as e:
            logger.error(f"Error getting scan history: {e}")
            return []
    
    async def get_user_preferences(self, user_id: int) -> Dict[str, Any]:
        """Get user preferences"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                db.row_factory = aiosqlite.Row
                async with db.execute(
                    'SELECT * FROM user_preferences WHERE user_id = ?',
                    (user_id,)
                ) as cursor:
                    row = await cursor.fetchone()
                    if row:
                        return dict(row)
                    else:
                        # Return defaults
                        return {
                            'show_detailed': True,
                            'auto_bookmark': False,
                            'notification_enabled': True
                        }
        except Exception as e:
            logger.error(f"Error getting user preferences: {e}")
            return {}
    
    async def update_user_preferences(
        self, 
        user_id: int, 
        preferences: Dict[str, Any]
    ) -> bool:
        """Update user preferences"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                # Insert or replace
                await db.execute(
                    '''INSERT OR REPLACE INTO user_preferences 
                       (user_id, show_detailed, auto_bookmark, notification_enabled)
                       VALUES (?, ?, ?, ?)''',
                    (
                        user_id,
                        preferences.get('show_detailed', True),
                        preferences.get('auto_bookmark', False),
                        preferences.get('notification_enabled', True)
                    )
                )
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Error updating user preferences: {e}")
            return False
    
    async def get_stats(self, user_id: int) -> Dict[str, Any]:
        """Get user statistics"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                # Total bookmarks
                async with db.execute(
                    'SELECT COUNT(*) FROM bookmarks WHERE user_id = ?',
                    (user_id,)
                ) as cursor:
                    bookmark_count = (await cursor.fetchone())[0]
                
                # Total scans
                async with db.execute(
                    'SELECT COUNT(*) FROM scan_history WHERE user_id = ?',
                    (user_id,)
                ) as cursor:
                    scan_count = (await cursor.fetchone())[0]
                
                # Most scanned URL
                async with db.execute(
                    '''SELECT url, COUNT(*) as count 
                       FROM scan_history 
                       WHERE user_id = ? 
                       GROUP BY url 
                       ORDER BY count DESC 
                       LIMIT 1''',
                    (user_id,)
                ) as cursor:
                    most_scanned = await cursor.fetchone()
                
                return {
                    'total_bookmarks': bookmark_count,
                    'total_scans': scan_count,
                    'most_scanned_url': most_scanned[0] if most_scanned else None,
                    'most_scanned_count': most_scanned[1] if most_scanned else 0
                }
        except Exception as e:
            logger.error(f"Error getting stats: {e}")
            return {}
    
    async def get_user_stats(self, user_id: int) -> Dict[str, Any]:
        """Get user statistics - alias for get_stats"""
        stats = await self.get_stats(user_id)
        return {
            'total_scans': stats.get('total_scans', 0),
            'bookmarks_count': stats.get('total_bookmarks', 0),
            'most_scanned': stats.get('most_scanned_url'),
            'scan_count': stats.get('most_scanned_count', 0)
        }
    
    async def register_user(self, user_id: int, username: str = None, 
                           first_name: str = None, last_name: str = None) -> bool:
        """Register or update user in database"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    '''INSERT INTO users (user_id, username, first_name, last_name, first_seen, last_seen)
                       VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
                       ON CONFLICT(user_id) DO UPDATE SET
                       username = excluded.username,
                       first_name = excluded.first_name,
                       last_name = excluded.last_name,
                       last_seen = CURRENT_TIMESTAMP''',
                    (user_id, username, first_name, last_name)
                )
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Error registering user: {e}")
            return False
    
    async def get_total_users(self) -> int:
        """Get total number of unique users who have used the bot"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                async with db.execute('SELECT COUNT(*) FROM users') as cursor:
                    count = await cursor.fetchone()
                    return count[0] if count else 0
        except Exception as e:
            logger.error(f"Error getting total users: {e}")
            return 0