"""Clean up inconsistent migration history and migration files.

This script:
1. Removes duplicate migration files
2. Updates django_migrations table to match the desired state
3. Fake-applies migrations to sync the state
"""
import os
import sys
import glob
import shutil
import hashlib
import argparse
import django
from pathlib import Path
from collections import defaultdict
from django.core.management import call_command


def init_django():
    """Initialize Django ORM so we can use models."""
    # Add the parent directory to Python path so Django can find the settings
    project_root = str(Path(__file__).resolve().parent.parent)
    if project_root not in sys.path:
        sys.path.append(project_root)
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
    django.setup()


def remove_duplicate_migrations(app_names):
    """Remove duplicate migration files keeping only one copy of each."""
    for app_name in app_names:
        print(f"\nChecking {app_name} migrations...")
        migrations_dir = Path(__file__).resolve().parent.parent / app_name / 'migrations'
        
        if not migrations_dir.exists():
            print(f"No migrations directory found for {app_name}")
            continue
        
        # First backup the migrations directory
        backup_dir = migrations_dir.parent / 'migrations.bak'
        if backup_dir.exists():
            shutil.rmtree(backup_dir)
        shutil.copytree(migrations_dir, backup_dir)
        
        # Find and remove duplicates based on filename
        seen_files = defaultdict(list)
        for f in migrations_dir.glob('*.py'):
            if f.name != '__init__.py':
                seen_files[f.name].append(f)
        
        # Keep only one copy of each migration
        for name, paths in seen_files.items():
            if len(paths) > 1:
                print(f"Found {len(paths)} copies of {name}")
                # Keep the first one, remove the rest
                for f in paths[1:]:
                    print(f"Removing duplicate: {f}")
                    f.unlink()
    
    return True


def clean_migrations(dry_run=False):
    """Remove problem migrations and re-add them in order."""
    from django.db.migrations.recorder import MigrationRecorder
    from django.db import transaction, connection

    # Create recorder
    recorder = MigrationRecorder(None)
    Migration = recorder.Migration

    # Define the order we want (these must match migration file names)
    desired_order = {
        'loans': [
            '0001_initial',
            '0002_add_boost_plus_product',
            '0002_add_penalty_charges',
            '0002_initial',
            '0003_loanproduct_available_repayment_methods_and_more',
            '0004_loanapplication_repayment_method',
            '0005_loanproduct_max_rollover_count_and_more',
            '0006_merge_20250822_2013',
            '0007_merge_20250822_2040',
            '0008_auto_20250823_0823',
            '0009_add_duration_options_and_payment_date',
            '0010_loan_deleted_at_loan_deleted_by_loan_is_deleted',
            '0011_loan_registration_fee_and_more'
        ],
        'utils': [
            # Keep order that matches migration dependencies
            '0001_initial',
            '0002_auditlog',
            '0003_systemsetting_document_documentshare_and_more',
            '0004_alter_document_document_type_and_more',
            '0005_alter_receipt_payment_method',
            '0006_add_related_application_to_notification',
            '0007_alter_notification_related_application',
            '0008_alter_notification_related_application_and_more',
            '0009_remove_notification_app_notification_loan_app',
            '0010_merge_20250827_0157'
        ],
        'reports': [
            '0001_initial',
            '0002_enhanced_reports_models',
            '0003_alter_notification_user',
            '0004_merge_20250824_2312',
            '0005_alter_notification_related_application_and_more',
            '0006_alter_notification_related_application',
            '0007_merge_20250827_0157'
        ],
        'users': [
            '0001_initial',
            '0002_alter_customuser_id_alter_customuser_monthly_income',
            '0003_customuser_logbook_customuser_signature_and_more',
            '0004_customuser_capital_invested_customuser_county_and_more',
            '0005_alter_customuser_role_alter_rolepermission_role',
            '0006_add_portfolio_management',
            '0007_customuser_profile_image',
            '0008_customuser_is_email_verified_and_more'
        ]
    }

    print("Current migrations in database:")
    for app in desired_order.keys():
        print(f"\n{app.upper()} migrations:")
        for m in Migration.objects.filter(app=app).order_by('name'):
            print(f" - {m.app}.{m.name}")

    # Backup current state
    print("\nBacking up current migration state...")
    backup = []
    for m in Migration.objects.all():
        backup.append({
            'app': m.app,
            'name': m.name,
            'applied': m.applied
        })
    
    try:
        with transaction.atomic():
            # Remove all targeted migrations
            print("\nRemoving existing migrations...")
            for app in desired_order.keys():
                print(f"\nRemoving {app} migrations...")
                Migration.objects.filter(app=app).delete()

            # Re-add in correct order
            print("\nRe-adding migrations in correct order:")
            for app, migrations in desired_order.items():
                print(f"\n{app.upper()} migrations:")
                for name in migrations:
                    m = Migration.objects.create(app=app, name=name)
                    print(f" + {m.app}.{m.name}")

            print("\nNew migration state in database:")
            for app in desired_order.keys():
                print(f"\n{app.upper()} migrations:")
                for m in Migration.objects.filter(app=app).order_by('name'):
                    print(f" - {m.app}.{m.name}")

            # Check if this is a dry run
            if dry_run:
                print("\nDRY RUN - Changes look correct but no changes will be made.")
                raise Exception("Dry run complete")

    except Exception as e:
        print(f"\nError or aborted: {e}")
        print("Restoring backup...")
        
        # Restore from backup
        Migration.objects.all().delete()
        for m in backup:
            Migration.objects.create(**m)
        
        print("Backup restored. No changes were made.")
        return False

    print("\nMigration history cleaned successfully.")
    return True


