database.py

"""
database.py — Banco de dados SQLite para logs de download e estado do bot.
"""

import os
import sqlite3
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

_DATA_DIR = "/data" if os.path.isdir("/data") else "."
DB_PATH = os.getenv("DB_PATH", os.path.join(_DATA_DIR, "downloads.db"))


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    """Cria as tabelas se não existirem."""
    with _connect() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS downloads (
                id             INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id        INTEGER,
                username       TEXT,
                first_name     TEXT,
                chat_id        INTEGER,
                request_time   TEXT,
                video_url      TEXT,
                video_title    TEXT,
                format_chosen  TEXT,
                file_size      INTEGER,
                status         TEXT,
                error_message  TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS bot_state (
                key   TEXT PRIMARY KEY,
                value TEXT
            )
        """)
        conn.execute("""
            INSERT OR IGNORE INTO bot_state (key, value) VALUES ('maintenance', '0')
        """)
        conn.commit()
    logger.info(f"[DB] Banco de dados em: {DB_PATH}")


def log_download(
    user_id: int,
    username: str | None,
    first_name: str | None,
    chat_id: int,
    video_url: str,
    video_title: str,
    format_chosen: str,
    file_size: int | None,
    status: str,
    error_message: str | None = None,
) -> None:
    with _connect() as conn:
        conn.execute(
            """
            INSERT INTO downloads
              (user_id, username, first_name, chat_id, request_time,
               video_url, video_title, format_chosen, file_size, status, error_message)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                user_id, username, first_name, chat_id,
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                video_url, video_title, format_chosen, file_size, status, error_message,
            ),
        )
        conn.commit()


def get_recent_downloads(limit: int = 10) -> list[dict]:
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT user_id, username, first_name, request_time,
                   video_url, video_title, format_chosen, file_size, status, error_message
            FROM downloads
            ORDER BY id DESC
            LIMIT ?
            """,
            (limit,),
        ).fetchall()
    return [dict(r) for r in rows]


def get_stats() -> dict:
    with _connect() as conn:
        total = conn.execute("SELECT COUNT(*) FROM downloads").fetchone()[0]
        success = conn.execute(
            "SELECT COUNT(*) FROM downloads WHERE status = 'sucesso'"
        ).fetchone()[0]
        errors = conn.execute(
            "SELECT COUNT(*) FROM downloads WHERE status = 'erro'"
        ).fetchone()[0]
    return {"total": total, "success": success, "errors": errors}


def is_maintenance() -> bool:
    with _connect() as conn:
        row = conn.execute(
            "SELECT value FROM bot_state WHERE key = 'maintenance'"
        ).fetchone()
    return row is not None and row[0] == "1"


def set_maintenance(enabled: bool) -> None:
    with _connect() as conn:
        conn.execute(
            "INSERT OR REPLACE INTO bot_state (key, value) VALUES ('maintenance', ?)",
            ("1" if enabled else "0",),
        )
        conn.commit()