database.py

# -----------------------------------------------------------------------------
# Copyright (c) 2025 tanbaycu. All rights reserved.
# Licensed under the MIT License. See LICENSE file for details.
# -----------------------------------------------------------------------------

import sqlite3

import datetime
import logging

DB_NAME = "bot_data.db"
logger = logging.getLogger(__name__)

def init_db():
    """Initializes the database with tables if they don't exist."""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    
    # Users Table
    c.execute('''CREATE TABLE IF NOT EXISTS users (
                    user_id INTEGER PRIMARY KEY,
                    username TEXT,
                    balance REAL DEFAULT 2.0,
                    is_premium INTEGER DEFAULT 0,
                    joined_date TEXT
                )''')
    
    # Transactions Table
    c.execute('''CREATE TABLE IF NOT EXISTS transactions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER,
                    amount REAL,
                    type TEXT,
                    timestamp TEXT,
                    note TEXT,
                    FOREIGN KEY(user_id) REFERENCES users(user_id)
                )''')
    
    conn.commit()
    conn.close()
    logger.info("Database initialized successfully.")

def get_user(user_id, username=None):
    """Gets user data, creating if not exists."""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    
    c.execute("SELECT * FROM users WHERE user_id=?", (user_id,))
    user = c.fetchone()
    
    if not user:
        # New User - Free 2.0 Xu
        joined_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        c.execute("INSERT INTO users (user_id, username, balance, is_premium, joined_date) VALUES (?, ?, ?, ?, ?)",
                  (user_id, username, 2.0, 0, joined_date))
        conn.commit()
        # Fetch the new user
        c.execute("SELECT * FROM users WHERE user_id=?", (user_id,))
        user = c.fetchone()
        
        # Log welcome bonus transaction
        add_transaction(user_id, 2.0, "BONUS", "Quà tặng thành viên mới")
    
    conn.close()
    return {
        "user_id": user[0],
        "username": user[1],
        "balance": user[2],
        "is_premium": bool(user[3]),
        "joined_date": user[4]
    }

def update_balance(user_id, amount, reason=""):
    """Updates user balance and logs transaction. Amount can be negative."""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    
    # Update Balance
    c.execute("UPDATE users SET balance = balance + ? WHERE user_id=?", (amount, user_id))
    
    # Log Transaction
    add_transaction_conn(c, user_id, amount, "USAGE" if amount < 0 else "DEPOSIT", reason)
    
    conn.commit()
    conn.close()

def add_transaction(user_id, amount, trans_type, note=""):
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    add_transaction_conn(c, user_id, amount, trans_type, note)
    conn.commit()
    conn.close()

def add_transaction_conn(cursor, user_id, amount, trans_type, note):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    cursor.execute("INSERT INTO transactions (user_id, amount, type, timestamp, note) VALUES (?, ?, ?, ?, ?)",
                   (user_id, amount, trans_type, timestamp, note))

def get_transactions(user_id, limit=5):
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute("SELECT * FROM transactions WHERE user_id=? ORDER BY id DESC LIMIT ?", (user_id, limit))
    rows = c.fetchall()
    conn.close()
    return rows

def get_all_users():
    """Returns a list of all user_ids."""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute("SELECT user_id FROM users")
    rows = c.fetchall()
    conn.close()
    return [row[0] for row in rows]