import sqlite3
import os

DB_FILE = "ovpn_profiler.db"

def migrate_db():
    if not os.path.exists(DB_FILE):
        print(f"Database file {DB_FILE} not found!")
        return

    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    
    try:
        cursor.execute("PRAGMA table_info(user_profiles)")
        columns = [info[1] for info in cursor.fetchall()]
        
        # Add is_revoked
        if "is_revoked" not in columns:
            print("Adding 'is_revoked' column...")
            cursor.execute("ALTER TABLE user_profiles ADD COLUMN is_revoked BOOLEAN DEFAULT 0")
        else:
            print("'is_revoked' column already exists.")

        # Add is_expired
        if "is_expired" not in columns:
            print("Adding 'is_expired' column...")
            cursor.execute("ALTER TABLE user_profiles ADD COLUMN is_expired BOOLEAN DEFAULT 0")
        else:
            print("'is_expired' column already exists.")
            
        # Ensure expiration_date exists
        if "expiration_date" not in columns:
             print("Adding 'expiration_date' column...")
             cursor.execute("ALTER TABLE user_profiles ADD COLUMN expiration_date DATETIME")
        else:
             print("'expiration_date' column already exists.")

        # Note: We do NOT remove 'expired_at' via ADD COLUMN script. 
        # SQLite does not support DROP COLUMN in older versions easily, 
        # and keeping it harmless is safer than complex migration logic.
        print("Migration successful.")
        conn.commit()

    except Exception as e:
        print(f"Migration failed: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_db()