#!/usr/bin/env python3
"""
Comprehensive Production Database Fix Script
Addresses multiple production database issues:
1. Missing penalty_charges table
2. Missing loan_app_id column in utils_notification table
3. Missing registration_fee_amount column in loan_applications table
4. Missing registration_fee column in loans table
"""

import os
import sys
import django
import pymysql
from decimal import Decimal
import logging

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

from django.db import connection, transaction
from django.core.management import execute_from_command_line

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

def check_table_exists(table_name):
    """Check if a table exists in the database"""
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.tables 
            WHERE table_schema = DATABASE() 
            AND table_name = %s
        """, [table_name])
        return cursor.fetchone()[0] > 0

def check_column_exists(table_name, column_name):
    """Check if a column exists in a table"""
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.columns 
            WHERE table_schema = DATABASE() 
            AND table_name = %s 
            AND column_name = %s
        """, [table_name, column_name])
        return cursor.fetchone()[0] > 0

def create_penalty_charges_table():
    """Create the penalty_charges table if it doesn't exist"""
    logger.info("Creating penalty_charges table...")
    
    # Create table without foreign key constraint first
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `penalty_charges` (
        `id` char(32) NOT NULL,
        `amount` decimal(12,2) NOT NULL,
        `penalty_rate` decimal(5,2) NOT NULL COMMENT 'Daily penalty rate used',
        `days_overdue` int unsigned NOT NULL COMMENT 'Number of days overdue when penalty was applied',
        `outstanding_amount` decimal(12,2) NOT NULL COMMENT 'Outstanding amount when penalty was calculated',
        `applied_date` datetime(6) NOT NULL,
        `created_at` datetime(6) NOT NULL,
        `loan_id` char(32) NOT NULL,
        PRIMARY KEY (`id`),
        KEY `penalty_charges_loan_id_idx` (`loan_id`),
        KEY `penalty_charges_applied_date_idx` (`applied_date`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("penalty_charges table created successfully")
        
        # Try to add foreign key constraint separately
        try:
            cursor.execute("""
                ALTER TABLE `penalty_charges` 
                ADD CONSTRAINT `penalty_charges_loan_id_fk` 
                FOREIGN KEY (`loan_id`) REFERENCES `loans` (`id`) ON DELETE CASCADE
            """)
            logger.info("Foreign key constraint added successfully")
        except Exception as fk_error:
            logger.warning(f"Could not add foreign key constraint: {fk_error}")
            logger.info("Table created without foreign key constraint - this is acceptable for functionality")

def create_customer_requests_table():
    """Create the customer_requests table if it doesn't exist"""
    logger.info("Creating customer_requests table...")
    
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `customer_requests` (
        `id` char(32) NOT NULL,
        `request_number` varchar(20) NOT NULL UNIQUE,
        `customer_id` char(32) NOT NULL,
        `request_type` varchar(30) NOT NULL,
        `subject` varchar(200) NOT NULL,
        `description` longtext NOT NULL,
        `priority` varchar(20) NOT NULL DEFAULT 'medium',
        `status` varchar(20) NOT NULL DEFAULT 'pending',
        `assigned_to_id` char(32) NULL,
        `related_loan_id` char(32) NULL,
        `related_application_id` char(32) NULL,
        `resolution_notes` longtext NULL,
        `resolved_by_id` char(32) NULL,
        `resolved_at` datetime(6) NULL,
        `created_at` datetime(6) NOT NULL,
        `updated_at` datetime(6) NOT NULL,
        PRIMARY KEY (`id`),
        KEY `customer_requests_customer_id_idx` (`customer_id`),
        KEY `customer_requests_status_idx` (`status`),
        KEY `customer_requests_created_at_idx` (`created_at`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("customer_requests table created successfully")

def create_registration_fees_table():
    """Create the registration_fees table if it doesn't exist"""
    logger.info("Creating registration_fees table...")
    
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `registration_fees` (
        `id` char(32) NOT NULL,
        `product_type` varchar(30) NOT NULL,
        `fee_name` varchar(200) NOT NULL,
        `amount` decimal(10,2) NOT NULL,
        `description` longtext NULL,
        `is_active` tinyint(1) NOT NULL DEFAULT 1,
        `effective_from` datetime(6) NOT NULL,
        `effective_to` datetime(6) NULL,
        `created_at` datetime(6) NOT NULL,
        `updated_at` datetime(6) NOT NULL,
        `created_by_id` char(32) NULL,
        PRIMARY KEY (`id`),
        KEY `registration_fees_product_type_idx` (`product_type`),
        KEY `registration_fees_is_active_idx` (`is_active`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("registration_fees table created successfully")

def create_registration_fee_payments_table():
    """Create the registration_fee_payments table if it doesn't exist"""
    logger.info("Creating registration_fee_payments table...")
    
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `registration_fee_payments` (
        `id` char(32) NOT NULL,
        `receipt_number` varchar(20) NOT NULL UNIQUE,
        `customer_id` char(32) NOT NULL,
        `registration_fee_id` char(32) NOT NULL,
        `amount_paid` decimal(10,2) NOT NULL,
        `payment_method` varchar(20) NOT NULL,
        `payment_date` datetime(6) NOT NULL,
        `transaction_reference` varchar(100) NULL,
        `payment_notes` longtext NULL,
        `related_loan_id` char(32) NULL,
        `related_application_id` char(32) NULL,
        `processed_by_id` char(32) NULL,
        `created_at` datetime(6) NOT NULL,
        PRIMARY KEY (`id`),
        KEY `registration_fee_payments_customer_id_idx` (`customer_id`),
        KEY `registration_fee_payments_payment_date_idx` (`payment_date`),
        KEY `registration_fee_payments_registration_fee_id_idx` (`registration_fee_id`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("registration_fee_payments table created successfully")

def add_repayment_id_column():
    """Add repayment_id column to receipts table if missing"""
    logger.info("Adding repayment_id column to receipts table...")
    
    with connection.cursor() as cursor:
        # Add column first
        cursor.execute("""
            ALTER TABLE `receipts` 
            ADD COLUMN `repayment_id` char(32) NULL
        """)
        logger.info("repayment_id column added successfully")
        
        # Add index
        cursor.execute("""
            ALTER TABLE `receipts` 
            ADD KEY `receipts_repayment_id_idx` (`repayment_id`)
        """)
        logger.info("Index added successfully")
        
        # Try to add foreign key constraint separately
        try:
            cursor.execute("""
                ALTER TABLE `receipts` 
                ADD CONSTRAINT `receipts_repayment_id_fk` 
                FOREIGN KEY (`repayment_id`) REFERENCES `repayments` (`id`) ON DELETE CASCADE
            """)
            logger.info("Foreign key constraint added successfully")
        except Exception as fk_error:
            logger.warning(f"Could not add foreign key constraint: {fk_error}")
            logger.info("Column created without foreign key constraint - this is acceptable for functionality")

def create_report_schedules_table():
    """Create the report_schedules table if it doesn't exist"""
    logger.info("Creating report_schedules table...")
    
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `report_schedules` (
        `id` char(32) NOT NULL,
        `name` varchar(200) NOT NULL,
        `report_type` varchar(30) NOT NULL,
        `frequency` varchar(20) NOT NULL,
        `is_active` tinyint(1) NOT NULL DEFAULT 1,
        `next_run` datetime(6) NOT NULL,
        `last_run` datetime(6) NULL,
        `parameters` json NOT NULL,
        `created_at` datetime(6) NOT NULL,
        `created_by_id` char(32) NOT NULL,
        PRIMARY KEY (`id`),
        KEY `report_schedules_report_type_idx` (`report_type`),
        KEY `report_schedules_is_active_idx` (`is_active`),
        KEY `report_schedules_next_run_idx` (`next_run`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("report_schedules table created successfully")

def create_report_schedules_recipients_table():
    """Create the report_schedules_recipients many-to-many table"""
    logger.info("Creating report_schedules_recipients table...")
    
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS `report_schedules_recipients` (
        `id` bigint NOT NULL AUTO_INCREMENT,
        `reportschedule_id` char(32) NOT NULL,
        `user_id` char(32) NOT NULL,
        PRIMARY KEY (`id`),
        UNIQUE KEY `report_schedules_recipients_unique` (`reportschedule_id`, `user_id`),
        KEY `report_schedules_recipients_reportschedule_id_idx` (`reportschedule_id`),
        KEY `report_schedules_recipients_user_id_idx` (`user_id`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    
    with connection.cursor() as cursor:
        cursor.execute(create_table_sql)
        logger.info("report_schedules_recipients table created successfully")

def add_loan_app_id_column():
    """Add loan_app_id column to utils_notification table"""
    logger.info("Adding loan_app_id column to utils_notification table...")
    
    with connection.cursor() as cursor:
        # Add column first
        cursor.execute("""
            ALTER TABLE `utils_notification` 
            ADD COLUMN `loan_app_id` char(32) NULL
        """)
        logger.info("loan_app_id column added successfully")
        
        # Add index
        cursor.execute("""
            ALTER TABLE `utils_notification` 
            ADD KEY `utils_notification_loan_app_id_idx` (`loan_app_id`)
        """)
        logger.info("Index added successfully")
        
        # Try to add foreign key constraint separately
        try:
            cursor.execute("""
                ALTER TABLE `utils_notification` 
                ADD CONSTRAINT `utils_notification_loan_app_id_fk` 
                FOREIGN KEY (`loan_app_id`) REFERENCES `loan_applications` (`id`) ON DELETE CASCADE
            """)
            logger.info("Foreign key constraint added successfully")
        except Exception as fk_error:
            logger.warning(f"Could not add foreign key constraint: {fk_error}")
            logger.info("Column created without foreign key constraint - this is acceptable for functionality")

def add_related_loan_id_column():
    """Add related_loan_id column to utils_notification table if missing"""
    logger.info("Adding related_loan_id column to utils_notification table...")
    
    with connection.cursor() as cursor:
        # Add column first
        cursor.execute("""
            ALTER TABLE `utils_notification` 
            ADD COLUMN `related_loan_id` char(32) NULL
        """)
        logger.info("related_loan_id column added successfully")
        
        # Add index
        cursor.execute("""
            ALTER TABLE `utils_notification` 
            ADD KEY `utils_notification_related_loan_id_idx` (`related_loan_id`)
        """)
        logger.info("Index added successfully")
        
        # Try to add foreign key constraint separately
        try:
            cursor.execute("""
                ALTER TABLE `utils_notification` 
                ADD CONSTRAINT `utils_notification_related_loan_id_fk` 
                FOREIGN KEY (`related_loan_id`) REFERENCES `loans` (`id`) ON DELETE CASCADE
            """)
            logger.info("Foreign key constraint added successfully")
        except Exception as fk_error:
            logger.warning(f"Could not add foreign key constraint: {fk_error}")
            logger.info("Column created without foreign key constraint - this is acceptable for functionality")

def add_registration_fee_columns():
    """Add missing registration fee columns"""
    logger.info("Adding registration fee columns...")
    
    # Add registration_fee_amount to loan_applications table
    if not check_column_exists('loan_applications', 'registration_fee_amount'):
        with connection.cursor() as cursor:
            cursor.execute("""
                ALTER TABLE `loan_applications` 
                ADD COLUMN `registration_fee_amount` decimal(12,2) NULL DEFAULT 0.00
            """)
        logger.info("Added registration_fee_amount column to loan_applications")
    
    # Add registration_fee to loans table
    if not check_column_exists('loans', 'registration_fee'):
        with connection.cursor() as cursor:
            cursor.execute("""
                ALTER TABLE `loans` 
                ADD COLUMN `registration_fee` decimal(12,2) NOT NULL DEFAULT 0.00
            """)
        logger.info("Added registration_fee column to loans")

def update_existing_data():
    """Update existing data to set default values for new columns"""
    logger.info("Updating existing data with default values...")
    
    with connection.cursor() as cursor:
        # Update loan_applications registration_fee_amount
        cursor.execute("""
            UPDATE loan_applications 
            SET registration_fee_amount = 0.00 
            WHERE registration_fee_amount IS NULL
        """)
        
        # Update loans registration_fee
        cursor.execute("""
            UPDATE loans 
            SET registration_fee = 0.00 
            WHERE registration_fee IS NULL OR registration_fee = 0
        """)
    
    logger.info("Existing data updated successfully")

def verify_fixes():
    """Verify that all fixes have been applied correctly"""
    logger.info("Verifying database fixes...")
    
    issues = []
    
    # Check penalty_charges table
    if not check_table_exists('penalty_charges'):
        issues.append("penalty_charges table still missing")
    else:
        logger.info("✓ penalty_charges table exists")
        # Check for created_at column
        if not check_column_exists('penalty_charges', 'created_at'):
            issues.append("created_at column missing from penalty_charges")
        else:
            logger.info("✓ penalty_charges.created_at column exists")
    
    # Check customer_requests table
    if not check_table_exists('customer_requests'):
        issues.append("customer_requests table still missing")
    else:
        logger.info("✓ customer_requests table exists")
    
    # Check registration_fees table
    if not check_table_exists('registration_fees'):
        issues.append("registration_fees table still missing")
    else:
        logger.info("✓ registration_fees table exists")
    
    # Check registration_fee_payments table
    if not check_table_exists('registration_fee_payments'):
        issues.append("registration_fee_payments table still missing")
    else:
        logger.info("✓ registration_fee_payments table exists")
    
    # Check report_schedules table
    if not check_table_exists('report_schedules'):
        issues.append("report_schedules table still missing")
    else:
        logger.info("✓ report_schedules table exists")
    
    # Check report_schedules_recipients table
    if not check_table_exists('report_schedules_recipients'):
        issues.append("report_schedules_recipients table still missing")
    else:
        logger.info("✓ report_schedules_recipients table exists")
    
    # Check repayment_id column in receipts
    if not check_column_exists('receipts', 'repayment_id'):
        issues.append("repayment_id column still missing from receipts")
    else:
        logger.info("✓ repayment_id column exists in receipts")
    
    # Check related_loan_id column in notifications
    if not check_column_exists('utils_notification', 'related_loan_id'):
        issues.append("related_loan_id column still missing from utils_notification")
    else:
        logger.info("✓ related_loan_id column exists in utils_notification")
    
    # Check loan_app_id column
    if not check_column_exists('utils_notification', 'loan_app_id'):
        issues.append("loan_app_id column still missing from utils_notification")
    else:
        logger.info("✓ loan_app_id column exists in utils_notification")
    
    # Check registration fee columns
    if not check_column_exists('loan_applications', 'registration_fee_amount'):
        issues.append("registration_fee_amount column missing from loan_applications")
    else:
        logger.info("✓ registration_fee_amount column exists in loan_applications")
    
    if not check_column_exists('loans', 'registration_fee'):
        issues.append("registration_fee column missing from loans")
    else:
        logger.info("✓ registration_fee column exists in loans")
    
    return issues

def main():
    """Main execution function"""
    logger.info("Starting comprehensive production database fix...")
    
    try:
        with transaction.atomic():
            # 1. Create penalty_charges table if missing
            if not check_table_exists('penalty_charges'):
                create_penalty_charges_table()
            else:
                logger.info("penalty_charges table already exists")
                # Check if created_at column exists and add if missing
                if not check_column_exists('penalty_charges', 'created_at'):
                    logger.info("Adding missing created_at column to penalty_charges table")
                    with connection.cursor() as cursor:
                        cursor.execute("ALTER TABLE penalty_charges ADD COLUMN created_at datetime(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6)")
                    logger.info("created_at column added to penalty_charges table")
            
            # 2. Create customer_requests table if missing
            if not check_table_exists('customer_requests'):
                create_customer_requests_table()
            else:
                logger.info("customer_requests table already exists")
            
            # 3. Create registration_fees table if missing
            if not check_table_exists('registration_fees'):
                create_registration_fees_table()
            else:
                logger.info("registration_fees table already exists")
            
            # 4. Create registration_fee_payments table if missing
            if not check_table_exists('registration_fee_payments'):
                create_registration_fee_payments_table()
            else:
                logger.info("registration_fee_payments table already exists")
            
            # 5. Create report_schedules table if missing
            if not check_table_exists('report_schedules'):
                create_report_schedules_table()
            else:
                logger.info("report_schedules table already exists")
            
            # 6. Create report_schedules_recipients table if missing
            if not check_table_exists('report_schedules_recipients'):
                create_report_schedules_recipients_table()
            else:
                logger.info("report_schedules_recipients table already exists")
            
            # 7. Add repayment_id column to receipts if missing
            if not check_column_exists('receipts', 'repayment_id'):
                add_repayment_id_column()
            else:
                logger.info("repayment_id column already exists")
            
            # 8. Add related_loan_id column to notifications if missing
            if not check_column_exists('utils_notification', 'related_loan_id'):
                add_related_loan_id_column()
            else:
                logger.info("related_loan_id column already exists")
            
            # 9. Add loan_app_id column if missing
            if not check_column_exists('utils_notification', 'loan_app_id'):
                add_loan_app_id_column()
            else:
                logger.info("loan_app_id column already exists")
            
            # 10. Add registration fee columns if missing
            add_registration_fee_columns()
            
            # 10. Update existing data
            update_existing_data()
            
            # 11. Verify all fixes
            issues = verify_fixes()
            
            if issues:
                logger.error("Some issues remain:")
                for issue in issues:
                    logger.error(f"  - {issue}")
                return False
            else:
                logger.info("All database fixes applied successfully!")
                return True
                
    except Exception as e:
        logger.error(f"Error during database fix: {str(e)}")
        logger.error("Rolling back changes...")
        raise

if __name__ == "__main__":
    try:
        success = main()
        if success:
            logger.info("Production database fix completed successfully!")
            sys.exit(0)
        else:
            logger.error("Production database fix failed!")
            sys.exit(1)
    except Exception as e:
        logger.error(f"Fatal error: {str(e)}")
        sys.exit(1)