database.py

import sqlite3
import os
import time

DB_PATH = "vqa.db"

def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # Users table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT
        )
    ''')
    
    # Logs table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            username TEXT,
            model TEXT,
            question TEXT
        )
    ''')

    # Gemini Usage table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS gemini_usage (
            username TEXT,
            date TEXT,
            call_count INTEGER,
            PRIMARY KEY (username, date)
        )
    ''')

    # Seed Admin
    cursor.execute("INSERT OR IGNORE INTO users (username, password) VALUES (?, ?)", ("vignesh", "vignesh"))
    
    conn.commit()
    conn.close()

def login_user(username, password):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT password FROM users WHERE username=?", (username,))
    row = cursor.fetchone()
    conn.close()
    
    if row and row[0] == password:
        return True
    return False

def signup_user(username, password):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    try:
        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password))
        conn.commit()
        success = True
    except sqlite3.IntegrityError:
        success = False
    conn.close()
    return success

def log_request_db(username, model, question):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    cursor.execute("INSERT INTO logs (timestamp, username, model, question) VALUES (?, ?, ?, ?)",
                   (timestamp, username, model, question))
    conn.commit()
    conn.close()

def check_gemini_limit(username):
    if username == "vignesh":
        return True # Admin bypass

    from datetime import datetime
    today_str = datetime.now().strftime("%Y-%m-%d")
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("SELECT call_count FROM gemini_usage WHERE username=? AND date=?", (username, today_str))
    row = cursor.fetchone()
    
    if row is None:
        count = 0
    else:
        count = row[0]
        
    conn.close()
    
    if count >= 10:
        return False
    return True

def increment_gemini_limit(username):
    if username == "vignesh":
        return

    from datetime import datetime
    today_str = datetime.now().strftime("%Y-%m-%d")
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("SELECT call_count FROM gemini_usage WHERE username=? AND date=?", (username, today_str))
    row = cursor.fetchone()
    
    if row is None:
        cursor.execute("INSERT INTO gemini_usage (username, date, call_count) VALUES (?, ?, ?)", (username, today_str, 1))
    else:
        cursor.execute("UPDATE gemini_usage SET call_count = call_count + 1 WHERE username=? AND date=?", (username, today_str))
        
    conn.commit()
    conn.close()

def login_or_create_google_user(username):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    # Check if user exists
    cursor.execute("SELECT username FROM users WHERE username=?", (username,))
    row = cursor.fetchone()
    if not row:
        # Create user with a dummy password since Google handles auth
        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, "GOOGLE_AUTH_USER"))
        conn.commit()
    conn.close()
    return True

def get_user_logs(username):
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    if username == "vignesh":
        cursor.execute("SELECT * FROM logs ORDER BY id DESC")
    else:
        cursor.execute("SELECT * FROM logs WHERE username=? ORDER BY id DESC", (username,))
        
    rows = cursor.fetchall()
    conn.close()
    
    return [dict(row) for row in rows]