#!/usr/bin/env python3
"""
Comprehensive Schema Migration Fix Script

This script adds all missing columns and tables identified in the verification.
It ensures the database schema matches what the application code expects.

Missing items to fix:
1. Tables: loans_repayment, receipts, loan_statements, offer_letters
2. Columns in loans: loan_officer_id, approved_by_id, disbursed_by_id, created_by_id, updated_by_id, outstanding_balance, registration_fee
3. Columns in loan_applications: loan_officer_id, approved_by_id, disbursed_by_id
"""

import os
import sys
import django
from datetime import datetime

# Setup Django
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
django.setup()

from django.db import connection, transaction


def check_table_exists(table_name):
    """Check if a table exists"""
    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"""
    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 add_column_if_missing(cursor, table_name, column_name, column_definition):
    """Add a column if it doesn't exist"""
    if not check_column_exists(table_name, column_name):
        print(f"  Adding column {table_name}.{column_name}...")
        cursor.execute(f"""
            ALTER TABLE `{table_name}`
            ADD COLUMN `{column_name}` {column_definition}
        """)
        print(f"  ✅ Added {table_name}.{column_name}")
        return True
    else:
        print(f"  ⏭️  Column {table_name}.{column_name} already exists")
        return False


def create_missing_tables():
    """Create all missing tables"""
    print("\n" + "="*80)
    print("CREATING MISSING TABLES")
    print("="*80)
    
    with connection.cursor() as cursor:
        # 1. Create loans_repayment table (if missing)
        if not check_table_exists('loans_repayment'):
            print("\nCreating loans_repayment table...")
            cursor.execute("""
                CREATE TABLE `loans_repayment` (
                    `id` bigint NOT NULL AUTO_INCREMENT,
                    `loan_id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `amount` decimal(10,2) NOT NULL,
                    `payment_date` datetime(6) NOT NULL,
                    `payment_method` varchar(50) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `transaction_reference` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `notes` longtext COLLATE utf8mb4_unicode_ci,
                    `recorded_by_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `created_at` datetime(6) NOT NULL,
                    `mpesa_transaction_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `payment_source` varchar(20) COLLATE utf8mb4_unicode_ci NOT NULL DEFAULT 'manual',
                    PRIMARY KEY (`id`),
                    KEY `loans_repayment_loan_id_idx` (`loan_id`),
                    KEY `loans_repayment_recorded_by_id_idx` (`recorded_by_id`),
                    KEY `loans_repayment_mpesa_transaction_id_idx` (`mpesa_transaction_id`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)
            print("  ✅ Created loans_repayment table")
        else:
            print("  ⏭️  loans_repayment table already exists")
        
        # 2. Create receipts table (if missing)
        if not check_table_exists('receipts'):
            print("\nCreating receipts table...")
            cursor.execute("""
                CREATE TABLE `receipts` (
                    `id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `receipt_number` varchar(20) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `repayment_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `loan_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `borrower_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `amount_paid` decimal(12,2) DEFAULT NULL,
                    `payment_method` varchar(20) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `payment_date` datetime(6) DEFAULT NULL,
                    `transaction_reference` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `receipt_date` datetime(6) NOT NULL,
                    `pdf_file` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `created_at` datetime(6) NOT NULL,
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `receipt_number` (`receipt_number`),
                    KEY `receipts_repayment_id_idx` (`repayment_id`),
                    KEY `receipts_loan_id_idx` (`loan_id`),
                    KEY `receipts_borrower_id_idx` (`borrower_id`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)
            print("  ✅ Created receipts table")
        else:
            print("  ⏭️  receipts table already exists")
        
        # 3. Create loan_statements table (if missing)
        if not check_table_exists('loan_statements'):
            print("\nCreating loan_statements table...")
            cursor.execute("""
                CREATE TABLE `loan_statements` (
                    `id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `statement_number` varchar(20) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `loan_id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `borrower_id` char(32) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `statement_date` datetime(6) NOT NULL,
                    `period_from` datetime(6) NOT NULL,
                    `period_to` datetime(6) NOT NULL,
                    `principal_amount` decimal(12,2) NOT NULL,
                    `interest_amount` decimal(12,2) NOT NULL,
                    `total_amount` decimal(12,2) NOT NULL,
                    `amount_paid` decimal(12,2) NOT NULL,
                    `outstanding_balance` decimal(12,2) NOT NULL,
                    `pdf_file` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `created_at` datetime(6) NOT NULL,
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `statement_number` (`statement_number`),
                    KEY `loan_statements_loan_id_idx` (`loan_id`),
                    KEY `loan_statements_borrower_id_idx` (`borrower_id`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)
            print("  ✅ Created loan_statements table")
        else:
            print("  ⏭️  loan_statements table already exists")
        
        # 4. Create offer_letters table (if missing)
        if not check_table_exists('offer_letters'):
            print("\nCreating offer_letters table...")
            cursor.execute("""
                CREATE TABLE `offer_letters` (
                    `id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `offer_number` varchar(20) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `application_id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `borrower_id` char(32) COLLATE utf8mb4_unicode_ci NOT NULL,
                    `loan_amount` decimal(12,2) NOT NULL,
                    `interest_rate` decimal(5,2) NOT NULL,
                    `duration_days` int NOT NULL,
                    `total_repayment` decimal(12,2) NOT NULL,
                    `status` varchar(20) COLLATE utf8mb4_unicode_ci NOT NULL DEFAULT 'pending',
                    `generated_at` datetime(6) NOT NULL,
                    `expires_at` datetime(6) DEFAULT NULL,
                    `accepted_at` datetime(6) DEFAULT NULL,
                    `pdf_file` varchar(100) COLLATE utf8mb4_unicode_ci DEFAULT NULL,
                    `borrower_signature` longtext COLLATE utf8mb4_unicode_ci,
                    `signed_at` datetime(6) DEFAULT NULL,
                    `created_at` datetime(6) NOT NULL,
                    PRIMARY KEY (`id`),
                    UNIQUE KEY `offer_number` (`offer_number`),
                    KEY `offer_letters_application_id_idx` (`application_id`),
                    KEY `offer_letters_borrower_id_idx` (`borrower_id`)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)
            print("  ✅ Created offer_letters table")
        else:
            print("  ⏭️  offer_letters table already exists")


def add_missing_loan_columns():
    """Add missing columns to loans table"""
    print("\n" + "="*80)
    print("ADDING MISSING COLUMNS TO LOANS TABLE")
    print("="*80)
    
    if not check_table_exists('loans'):
        print("❌ Loans table does not exist!")
        return
    
    with connection.cursor() as cursor:
        columns_to_add = [
            ('loan_officer_id', 'CHAR(32) NULL'),
            ('approved_by_id', 'CHAR(32) NULL'),
            ('disbursed_by_id', 'CHAR(32) NULL'),
            ('created_by_id', 'CHAR(32) NULL'),
            ('updated_by_id', 'CHAR(32) NULL'),
            ('outstanding_balance', 'DECIMAL(10,2) DEFAULT NULL'),
            ('registration_fee', 'DECIMAL(12,2) NOT NULL DEFAULT 0.00'),
        ]
        
        for column_name, column_def in columns_to_add:
            add_column_if_missing(cursor, 'loans', column_name, column_def)


def add_missing_loan_application_columns():
    """Add missing columns to loan_applications table"""
    print("\n" + "="*80)
    print("ADDING MISSING COLUMNS TO LOAN_APPLICATIONS TABLE")
    print("="*80)
    
    if not check_table_exists('loan_applications'):
        print("❌ Loan_applications table does not exist!")
        return
    
    with connection.cursor() as cursor:
        columns_to_add = [
            ('loan_officer_id', 'CHAR(32) NULL'),
            ('approved_by_id', 'CHAR(32) NULL'),
            ('disbursed_by_id', 'CHAR(32) NULL'),
        ]
        
        for column_name, column_def in columns_to_add:
            add_column_if_missing(cursor, 'loan_applications', column_name, column_def)


def update_outstanding_balance():
    """Update outstanding_balance for existing loans"""
    print("\n" + "="*80)
    print("UPDATING OUTSTANDING BALANCE FOR EXISTING LOANS")
    print("="*80)
    
    if not check_column_exists('loans', 'outstanding_balance'):
        print("⏭️  outstanding_balance column doesn't exist yet, skipping update")
        return
    
    with connection.cursor() as cursor:
        print("Calculating outstanding balance for all loans...")
        cursor.execute("""
            UPDATE loans
            SET outstanding_balance = total_amount - COALESCE(amount_paid, 0)
            WHERE outstanding_balance IS NULL
        """)
        affected = cursor.rowcount
        print(f"✅ Updated outstanding_balance for {affected} loans")


def create_indexes():
    """Create indexes for foreign key columns"""
    print("\n" + "="*80)
    print("CREATING INDEXES FOR FOREIGN KEY COLUMNS")
    print("="*80)
    
    indexes = [
        ('loans', 'idx_loans_borrower_id', 'borrower_id'),
        ('loans', 'idx_loans_application_id', 'application_id'),
        ('loans', 'idx_loans_loan_officer_id', 'loan_officer_id'),
        ('loans', 'idx_loans_approved_by_id', 'approved_by_id'),
        ('loans', 'idx_loans_disbursed_by_id', 'disbursed_by_id'),
        ('loan_applications', 'idx_loan_apps_borrower_id', 'borrower_id'),
        ('loan_applications', 'idx_loan_apps_loan_product_id', 'loan_product_id'),
        ('loan_applications', 'idx_loan_apps_reviewed_by_id', 'reviewed_by_id'),
        ('loan_applications', 'idx_loan_apps_loan_officer_id', 'loan_officer_id'),
        ('loan_applications', 'idx_loan_apps_approved_by_id', 'approved_by_id'),
    ]
    
    with connection.cursor() as cursor:
        for table, index_name, column_name in indexes:
            if not check_table_exists(table):
                print(f"  ⏭️  Table {table} doesn't exist, skipping index {index_name}")
                continue
            
            if not check_column_exists(table, column_name):
                print(f"  ⏭️  Column {table}.{column_name} doesn't exist, skipping index {index_name}")
                continue
            
            # Check if index already exists
            cursor.execute("""
                SELECT COUNT(*)
                FROM information_schema.statistics
                WHERE table_schema = DATABASE()
                AND table_name = %s
                AND index_name = %s
            """, [table, index_name])
            
            if cursor.fetchone()[0] == 0:
                print(f"  Creating index {index_name} on {table}.{column_name}...")
                try:
                    cursor.execute(f"""
                        CREATE INDEX `{index_name}` ON `{table}` (`{column_name}`)
                    """)
                    print(f"  ✅ Created index {index_name}")
                except Exception as e:
                    print(f"  ⚠️  Could not create index {index_name}: {e}")
            else:
                print(f"  ⏭️  Index {index_name} already exists")


def main():
    """Run all migration steps"""
    print("="*80)
    print("COMPREHENSIVE SCHEMA MIGRATION FIX")
    print("="*80)
    print(f"Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    try:
        with transaction.atomic():
            # Step 1: Create missing tables
            create_missing_tables()
            
            # Step 2: Add missing columns to loans table
            add_missing_loan_columns()
            
            # Step 3: Add missing columns to loan_applications table
            add_missing_loan_application_columns()
            
            # Step 4: Update outstanding balance
            update_outstanding_balance()
            
            # Step 5: Create indexes
            create_indexes()
        
        print("\n" + "="*80)
        print("✅ MIGRATION COMPLETED SUCCESSFULLY")
        print("="*80)
        print(f"Finished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("\nNext steps:")
        print("1. Run: python verify_schema_migration.py")
        print("2. Check that all issues are resolved")
        print("3. Test the application thoroughly")
        
        return 0
        
    except Exception as e:
        print("\n" + "="*80)
        print("❌ MIGRATION FAILED")
        print("="*80)
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == '__main__':
    sys.exit(main())