import pymysql
import os
import sys
from pymysql.cursors import DictCursor

def get_db_connection():
    try:
        return pymysql.connect(
            host=os.getenv("MARIADB_HOST", "127.0.0.1"),
            port=int(os.getenv("MARIADB_PORT", 3306)),
            user=os.getenv("MARIADB_USER", "crumb"),
            password=os.getenv("MARIADB_PASSWORD", "secret"),
            database=os.getenv("MARIADB_DATABASE", "crumbcrm"),
            autocommit=True,
            charset="utf8mb4",
            cursorclass=DictCursor,
        )
    except Exception as e:
        print(f"Error connecting to database: {e}")
        return None

def migrate():
    print("🚀 Starting user migration...")
    conn = get_db_connection()
    if not conn:
        print("❌ Could not connect to database. Please ensure MariaDB is running and accessible.")
        sys.exit(1)

    try:
        with conn.cursor() as cursor:
            # Check if columns exist
            cursor.execute("DESCRIBE users")
            columns = [row['Field'] for row in cursor.fetchall()]
            
            alter_statements = []
            
            if 'user_group' not in columns:
                alter_statements.append("ADD COLUMN user_group VARCHAR(50) DEFAULT 'demo'")
            
            if 'theme' not in columns:
                alter_statements.append("ADD COLUMN theme VARCHAR(50) DEFAULT 'pico-default'")
            
            if 'accessibility' not in columns:
                alter_statements.append("ADD COLUMN accessibility JSON DEFAULT NULL")
            
            if alter_statements:
                sql = f"ALTER TABLE users {', '.join(alter_statements)}"
                print(f"Executing: {sql}")
                cursor.execute(sql)
                print("✅ Migration successful: Added new columns.")
            else:
                print("ℹ️  Migration skipped: Columns already exist.")

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

if __name__ == "__main__":
    migrate()