#!/usr/bin/env python3
"""
Database Migration Script

This script migrates the existing SQLite database to the new schema that includes
the students table and relationships.

Usage:
    python migrate_database.py [--database <db_file>] [--dry-run]

Examples:
    python migrate_database.py
    python migrate_database.py --database tablets.db
    python migrate_database.py --dry-run
"""

import argparse
import os
import sqlite3
import sys
from datetime import datetime


DEFAULT_DATABASE = 'tablets.db'


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Migrate database to new schema with students table'
    )
    parser.add_argument(
        '--database', '-d',
        default=DEFAULT_DATABASE,
        help=f'Path to SQLite database file (default: {DEFAULT_DATABASE})'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Show migration steps without executing'
    )
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Show detailed migration information'
    )
    return parser.parse_args()


def backup_database(database_path: str) -> str:
    """
    Create a backup of the database before migration.
    
    Args:
        database_path: Path to the database file
        
    Returns:
        Path to the backup file
    """
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_path = f"{database_path}.backup_{timestamp}"
    
    if os.path.exists(database_path):
        import shutil
        shutil.copy2(database_path, backup_path)
        print(f"Created backup: {backup_path}")
    
    return backup_path


def check_table_exists(conn: sqlite3.Connection, table_name: str) -> bool:
    """Check if a table exists in the database."""
    cursor = conn.cursor()
    cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
        (table_name,)
    )
    return cursor.fetchone() is not None


def check_column_exists(conn: sqlite3.Connection, table_name: str, column_name: str) -> bool:
    """Check if a column exists in a table."""
    cursor = conn.cursor()
    cursor.execute(f"PRAGMA table_info({table_name})")
    columns = [col[1] for col in cursor.fetchall()]
    return column_name in columns


def create_students_table(conn: sqlite3.Connection) -> bool:
    """Create the students table."""
    try:
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS students (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                order_number INTEGER,
                cial_code TEXT UNIQUE NOT NULL,
                nif_nie_passport TEXT UNIQUE,
                registration_number INTEGER,
                file_number TEXT,
                first_name TEXT NOT NULL,
                last_name TEXT NOT NULL,
                full_name TEXT NOT NULL,
                birth_date TEXT,
                gender TEXT CHECK(gender IN ('M', 'F', 'O')),
                study_group TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP,
                CONSTRAINT unique_identification UNIQUE (cial_code, nif_nie_passport)
            )
        ''')
        
        # Create indexes
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_cial ON students(cial_code)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_nif ON students(nif_nie_passport)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_name ON students(last_name, first_name)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_birth_date ON students(birth_date)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_gender ON students(gender)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_students_group ON students(study_group)
        ''')
        
        conn.commit()
        print("✓ Created students table")
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error creating students table: {e}")
        conn.rollback()
        return False


def add_student_relationship_to_tablets(conn: sqlite3.Connection) -> bool:
    """Add assigned_to_student column to tablets table."""
    try:
        cursor = conn.cursor()
        
        if not check_column_exists(conn, 'tablets', 'assigned_to_student'):
            cursor.execute('''
                ALTER TABLE tablets ADD COLUMN assigned_to_student INTEGER
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_tablets_student ON tablets(assigned_to_student)
            ''')
            conn.commit()
            print("✓ Added assigned_to_student column to tablets table")
        else:
            print("✓ assigned_to_student column already exists in tablets table")
        
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error adding assigned_to_student to tablets: {e}")
        conn.rollback()
        return False


def add_student_relationship_to_non_loanable_devices(conn: sqlite3.Connection) -> bool:
    """Add assigned_to_student column to non_loanable_devices table."""
    try:
        cursor = conn.cursor()
        
        if not check_column_exists(conn, 'non_loanable_devices', 'assigned_to_student'):
            cursor.execute('''
                ALTER TABLE non_loanable_devices ADD COLUMN assigned_to_student INTEGER
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_non_loanable_student ON non_loanable_devices(assigned_to_student)
            ''')
            conn.commit()
            print("✓ Added assigned_to_student column to non_loanable_devices table")
        else:
            print("✓ assigned_to_student column already exists in non_loanable_devices table")
        
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error adding assigned_to_student to non_loanable_devices: {e}")
        conn.rollback()
        return False


def add_student_relationship_to_loans(conn: sqlite3.Connection) -> bool:
    """Add student_id column to loans table."""
    try:
        cursor = conn.cursor()
        
        if not check_column_exists(conn, 'loans', 'student_id'):
            cursor.execute('''
                ALTER TABLE loans ADD COLUMN student_id INTEGER
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_loans_student ON loans(student_id)
            ''')
            conn.commit()
            print("✓ Added student_id column to loans table")
        else:
            print("✓ student_id column already exists in loans table")
        
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error adding student_id to loans: {e}")
        conn.rollback()
        return False


def add_foreign_keys(conn: sqlite3.Connection) -> bool:
    """Add foreign key constraints to tables."""
    try:
        cursor = conn.cursor()
        
        # For tablets table
        if check_column_exists(conn, 'tablets', 'assigned_to_student'):
            # SQLite doesn't support adding foreign keys with ALTER TABLE easily
            # We'll rely on application-level enforcement
            print("✓ Foreign key for tablets.assigned_to_student will be enforced at application level")
        
        # For non_loanable_devices table
        if check_column_exists(conn, 'non_loanable_devices', 'assigned_to_student'):
            print("✓ Foreign key for non_loanable_devices.assigned_to_student will be enforced at application level")
        
        # For loans table
        if check_column_exists(conn, 'loans', 'student_id'):
            print("✓ Foreign key for loans.student_id will be enforced at application level")
        
        conn.commit()
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error adding foreign keys: {e}")
        conn.rollback()
        return False


