db/migration.py

import os
from dotenv import load_dotenv
from sqlalchemy import create_engine, text
import sys

# Add the parent directory to Python path to import your connection module
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Import your existing database connection logic
from db.connection import get_database_config

def get_database_url():
    """Get database URL using the same logic as your main app"""
    db_config = get_database_config()
    
    if not db_config:
        print("❌ No database configuration found")
        return None
    
    # Construct connection URL
    DATABASE_URL = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"
    
    # Add SSL mode for Supabase
    if db_config['host'].endswith('supabase.co'):
        DATABASE_URL += "?sslmode=require"
    
    print(f"🔗 Using database: {db_config['host']}:{db_config['port']}")
    return DATABASE_URL

def add_generated_model_column():
    """Add generated_model column to content_history table"""
    print("🔄 Running database migration: adding generated_model column...")
    
    DATABASE_URL = get_database_url()
    if not DATABASE_URL:
        print("❌ Cannot proceed without database configuration")
        return
    
    # Create engine
    engine = create_engine(DATABASE_URL)
    
    with engine.connect() as conn:
        try:
            # Check if column already exists
            result = conn.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name='content_history' AND column_name='generated_model'
            """))
            
            if result.fetchone() is None:
                print("📝 Adding generated_model column...")
                
                # Add the column
                conn.execute(text("""
                    ALTER TABLE content_history 
                    ADD COLUMN generated_model VARCHAR DEFAULT 'groq'
                """))
                print("✅ Added generated_model column to content_history table")
                
                # Update existing records (set all existing content to groq)
                conn.execute(text("""
                    UPDATE content_history 
                    SET generated_model = 'groq' 
                    WHERE generated_model IS NULL
                """))
                print("✅ Set existing records to 'groq'")
                
                conn.commit()
                print("🎉 Migration completed successfully!")
            else:
                print("✅ generated_model column already exists")
                
        except Exception as e:
            print(f"❌ Migration failed: {e}")
            conn.rollback()
            raise

def add_auth_columns_to_users():
    """Add authentication columns to users table - FIXED for existing data"""
    print("🔄 Running database migration: adding authentication columns to users table...")
    
    DATABASE_URL = get_database_url()
    if not DATABASE_URL:
        print("❌ Cannot proceed without database configuration")
        return
    
    # Create engine
    engine = create_engine(DATABASE_URL)
    
    with engine.connect() as conn:
        try:
            # Check if users table exists first
            result = conn.execute(text("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_name='users'
            """))
            
            if result.fetchone() is None:
                print("❌ Users table doesn't exist. Please run the initial database setup first.")
                return
            
            # List of columns to add with their definitions
            columns_to_add = [
                {
                    'name': 'fullname',
                    'definition': 'VARCHAR(200) NOT NULL DEFAULT \'\'',
                    'update_existing': "UPDATE users SET fullname = 'User' WHERE fullname IS NULL OR fullname = ''"
                },
                {
                    'name': 'username', 
                    'definition': 'VARCHAR(50) NOT NULL DEFAULT \'\'',  # Remove UNIQUE constraint initially
                    'update_existing': "UPDATE users SET username = 'user_' || id::text WHERE username IS NULL OR username = ''"
                },
                {
                    'name': 'password_hash',
                    'definition': 'VARCHAR(128) NOT NULL DEFAULT \'\'',
                    'update_existing': None
                },
                {
                    'name': 'salt',
                    'definition': 'VARCHAR(32) NOT NULL DEFAULT \'\'',
                    'update_existing': None
                },
                {
                    'name': 'last_login',
                    'definition': 'TIMESTAMP',
                    'update_existing': None
                }
            ]
            
            added_columns = []
            
            for column in columns_to_add:
                # Check if column already exists
                result = conn.execute(text(f"""
                    SELECT column_name 
                    FROM information_schema.columns 
                    WHERE table_name='users' AND column_name='{column['name']}'
                """))
                
                if result.fetchone() is None:
                    print(f"📝 Adding {column['name']} column...")
                    
                    try:
                        # Add the column without UNIQUE constraint initially
                        conn.execute(text(f"""
                            ALTER TABLE users 
                            ADD COLUMN {column['name']} {column['definition']}
                        """))
                        
                        # Update existing records if needed
                        if column['update_existing']:
                            conn.execute(text(column['update_existing']))
                            print(f"✅ Updated existing records for {column['name']}")
                        
                        added_columns.append(column['name'])
                        print(f"✅ Added {column['name']} column to users table")
                        
                    except Exception as column_error:
                        print(f"⚠️ Failed to add {column['name']} column: {column_error}")
                        continue
                else:
                    print(f"✅ {column['name']} column already exists")
            
            # NOW add the UNIQUE constraint after all usernames are populated
            if 'username' in added_columns:
                print("🔧 Adding UNIQUE constraint to username column...")
                try:
                    conn.execute(text("""
                        ALTER TABLE users 
                        ADD CONSTRAINT users_username_unique UNIQUE (username)
                    """))
                    print("✅ Added UNIQUE constraint to username column")
                except Exception as e:
                    print(f"⚠️ Could not add UNIQUE constraint: {e}")
                    # This might happen if there are duplicate usernames, but we generated unique ones
            
            if added_columns:
                conn.commit()
                print(f"🎉 Migration completed successfully! Added columns: {', '.join(added_columns)}")
            else:
                print("✅ All authentication columns already exist")
                
            # Create index on username for faster lookups
            print("🔧 Creating index on username column...")
            conn.execute(text("""
                CREATE INDEX IF NOT EXISTS idx_users_username ON users(username)
            """))
            print("✅ Username index created")
            
            conn.commit()
                
        except Exception as e:
            print(f"❌ Migration failed: {e}")
            conn.rollback()
            raise

def migrate_existing_users():
    """Migrate existing users to have proper usernames and fullnames"""
    print("🔄 Migrating existing users to new authentication system...")
    
    DATABASE_URL = get_database_url()
    if not DATABASE_URL:
        print("❌ Cannot proceed without database configuration")
        return
    
    engine = create_engine(DATABASE_URL)
    
    with engine.connect() as conn:
        try:
            # Check if we have users without proper usernames
            result = conn.execute(text("""
                SELECT COUNT(*) as count 
                FROM users 
                WHERE username = '' OR username IS NULL
            """))
            
            users_without_username = result.fetchone()[0]
            
            if users_without_username > 0:
                print(f"📝 Found {users_without_username} users without usernames, migrating...")
                
                # Generate unique usernames for existing users
                conn.execute(text("""
                    UPDATE users 
                    SET 
                        username = 'user_' || id::text,
                        fullname = 'User ' || id::text
                    WHERE username = '' OR username IS NULL
                """))
                
                print(f"✅ Migrated {users_without_username} users to new authentication system")
            else:
                print("✅ All users already have usernames")
                
            conn.commit()
                
        except Exception as e:
            print(f"❌ User migration failed: {e}")
            conn.rollback()
            raise

def run_all_migrations():
    """Run all database migrations"""
    print("🚀 Starting database migrations...")
    
    try:
        # 1. Add generated_model column to content_history
        add_generated_model_column()
        
        print("\n" + "="*50 + "\n")
        
        # 2. Add authentication columns to users table
        add_auth_columns_to_users()
        
        print("\n" + "="*50 + "\n")
        
        # 3. Migrate existing users
        migrate_existing_users()
        
        print("\n🎉 All migrations completed successfully!")
        
    except Exception as e:
        print(f"❌ Migration process failed: {e}")
        raise

if __name__ == "__main__":
    # Load environment variables
    load_dotenv()
    
    # Run migrations
    run_all_migrations()