import mysql.connector
import os
from urllib.parse import urlparse

# Get DB config from connection URL or assume defaults
db_url = os.environ.get('SQLALCHEMY_DATABASE_URI', 'mysql+mysqlconnector://root:@localhost:3306/ashlya_academy_lms')
parsed = urlparse(db_url)

try:
    conn = mysql.connector.connect(
        host=parsed.hostname or 'localhost',
        user=parsed.username or 'root',
        password=parsed.password or '',
        database=parsed.path[1:] if parsed.path else 'ashlya_academy_lms',
        port=parsed.port or 3306
    )
    cursor = conn.cursor()
    
    # Try adding the column
    try:
        cursor.execute("ALTER TABLE ai_models ADD COLUMN is_free BOOLEAN DEFAULT FALSE;")
        print("Successfully added is_free column to ai_models.")
    except Exception as e:
        print(f"Adding column failed (maybe it already exists?): {str(e)}")

    # Also try adding auth_required since it's around the same place in the model
    try:
        cursor.execute("ALTER TABLE ai_models ADD COLUMN auth_required BOOLEAN DEFAULT FALSE;")
        print("Successfully added auth_required column to ai_models.")
    except Exception as e:
        print(f"Adding auth_required column failed: {str(e)}")
        
    conn.commit()
    cursor.close()
    conn.close()
    
except Exception as e:
    print(f"Connection failed: {str(e)}")
