#!/usr/bin/env python3
"""
Database Fix Script: User Permissions Table

This script fixes the database schema for the user_permissions table to match
the current model definition. The error indicates that the table structure
doesn't match the expected schema.

Usage:
    python fix_user_permissions_database.py [--dry-run] [--backup]
"""

import os
import sys
import logging
import argparse
import subprocess
from datetime import datetime
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('database_fix.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class DatabaseFixer:
    """Database fixer for user permissions table"""
    
    def __init__(self, dry_run=False, backup=True):
        self.dry_run = dry_run
        self.backup = backup
        self.fix_log = []
    
    def check_database_connection(self):
        """Check if we can connect to the database"""
        try:
            # Try to import Django settings
            import django
            from django.conf import settings
            
            # Configure Django
            os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
            django.setup()
            
            from django.db import connection
            
            # Test connection
            with connection.cursor() as cursor:
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                if result:
                    logger.info("Database connection successful")
                    return True
                    
        except Exception as e:
            logger.error(f"Database connection failed: {str(e)}")
            return False
        
        return False
    
    def check_table_exists(self, table_name):
        """Check if a table exists in the database"""
        try:
            from django.db import connection
            
            with connection.cursor() as cursor:
                cursor.execute(f"SHOW TABLES LIKE '{table_name}'")
                result = cursor.fetchone()
                return result is not None
                
        except Exception as e:
            logger.error(f"Error checking table {table_name}: {str(e)}")
            return False
    
    def get_table_structure(self, table_name):
        """Get the current structure of a table"""
        try:
            from django.db import connection
            
            with connection.cursor() as cursor:
                cursor.execute(f"DESCRIBE {table_name}")
                columns = cursor.fetchall()
                return columns
                
        except Exception as e:
            logger.error(f"Error getting table structure for {table_name}: {str(e)}")
            return []
    
    def backup_table(self, table_name):
        """Create a backup of the table"""
        if not self.backup:
            return True
            
        try:
            from django.db import connection
            
            backup_table_name = f"{table_name}_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            
            with connection.cursor() as cursor:
                # Create backup table
                cursor.execute(f"CREATE TABLE {backup_table_name} AS SELECT * FROM {table_name}")
                logger.info(f"Created backup table: {backup_table_name}")
                self.fix_log.append(f"Backup created: {backup_table_name}")
                return True
                
        except Exception as e:
            logger.error(f"Error creating backup for {table_name}: {str(e)}")
            return False
    
    def fix_user_permissions_table(self):
        """Fix the user_permissions table structure"""
        logger.info("Fixing user_permissions table...")
        
        try:
            from django.db import connection
            
            # Check if table exists
            if not self.check_table_exists('user_permissions'):
                logger.info("user_permissions table does not exist, creating it...")
                return self.create_user_permissions_table()
            
            # Get current table structure
            columns = self.get_table_structure('user_permissions')
            logger.info(f"Current table structure: {columns}")
            
            # Check if module column exists and has correct type
            module_column = None
            for column in columns:
                if column[0] == 'module':
                    module_column = column
                    break
            
            if not module_column:
                logger.info("Module column does not exist, adding it...")
                return self.add_module_column()
            
            # Check if module column has correct length
            if 'varchar(20)' in module_column[1].lower():
                logger.info("Module column has incorrect length (20), updating to 50...")
                return self.update_module_column_length()
            
            logger.info("user_permissions table structure is correct")
            self.fix_log.append("user_permissions table: Structure is correct")
            return True
            
        except Exception as e:
            logger.error(f"Error fixing user_permissions table: {str(e)}")
            return False
    
    def create_user_permissions_table(self):
        """Create the user_permissions table with correct structure"""
        logger.info("Creating user_permissions table...")
        
        if self.dry_run:
            logger.info("[DRY RUN] Would create user_permissions table")
            self.fix_log.append("user_permissions table: Would be created")
            return True
        
        try:
            from django.db import connection
            
            with connection.cursor() as cursor:
                # Create the table with correct structure
                create_sql = """
                CREATE TABLE user_permissions (
                    id BIGINT AUTO_INCREMENT PRIMARY KEY,
                    user_id BIGINT NOT NULL,
                    module VARCHAR(50) NOT NULL,
                    action VARCHAR(30) NOT NULL,
                    is_allowed BOOLEAN NOT NULL DEFAULT FALSE,
                    granted_by_id BIGINT NULL,
                    reason TEXT NULL,
                    expires_at DATETIME NULL,
                    created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    UNIQUE KEY unique_user_permission (user_id, module, action),
                    KEY idx_user_expires (user_id, expires_at),
                    KEY idx_module_action (module, action),
                    CONSTRAINT fk_user_permissions_user 
                        FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
                    CONSTRAINT fk_user_permissions_granted_by 
                        FOREIGN KEY (granted_by_id) REFERENCES users(id) ON DELETE SET NULL
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                """
                
                cursor.execute(create_sql)
                logger.info("Created user_permissions table successfully")
                self.fix_log.append("user_permissions table: Created successfully")
                return True
                
        except Exception as e:
            logger.error(f"Error creating user_permissions table: {str(e)}")
            return False
    
    def add_module_column(self):
        """Add the module column to existing table"""
        logger.info("Adding module column...")
        
        if self.dry_run:
            logger.info("[DRY RUN] Would add module column")
            self.fix_log.append("user_permissions table: Would add module column")
            return True
        
        try:
            from django.db import connection
            
            with connection.cursor() as cursor:
                # Add module column
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN module VARCHAR(50) NOT NULL AFTER user_id")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN action VARCHAR(30) NOT NULL AFTER module")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN is_allowed BOOLEAN NOT NULL DEFAULT FALSE AFTER action")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN granted_by_id BIGINT NULL AFTER is_allowed")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN reason TEXT NULL AFTER granted_by_id")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN expires_at DATETIME NULL AFTER reason")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP AFTER expires_at")
                cursor.execute("ALTER TABLE user_permissions ADD COLUMN updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP AFTER created_at")
                
                # Add indexes and constraints
                cursor.execute("ALTER TABLE user_permissions ADD UNIQUE KEY unique_user_permission (user_id, module, action)")
                cursor.execute("ALTER TABLE user_permissions ADD KEY idx_user_expires (user_id, expires_at)")
                cursor.execute("ALTER TABLE user_permissions ADD KEY idx_module_action (module, action)")
                
                logger.info("Added module column and other required columns successfully")
                self.fix_log.append("user_permissions table: Added required columns")
                return True
                
        except Exception as e:
            logger.error(f"Error adding module column: {str(e)}")
            return False
    
    def update_module_column_length(self):
        """Update module column length from 20 to 50"""
        logger.info("Updating module column length...")
        
        if self.dry_run:
            logger.info("[DRY RUN] Would update module column length")
            self.fix_log.append("user_permissions table: Would update module column length")
            return True
        
        try:
            from django.db import connection
            
            with connection.cursor() as cursor:
                # Update module column length
                cursor.execute("ALTER TABLE user_permissions MODIFY COLUMN module VARCHAR(50) NOT NULL")
                logger.info("Updated module column length successfully")
                self.fix_log.append("user_permissions table: Updated module column length")
                return True
                
        except Exception as e:
            logger.error(f"Error updating module column length: {str(e)}")
            return False
    
    def run_django_migrations(self):
        """Run Django migrations to ensure database is up to date"""
        logger.info("Running Django migrations...")
        
        if self.dry_run:
            logger.info("[DRY RUN] Would run Django migrations")
            self.fix_log.append("Django migrations: Would be run")
            return True
        
        try:
            # Run migrations
            result = subprocess.run([
                sys.executable, 'manage.py', 'migrate', 'users'
            ], capture_output=True, text=True, cwd=Path.cwd())
            
            if result.returncode == 0:
                logger.info("Django migrations completed successfully")
                self.fix_log.append("Django migrations: Completed successfully")
                return True
            else:
                logger.error(f"Django migrations failed: {result.stderr}")
                self.fix_log.append("Django migrations: Failed")
                return False
                
        except Exception as e:
            logger.error(f"Error running Django migrations: {str(e)}")
            return False
    
    def fix(self):
        """Main fix method"""
        logger.info("=" * 60)
        logger.info("DATABASE FIX: USER PERMISSIONS TABLE")
        logger.info("=" * 60)
        logger.info(f"Dry run: {self.dry_run}")
        logger.info(f"Backup: {self.backup}")
        logger.info("=" * 60)
        
        try:
            # Check database connection
            if not self.check_database_connection():
                logger.error("Cannot connect to database")
                return False
            
            # Fix user_permissions table
            if not self.fix_user_permissions_table():
                logger.error("Failed to fix user_permissions table")
                return False
            
            # Run Django migrations
            if not self.run_django_migrations():
                logger.error("Failed to run Django migrations")
                return False
            
            # Log summary
            logger.info("=" * 60)
            logger.info("FIX SUMMARY")
            logger.info("=" * 60)
            for log_entry in self.fix_log:
                logger.info(f"  {log_entry}")
            logger.info("=" * 60)
            
            if self.dry_run:
                logger.info("DRY RUN COMPLETED - No changes were made")
            else:
                logger.info("DATABASE FIX COMPLETED SUCCESSFULLY!")
            
            return True
            
        except Exception as e:
            logger.error(f"Database fix failed: {str(e)}")
            return False


def main():
    """Main function with argument parsing"""
    parser = argparse.ArgumentParser(
        description='Fix user_permissions database table',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  python fix_user_permissions_database.py                    # Fix with backup
  python fix_user_permissions_database.py --dry-run          # Test fix
  python fix_user_permissions_database.py --no-backup        # Fix without backup
        '''
    )
    
    parser.add_argument('--dry-run', action='store_true',
                       help='Show what would be done without making changes')
    parser.add_argument('--no-backup', action='store_true',
                       help='Skip creating backup (not recommended)')
    
    args = parser.parse_args()
    
    try:
        fixer = DatabaseFixer(
            dry_run=args.dry_run,
            backup=not args.no_backup
        )
        
        success = fixer.fix()
        sys.exit(0 if success else 1)
        
    except Exception as e:
        logger.error(f"Script failed: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()