#!/usr/bin/env python3
"""
Fix for corrupted users_user_permissions table in production.
This handles cases where the table exists but has missing columns or corrupted structure.
"""

import os
import sys
import django
import logging
from django.db import connection, transaction

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

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def check_table_structure(table_name):
    """Check the structure of a table and return column info"""
    with connection.cursor() as cursor:
        cursor.execute(f"DESCRIBE {table_name}")
        columns = cursor.fetchall()
        return {col[0]: col[1] for col in columns}

def backup_existing_data(table_name):
    """Backup existing data from a table"""
    logger.info(f"Backing up existing data from {table_name}...")
    
    try:
        with connection.cursor() as cursor:
            cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
            count = cursor.fetchone()[0]
            
            if count > 0:
                logger.info(f"Found {count} existing records in {table_name}")
                cursor.execute(f"SELECT * FROM {table_name}")
                data = cursor.fetchall()
                return data
            else:
                logger.info(f"No existing data in {table_name}")
                return []
    except Exception as e:
        logger.warning(f"Could not backup data from {table_name}: {e}")
        return []

def recreate_users_user_permissions_table():
    """Completely recreate the users_user_permissions table"""
    logger.info("Recreating users_user_permissions table...")
    
    try:
        with connection.cursor() as cursor:
            # Check if table exists
            cursor.execute("SHOW TABLES LIKE 'users_user_permissions'")
            table_exists = cursor.fetchone() is not None
            
            if table_exists:
                logger.info("Table exists, backing up data and dropping...")
                # Backup existing data
                existing_data = backup_existing_data('users_user_permissions')
                
                # Drop the corrupted table
                cursor.execute("DROP TABLE users_user_permissions")
                logger.info("Dropped existing corrupted table")
            else:
                logger.info("Table doesn't exist, will create new one")
                existing_data = []
            
            # Create the table with correct structure
            create_sql = """
            CREATE TABLE users_user_permissions (
                id INT AUTO_INCREMENT PRIMARY KEY,
                customuser_id CHAR(36) NOT NULL,
                permission_id INT NOT NULL,
                UNIQUE KEY users_user_permissions_customuser_id_permission_id_uniq (customuser_id, permission_id),
                KEY users_user_permissions_customuser_id (customuser_id),
                KEY users_user_permissions_permission_id (permission_id),
                CONSTRAINT users_user_permissions_customuser_id_fk 
                    FOREIGN KEY (customuser_id) REFERENCES users (id) ON DELETE CASCADE,
                CONSTRAINT users_user_permissions_permission_id_fk 
                    FOREIGN KEY (permission_id) REFERENCES auth_permission (id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """
            
            cursor.execute(create_sql)
            logger.info("✅ Successfully created users_user_permissions table")
            
            # Restore backed up data if any
            if existing_data:
                logger.info(f"Restoring {len(existing_data)} records...")
                for record in existing_data:
                    try:
                        # Skip the id field (auto-increment) and insert the rest
                        cursor.execute("""
                            INSERT INTO users_user_permissions (customuser_id, permission_id) 
                            VALUES (%s, %s)
                        """, [record[1], record[2]])  # Assuming customuser_id is index 1, permission_id is index 2
                    except Exception as e:
                        logger.warning(f"Could not restore record {record}: {e}")
                
                logger.info("✅ Data restoration completed")
            
            return True
            
    except Exception as e:
        logger.error(f"❌ Error recreating users_user_permissions table: {e}")
        return False