def add_timestamp_columns(conn: sqlite3.Connection) -> bool:
    """Add created_at and updated_at columns to existing tables if they don't exist."""
    try:
        cursor = conn.cursor()
        
        tables_to_update = ['tablets', 'users', 'loans', 'non_loanable_devices']
        
        for table in tables_to_update:
            if not check_column_exists(conn, table, 'created_at'):
                cursor.execute(f'''
                    ALTER TABLE {table} ADD COLUMN created_at TEXT
                ''')
                # Update existing rows
                cursor.execute(f'''
                    UPDATE {table} SET created_at = CURRENT_TIMESTAMP WHERE created_at IS NULL
                ''')
                print(f"✓ Added created_at column to {table} table")
            
            if not check_column_exists(conn, table, 'updated_at'):
                cursor.execute(f'''
                    ALTER TABLE {table} ADD COLUMN updated_at TEXT
                ''')
                # Update existing rows
                cursor.execute(f'''
                    UPDATE {table} SET updated_at = CURRENT_TIMESTAMP WHERE updated_at IS NULL
                ''')
                print(f"✓ Added updated_at column to {table} table")
        
        conn.commit()
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Error adding timestamp columns: {e}")
        conn.rollback()
        return False


def migrate_database(database_path: str, dry_run: bool = False, verbose: bool = False) -> bool:
    """
    Perform the database migration.
    
    Args:
        database_path: Path to the database file
        dry_run: If True, show steps without executing
        verbose: If True, show detailed information
        
    Returns:
        True if migration succeeded, False otherwise
    """
    if dry_run:
        print("[DRY RUN] Migration steps:")
    else:
        print("Starting database migration...")
        # Create backup
        backup_path = backup_database(database_path)
    
    try:
        conn = sqlite3.connect(database_path)
        conn.execute("PRAGMA foreign_keys = ON")
        
        if dry_run:
            print("  1. Would create students table")
            print("  2. Would add assigned_to_student to tablets table")
            print("  3. Would add assigned_to_student to non_loanable_devices table")
            print("  4. Would add student_id to loans table")
            print("  5. Would add timestamp columns to existing tables")
            print("  6. Would create indexes")
        else:
            print("\nRunning migration steps:")
            
            # Step 1: Create students table
            if not create_students_table(conn):
                return False
            
            # Step 2: Add student relationship to tablets
            if not add_student_relationship_to_tablets(conn):
                return False
            
            # Step 3: Add student relationship to non_loanable_devices
            if not add_student_relationship_to_non_loanable_devices(conn):
                return False
            
            # Step 4: Add student relationship to loans
            if not add_student_relationship_to_loans(conn):
                return False
            
            # Step 5: Add timestamp columns
            if not add_timestamp_columns(conn):
                return False
            
            # Step 6: Add foreign keys (enforced at application level)
            if not add_foreign_keys(conn):
                return False
        
        conn.close()
        
        if dry_run:
            print("\n[DRY RUN] No changes were made to the database.")
        else:
            print("\n✓ Database migration completed successfully!")
        
        return True
        
    except Exception as e:
        print(f"\n✗ Migration failed: {e}")
        if not dry_run and 'conn' in locals():
            conn.rollback()
            conn.close()
        return False


def verify_migration(database_path: str) -> bool:
    """Verify that the migration was successful."""
    try:
        conn = sqlite3.connect(database_path)
        cursor = conn.cursor()
        
        print("\nVerifying migration:")
        
        # Check students table exists
        if check_table_exists(conn, 'students'):
            print("✓ students table exists")
            cursor.execute("SELECT COUNT(*) FROM students")
            count = cursor.fetchone()[0]
            print(f"  - Contains {count} students")
        else:
            print("✗ students table missing")
            return False
        
        # Check columns in tablets
        if check_column_exists(conn, 'tablets', 'assigned_to_student'):
            print("✓ tablets.assigned_to_student column exists")
        else:
            print("✗ tablets.assigned_to_student column missing")
            return False
        
        # Check columns in non_loanable_devices
        if check_column_exists(conn, 'non_loanable_devices', 'assigned_to_student'):
            print("✓ non_loanable_devices.assigned_to_student column exists")
        else:
            print("✗ non_loanable_devices.assigned_to_student column missing")
            return False
        
        # Check columns in loans
        if check_column_exists(conn, 'loans', 'student_id'):
            print("✓ loans.student_id column exists")
        else:
            print("✗ loans.student_id column missing")
            return False
        
        # Check indexes
        cursor.execute("SELECT name FROM sqlite_master WHERE type='index'")
        indexes = [row[0] for row in cursor.fetchall()]
        required_indexes = [
            'idx_students_cial', 'idx_students_nif', 'idx_students_name',
            'idx_students_birth_date', 'idx_students_gender', 'idx_students_group',
            'idx_tablets_student', 'idx_non_loanable_student', 'idx_loans_student'
        ]
        for idx in required_indexes:
            if idx in indexes:
                print(f"✓ Index {idx} exists")
            else:
                print(f"✗ Index {idx} missing")
                return False
        
        conn.close()
        print("\n✓ All migration verifications passed!")
        return True
        
    except Exception as e:
        print(f"✗ Verification failed: {e}")
        return False


def main():
    """Main entry point."""
    args = parse_arguments()
    
    print(f"Database: {args.database}")
    print(f"Dry run: {args.dry_run}")
    print()
    
    # Perform migration
    success = migrate_database(
        database_path=args.database,
        dry_run=args.dry_run,
        verbose=args.verbose
    )
    
    if success and not args.dry_run:
        # Verify migration
        verify_migration(args.database)
    
    # Exit with appropriate code
    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()