def create_merge_migration():
    """Create a new merge migration for utils app."""
    from django.core.management import call_command
    
    # Create a new merge migration
    try:
        call_command('makemigrations', 'utils', merge=True, name='merge_20250827')
        return True
    except Exception as e:
        print(f"Error creating merge migration: {e}")
        return False


def check_and_add_missing_columns():
    """Check for and add missing database columns that cause OperationalError."""
    from django.db import connection
    
    with connection.cursor() as cursor:
        # Check if registration_fee_amount column exists in loan_applications
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'loan_applications' 
            AND COLUMN_NAME = 'registration_fee_amount'
        """)
        
        if cursor.fetchone()[0] == 0:
            print("Adding missing registration_fee_amount column to loan_applications...")
            cursor.execute("""
                ALTER TABLE loan_applications 
                ADD COLUMN registration_fee_amount DECIMAL(10,2) DEFAULT 0.00
            """)
            print("✓ Added registration_fee_amount column")
        else:
            print("✓ registration_fee_amount column already exists")
            
        # Check if registration_fee column exists in loans
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'loans' 
            AND COLUMN_NAME = 'registration_fee'
        """)
        
        if cursor.fetchone()[0] == 0:
            print("Adding missing registration_fee column to loans...")
            cursor.execute("""
                ALTER TABLE loans 
                ADD COLUMN registration_fee DECIMAL(10,2) DEFAULT 0.00
            """)
            print("✓ Added registration_fee column")
        else:
            print("✓ registration_fee column already exists")