def recreate_users_groups_table():
    """Completely recreate the users_groups table"""
    logger.info("Recreating users_groups table...")
    
    try:
        with connection.cursor() as cursor:
            # Check if table exists
            cursor.execute("SHOW TABLES LIKE 'users_groups'")
            table_exists = cursor.fetchone() is not None
            
            if table_exists:
                logger.info("Table exists, backing up data and dropping...")
                # Backup existing data
                existing_data = backup_existing_data('users_groups')
                
                # Drop the existing table
                cursor.execute("DROP TABLE users_groups")
                logger.info("Dropped existing table")
            else:
                logger.info("Table doesn't exist, will create new one")
                existing_data = []
            
            # Create the table with correct structure
            create_sql = """
            CREATE TABLE users_groups (
                id INT AUTO_INCREMENT PRIMARY KEY,
                customuser_id CHAR(36) NOT NULL,
                group_id INT NOT NULL,
                UNIQUE KEY users_groups_customuser_id_group_id_uniq (customuser_id, group_id),
                KEY users_groups_customuser_id (customuser_id),
                KEY users_groups_group_id (group_id),
                CONSTRAINT users_groups_customuser_id_fk 
                    FOREIGN KEY (customuser_id) REFERENCES users (id) ON DELETE CASCADE,
                CONSTRAINT users_groups_group_id_fk 
                    FOREIGN KEY (group_id) REFERENCES auth_group (id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """
            
            cursor.execute(create_sql)
            logger.info("✅ Successfully created users_groups table")
            
            # Restore backed up data if any
            if existing_data:
                logger.info(f"Restoring {len(existing_data)} records...")
                for record in existing_data:
                    try:
                        # Skip the id field (auto-increment) and insert the rest
                        cursor.execute("""
                            INSERT INTO users_groups (customuser_id, group_id) 
                            VALUES (%s, %s)
                        """, [record[1], record[2]])  # Assuming customuser_id is index 1, group_id is index 2
                    except Exception as e:
                        logger.warning(f"Could not restore record {record}: {e}")
                
                logger.info("✅ Data restoration completed")
            
            return True
            
    except Exception as e:
        logger.error(f"❌ Error recreating users_groups table: {e}")
        return False

def test_permission_system():
    """Test that the permission system works correctly"""
    logger.info("Testing permission system...")
    
    try:
        from django.contrib.auth.models import Permission
        from users.models import CustomUser
        
        # Test basic queries
        permissions = Permission.objects.all()[:3]
        users = CustomUser.objects.all()[:3]
        
        if not users.exists() or not permissions.exists():
            logger.warning("No users or permissions found for testing")
            return True
        
        user = users.first()
        permission = permissions.first()
        
        # Test adding permission
        user.user_permissions.add(permission)
        logger.info("✅ Permission add test successful")
        
        # Test removing permission
        user.user_permissions.remove(permission)
        logger.info("✅ Permission remove test successful")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ Permission system test failed: {e}")
        return False

def main():
    """Main function to fix the corrupted permission tables"""
    logger.info("=" * 60)
    logger.info("FIXING CORRUPTED PERMISSION TABLES")
    logger.info("=" * 60)
    
    try:
        # Check current table structures
        logger.info("Checking current table structures...")
        
        try:
            users_perms_structure = check_table_structure('users_user_permissions')
            logger.info(f"users_user_permissions columns: {list(users_perms_structure.keys())}")
        except Exception as e:
            logger.info(f"users_user_permissions table issue: {e}")
        
        try:
            users_groups_structure = check_table_structure('users_groups')
            logger.info(f"users_groups columns: {list(users_groups_structure.keys())}")
        except Exception as e:
            logger.info(f"users_groups table issue: {e}")
        
        # Recreate the tables
        success1 = recreate_users_user_permissions_table()
        success2 = recreate_users_groups_table()
        
        if not success1 or not success2:
            logger.error("❌ Failed to recreate required tables")
            return False
        
        # Test the permission system
        success3 = test_permission_system()
        
        if success3:
            logger.info("🎉 SUCCESS: Permission system is now working correctly!")
            logger.info("The 'Field permission_id doesn't have a default value' error has been resolved.")
            return True
        else:
            logger.error("❌ Permission system test failed")
            return False
            
    except Exception as e:
        logger.error(f"❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)