stickerbot/database.py

"""Database module for sticker bot."""
import aiosqlite
from datetime import datetime
from typing import List, Optional, Dict, Any
from config import DATABASE_PATH, logger


class Database:
    """SQLite database handler for sticker bot."""
    
    def __init__(self, db_path: str = DATABASE_PATH):
        self.db_path = db_path
    
    async def init(self):
        """Initialize database tables."""
        async with aiosqlite.connect(self.db_path) as db:
            # Sticker sets table
            await db.execute("""
                CREATE TABLE IF NOT EXISTS sticker_sets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    set_name TEXT NOT NULL UNIQUE,
                    set_title TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # User premium status table
            await db.execute("""
                CREATE TABLE IF NOT EXISTS user_premium (
                    user_id INTEGER PRIMARY KEY,
                    is_premium BOOLEAN DEFAULT FALSE,
                    last_checked TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Track stickers for each set
            await db.execute("""
                CREATE TABLE IF NOT EXISTS stickers (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    set_name TEXT NOT NULL,
                    file_id TEXT NOT NULL,
                    emoji TEXT,
                    added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (set_name) REFERENCES sticker_sets(set_name) ON DELETE CASCADE
                )
            """)
            
            # API usage tracking (for remove.bg and other external services)
            await db.execute("""
                CREATE TABLE IF NOT EXISTS api_usage (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    service TEXT NOT NULL,
                    year_month TEXT NOT NULL,
                    usage_count INTEGER DEFAULT 0,
                    last_used TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(service, year_month)
                )
            """)
            
            await db.commit()
            logger.info("Database initialized successfully")
    
    async def add_sticker_set(self, user_id: int, set_name: str, set_title: str):
        """Add a new sticker set to tracking."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """INSERT OR REPLACE INTO sticker_sets (user_id, set_name, set_title, updated_at)
                   VALUES (?, ?, ?, CURRENT_TIMESTAMP)""",
                (user_id, set_name, set_title)
            )
            await db.commit()
    
    async def get_user_sticker_sets(self, user_id: int) -> List[Dict[str, Any]]:
        """Get all sticker sets for a user."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM sticker_sets WHERE user_id = ? ORDER BY created_at DESC",
                (user_id,)
            ) as cursor:
                rows = await cursor.fetchall()
                return [dict(row) for row in rows]
    
    async def delete_sticker_set(self, set_name: str):
        """Delete a sticker set from tracking."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("DELETE FROM sticker_sets WHERE set_name = ?", (set_name,))
            await db.commit()
    
    async def rename_sticker_set(self, set_name: str, new_title: str):
        """Update sticker set title."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """UPDATE sticker_sets SET set_title = ?, updated_at = CURRENT_TIMESTAMP
                   WHERE set_name = ?""",
                (new_title, set_name)
            )
            await db.commit()
    
    async def update_premium_status(self, user_id: int, is_premium: bool):
        """Update user's premium status."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """INSERT OR REPLACE INTO user_premium (user_id, is_premium, last_checked)
                   VALUES (?, ?, CURRENT_TIMESTAMP)""",
                (user_id, is_premium)
            )
            await db.commit()
    
    async def get_premium_status(self, user_id: int) -> Optional[bool]:
        """Get user's premium status."""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute(
                "SELECT is_premium FROM user_premium WHERE user_id = ?",
                (user_id,)
            ) as cursor:
                row = await cursor.fetchone()
                return row[0] if row else None
    
    async def add_sticker(self, set_name: str, file_id: str, emoji: Optional[str] = None):
        """Track a sticker in the database."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                "INSERT INTO stickers (set_name, file_id, emoji) VALUES (?, ?, ?)",
                (set_name, file_id, emoji)
            )
            await db.commit()
    
    async def delete_sticker(self, file_id: str):
        """Remove a sticker from tracking."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("DELETE FROM stickers WHERE file_id = ?", (file_id,))
            await db.commit()
    
    async def get_stickers_in_set(self, set_name: str) -> List[Dict[str, Any]]:
        """Get all stickers in a set."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM stickers WHERE set_name = ?",
                (set_name,)
            ) as cursor:
                rows = await cursor.fetchall()
                return [dict(row) for row in rows]
    
    async def track_api_call(self, service: str):
        """Track an API call for usage statistics."""
        year_month = datetime.now().strftime("%Y-%m")
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """INSERT INTO api_usage (service, year_month, usage_count, last_used)
                   VALUES (?, ?, 1, CURRENT_TIMESTAMP)
                   ON CONFLICT(service, year_month) 
                   DO UPDATE SET usage_count = usage_count + 1, last_used = CURRENT_TIMESTAMP""",
                (service, year_month)
            )
            await db.commit()
    
    async def get_monthly_api_usage(self, service: str) -> int:
        """Get current month's API usage count."""
        year_month = datetime.now().strftime("%Y-%m")
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute(
                "SELECT usage_count FROM api_usage WHERE service = ? AND year_month = ?",
                (service, year_month)
            ) as cursor:
                row = await cursor.fetchone()
                return row[0] if row else 0


# Global database instance
db = Database()