def fake_apply_migrations():
    """Fake apply migrations to sync the state."""
    from django.core.management import call_command
    from django.db import connection
    from django.db.migrations.executor import MigrationExecutor
    from django.db.migrations.recorder import MigrationRecorder

    def clean_stale_migrations():
        """Remove any stale migration entries from the database."""
        stale_migrations = [
            ('utils', '0010_remove_notification_app_notification_loan_app'),
            ('utils', '0011_merge_0004_0010'),
            ('utils', '0012_merge_0004_0011'),
            ('utils', '0013_merge_20250827_0152'),
            ('utils', '0014_merge_20250827_0154'),
            ('utils', '0010_merge_20250827_0157'),
            ('utils', '06_add_related_application_to_notification'),
            ('reports', '0002_initial'),
            ('reports', '0007_merge_20250827_0152'),
            ('reports', '0007_merge_20250827_0157'),
        ]
        
        recorder = MigrationRecorder(connection)
        Migration = recorder.Migration
        
        for app, name in stale_migrations:
            Migration.objects.filter(app=app, name=name).delete()
            print(f"Removed stale migration: {app}.{name}")

    def manually_apply_migrations():
        """Manually insert migration records in correct order."""
        recorder = MigrationRecorder(connection)
        Migration = recorder.Migration
        
        migrations_to_apply = [
            ('utils', '0009_remove_notification_app_notification_loan_app'),
            ('utils', '0010_merge_20250827_0157'),
            ('reports', '0006_alter_notification_related_application'),
            ('reports', '0007_merge_20250827_0157'),
        ]
        
        for app, name in migrations_to_apply:
            Migration.objects.get_or_create(app=app, name=name)
            print(f"Applied migration: {app}.{name}")
    
    # First, check if there are any migrations to apply
    executor = MigrationExecutor(connection)
    plan = executor.migration_plan(executor.loader.graph.leaf_nodes())
    
    if not plan:
        print("No migrations to apply.")
        return True
        
    # Get list of migrations in order they will be applied
    migrations_to_apply = [migration for migration, _ in plan]
    print("\nMigrations that will be applied:")
    for migration in migrations_to_apply:
        print(f" - {migration.app_label}.{migration.name}")
    
    try:
        print("\nApplying migrations in dependency order...")
        
        # First clean up any stale migrations
        clean_stale_migrations()
        
        # First, fake revert all migrations to clean state
        print("\nReverting all migrations to clean state...")
        try:
            for app in ['reports', 'utils', 'loans', 'users']:
                print(f"\nReverting {app} migrations...")
                call_command('migrate', app, 'zero', '--fake')
        except Exception as e:
            print(f"Error reverting migrations: {e}")
            return False
            
        # Then apply migrations in correct dependency order
        print("\nApplying migrations in correct order...")
        try:
            # First apply utils and reports initial migrations
            migrations_order = [
                ('utils', '0001_initial'),
                ('utils', '0002_auditlog'),
                ('utils', '0003_systemsetting_document_documentshare_and_more'),
                ('utils', '0004_alter_document_document_type_and_more'),
                ('utils', '0005_alter_receipt_payment_method'),
                ('reports', '0001_initial'),
                ('reports', '0002_enhanced_reports_models'),
            ]
            
            for app, migration in migrations_order:
                print(f"\nApplying {app}.{migration}...")
                call_command('migrate', app, migration, '--fake')
            
            # Then apply remaining migrations
            for app in ['loans', 'users']:
                print(f"\nApplying remaining {app} migrations...")
                call_command('migrate', app, '--fake')

            # Apply merge migrations manually
            manually_apply_migrations()
            
            return True
        except Exception as e:
            print(f"Error applying migrations: {e}")
            return False
    except Exception as e:
        print(f"Error applying migrations: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description='Clean up Django migration history.')
    parser.add_argument('--dry-run', '-n', action='store_true',
                      help='Show what would be done without making changes')
    args = parser.parse_args()

    # Initialize Django
    init_django()
    
    if args.dry_run:
        print("*** DRY RUN - no changes will be made ***\n")
    
    try:
        # 1. Clean up migration history in database first
        if clean_migrations(dry_run=args.dry_run):
            if not args.dry_run:
                # 2. Remove duplicate migration files
                print("\nRemoving duplicate migration files...")
                remove_duplicate_migrations(['utils', 'loans', 'reports', 'users'])
                
                # 3. Remove merge migrations
                print("\nCleaning up merge migrations...")
                merge_files = [
                    'utils/migrations/0013_merge_20250827_0152.py',
                    'utils/migrations/0014_merge_20250827_0154.py',
                    'reports/migrations/0007_merge_20250827_0152.py'
                ]
                for path in merge_files:
                    try:
                        if os.path.exists(path):
                            os.unlink(path)
                            print(f"Removed {path}")
                    except Exception as e:
                        print(f"Warning: Could not remove {path}: {e}")
                
                # 4. Create fresh merge migrations (skip if problematic)
                print("\nSkipping automatic merge migration creation for production safety...")
                # call_command('makemigrations', '--merge')  # Commented out for production safety
                
                # 5. Add missing database columns
                print("\nChecking and adding missing database columns...")
                check_and_add_missing_columns()
                
                # 6. Fake apply migrations to sync state
                print("\nFake applying migrations...")
                fake_apply_migrations()
                
                print("\nMigration cleanup complete!")
        
    except Exception as e:
        print(f"\nError during cleanup: {e}")
        print("You may need to restore from backup.")
        sys.exit(1)


if __name__ == '__main__':
    main()