#!/usr/bin/env python3
"""
COMPLETE SCRIPT TO FIX ALL CLIENT ASSIGNMENT DATABASE ISSUES
This script fixes:
1. Missing columns (is_active, unassigned_date, reason)
2. Wrong column types (UUID fields as integers)
3. ID field default value issue
4. Any other database schema issues
"""

import os
import sys
import django
from django.db import connection

def setup_django():
    """Setup Django environment"""
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
    django.setup()

def fix_portfolio_assignments_table():
    """Fix all issues with portfolio_assignments table"""
    
    print("🔧 Fixing portfolio_assignments table...")
    
    with connection.cursor() as cursor:
        # Check if table exists
        cursor.execute("""
            SELECT COUNT(*) 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'portfolio_assignments'
        """)
        
        if cursor.fetchone()[0] == 0:
            print("❌ portfolio_assignments table does not exist!")
            return False
        
        print("✅ portfolio_assignments table exists")
        
        # Get current table structure
        cursor.execute("""
            SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'portfolio_assignments'
            ORDER BY ORDINAL_POSITION
        """)
        
        current_columns = {row[0]: {'type': row[1], 'nullable': row[2], 'default': row[3]} 
                          for row in cursor.fetchall()}
        
        print(f"Current columns: {list(current_columns.keys())}")
        
        # Fix 1: Add missing columns and fix existing ones
        missing_columns = {
            'is_active': 'BOOLEAN DEFAULT TRUE',
            'unassigned_date': 'DATETIME NULL',
            'reason': 'TEXT NULL'
        }
        
        # Fix existing columns that need default values
        columns_to_fix = {
            'assignment_number': 'VARCHAR(50) NULL'
        }
        
        for col_name, col_def in missing_columns.items():
            if col_name not in current_columns:
                print(f"Adding missing column: {col_name}")
                try:
                    with connection.cursor() as add_cursor:
                        add_cursor.execute(f"""
                            ALTER TABLE portfolio_assignments 
                            ADD COLUMN {col_name} {col_def}
                        """)
                    print(f"✅ Added column: {col_name}")
                except Exception as e:
                    print(f"❌ Error adding {col_name}: {e}")
            else:
                print(f"✅ Column exists: {col_name}")
        
        # Fix existing columns that need default values
        for col_name, col_def in columns_to_fix.items():
            if col_name in current_columns:
                col_info = current_columns[col_name]
                if col_info['default'] is None:
                    print(f"Fixing column default value: {col_name}")
                    try:
                        with connection.cursor() as fix_cursor:
                            fix_cursor.execute(f"""
                                ALTER TABLE portfolio_assignments 
                                MODIFY COLUMN {col_name} {col_def}
                            """)
                        print(f"✅ Fixed column: {col_name}")
                    except Exception as e:
                        print(f"❌ Error fixing {col_name}: {e}")
                else:
                    print(f"✅ Column already has default: {col_name}")
            else:
                print(f"ℹ️  Column doesn't exist: {col_name}")
        
        # Fix 2: Fix UUID column types
        uuid_columns = ['client_id', 'portfolio_manager_id', 'assigned_by_id']
        
        for col_name in uuid_columns:
            if col_name in current_columns:
                current_type = current_columns[col_name]['type']
                if current_type == 'int':
                    print(f"Fixing column type: {col_name} (int -> char)")
                    try:
                        with connection.cursor() as fix_cursor:
                            fix_cursor.execute(f"""
                                ALTER TABLE portfolio_assignments 
                                MODIFY COLUMN {col_name} CHAR(36) NULL
                            """)
                        print(f"✅ Fixed column type: {col_name}")
                    except Exception as e:
                        print(f"❌ Error fixing {col_name}: {e}")
                else:
                    print(f"✅ Column type correct: {col_name} ({current_type})")
        
        # Fix 3: Fix ID column - make it auto-increment integer
        if 'id' in current_columns:
            id_info = current_columns['id']
            if id_info['type'] == 'char' or id_info['default'] is None:
                print("Fixing ID column to auto-increment integer...")
                try:
                    with connection.cursor() as id_cursor:
                        # First, drop existing primary key if it exists
                        try:
                            id_cursor.execute("ALTER TABLE portfolio_assignments DROP PRIMARY KEY")
                            print("✅ Dropped existing primary key")
                        except Exception as drop_error:
                            print(f"ℹ️  No existing primary key to drop: {drop_error}")
                        
                        # Now modify ID column to auto-increment integer with primary key
                        id_cursor.execute("""
                            ALTER TABLE portfolio_assignments 
                            MODIFY COLUMN id INT AUTO_INCREMENT PRIMARY KEY
                        """)
                    print("✅ Fixed ID column to auto-increment integer")
                except Exception as e:
                    print(f"❌ Error fixing ID column: {e}")
            else:
                print(f"✅ ID column correct: {id_info}")
        
        return True

def create_portfolio_assignments_table():
    """Create portfolio_assignments table if it doesn't exist"""
    
    print("🔧 Creating portfolio_assignments table...")
    
    with connection.cursor() as cursor:
        try:
            cursor.execute("""
                CREATE TABLE portfolio_assignments (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    client_id CHAR(36) NOT NULL,
                    portfolio_manager_id CHAR(36) NOT NULL,
                    assigned_by_id CHAR(36) NULL,
                    assigned_date DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                    unassigned_date DATETIME NULL,
                    is_active BOOLEAN DEFAULT TRUE,
                    reason TEXT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    INDEX idx_client_active (client_id, is_active),
                    INDEX idx_manager_active (portfolio_manager_id, is_active),
                    INDEX idx_assigned_date (assigned_date)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """)
            print("✅ Created portfolio_assignments table")
            return True
        except Exception as e:
            print(f"❌ Error creating table: {e}")
            return False

def fake_migrations():
    """Mark migrations as applied"""
    print("\n🔧 Marking migrations as applied...")
    
    try:
        from django.core.management import call_command
        call_command('migrate', 'users', '0018', '--fake', verbosity=0)
        print("✅ Migration 0018 marked as applied!")
        return True
    except Exception as e:
        print(f"❌ Error marking migration: {e}")
        return False

def main():
    """Main function to fix all client assignment issues"""
    print("🔧 COMPLETE CLIENT ASSIGNMENT DATABASE FIX")
    print("=" * 60)
    
    try:
        # Setup Django
        setup_django()
        print("✅ Django environment setup complete")
        
        # Check if table exists and fix it
        if not fix_portfolio_assignments_table():
            # If table doesn't exist, create it
            if not create_portfolio_assignments_table():
                print("❌ Failed to create portfolio_assignments table")
                sys.exit(1)
        
        # Mark migrations as applied
        if fake_migrations():
            print("\n🎉 ALL CLIENT ASSIGNMENT ISSUES FIXED!")
            print("=" * 50)
            print("✅ Missing columns added")
            print("✅ UUID column types fixed")
            print("✅ ID auto-increment fixed")
            print("✅ Migrations marked as applied")
            print("\nThe client assignment interface should now work perfectly!")
        else:
            print("\n❌ Failed to mark migrations as applied")
            sys.exit(1)
            
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()