#!/usr/bin/env python3
"""
Fix Portfolio Manager Column
Quick fix for missing portfolio_manager_id and verified_by_id columns
"""

import os
import sys
import django
from pathlib import Path

# Setup Django environment
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
django.setup()

from django.db import connection

def add_missing_portfolio_columns():
    """Add missing portfolio_manager_id and verified_by_id columns"""
    print("🔧 Adding missing portfolio columns...")
    
    with connection.cursor() as cursor:
        # Get current database name
        cursor.execute("SELECT DATABASE()")
        db_name = cursor.fetchone()[0]
        print(f"Working on database: {db_name}")
        
        # Add portfolio_manager_id column
        try:
            cursor.execute("""
                SELECT COLUMN_NAME 
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = %s 
                AND TABLE_NAME = 'users' 
                AND COLUMN_NAME = 'portfolio_manager_id'
            """, [db_name])
            
            if not cursor.fetchone():
                cursor.execute("ALTER TABLE users ADD COLUMN portfolio_manager_id CHAR(32) NULL")
                print("✅ Added users.portfolio_manager_id")
            else:
                print("✓ users.portfolio_manager_id already exists")
                
        except Exception as e:
            print(f"❌ Failed to add portfolio_manager_id: {e}")
        
        # Add verified_by_id column
        try:
            cursor.execute("""
                SELECT COLUMN_NAME 
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = %s 
                AND TABLE_NAME = 'users' 
                AND COLUMN_NAME = 'verified_by_id'
            """, [db_name])
            
            if not cursor.fetchone():
                cursor.execute("ALTER TABLE users ADD COLUMN verified_by_id CHAR(32) NULL")
                print("✅ Added users.verified_by_id")
            else:
                print("✓ users.verified_by_id already exists")
                
        except Exception as e:
            print(f"❌ Failed to add verified_by_id: {e}")
        
        # Add foreign key constraint for verified_by_id
        try:
            cursor.execute("ALTER TABLE users ADD CONSTRAINT fk_users_verified_by FOREIGN KEY (verified_by_id) REFERENCES users (id) ON DELETE SET NULL")
            print("✅ Added foreign key constraint for verified_by_id")
        except Exception as e:
            if "Duplicate key name" in str(e) or "already exists" in str(e):
                print("✓ Foreign key constraint already exists")
            else:
                print(f"⚠️ Could not add foreign key constraint: {e}")
        
        # Add index for portfolio_manager_id
        try:
            cursor.execute("CREATE INDEX idx_users_portfolio_manager ON users (portfolio_manager_id)")
            print("✅ Added index for portfolio_manager_id")
        except Exception as e:
            if "Duplicate key name" in str(e):
                print("✓ Index for portfolio_manager_id already exists")
            else:
                print(f"⚠️ Could not add index: {e}")

def test_login_access():
    """Test if login will work now"""
    print("\n🧪 Testing login access...")
    
    try:
        from users.models import CustomUser
        
        # Test query that includes portfolio_manager_id
        user = CustomUser.objects.first()
        if user:
            print(f"✅ Login test passed - can access user: {user.username}")
            print(f"   Portfolio manager ID: {getattr(user, 'portfolio_manager_id', 'None')}")
            print(f"   Verified by ID: {getattr(user, 'verified_by_id', 'None')}")
        else:
            print("✅ Login test passed - no users found but query works")
        
        return True
        
    except Exception as e:
        print(f"❌ Login test failed: {e}")
        return False

def main():
    """Main function to fix portfolio manager column issue"""
    print("🔧 HAVEN GRAZURI INVESTMENT LIMITED- Portfolio Manager Column Fix")
    print("=" * 60)
    
    # Add missing columns
    add_missing_portfolio_columns()
    
    # Test login access
    test_passed = test_login_access()
    
    print("\n" + "=" * 60)
    print("🎉 Portfolio manager column fix completed!")
    
    print(f"\n📊 SUMMARY:")
    print(f"   • Added portfolio_manager_id column")
    print(f"   • Added verified_by_id column")
    print(f"   • Added foreign key constraint")
    print(f"   • Added database index")
    print(f"   • Login test: {'PASSED' if test_passed else 'FAILED'}")
    
    print(f"\n🎯 NEXT STEPS:")
    print("   1. Test login page immediately")
    print("   2. Check user profile pages")
    print("   3. Test portfolio management features")
    
    if test_passed:
        print(f"\n✅ SUCCESS: Login should work now!")
    else:
        print(f"\n⚠️ WARNING: Login may still have issues")

if __name__ == "__main__":
    main()