#!/usr/bin/env python3
"""
Migration script to add soft deletion columns to related tables.
Adds deleted_at and deleted_by columns to submissions, task_attachments, production_notes, and reviews tables.
Creates appropriate partial indexes for each table.
"""

import sqlite3
import sys
from pathlib import Path

def migrate_related_tables_soft_deletion():
    """Add soft deletion columns to related tables and create partial indexes."""
    
    # Database path
    db_path = Path(__file__).parent / "vfx_project_management.db"
    
    if not db_path.exists():
        print(f"Database file not found at {db_path}")
        return False
    
    try:
        # Connect to database
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        print("Starting related tables soft deletion migration...")
        
        # Tables to migrate
        tables = [
            ("submissions", "task_id"),
            ("task_attachments", "task_id"),
            ("production_notes", "task_id"),
            ("reviews", "submission_id")
        ]
        
        for table_name, index_column in tables:
            print(f"\n--- Migrating {table_name} table ---")
            
            # Check if columns already exist
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = [column[1] for column in cursor.fetchall()]
            
            if 'deleted_at' in columns:
                print(f"Soft deletion columns already exist in {table_name} table")
                continue
            
            # Add deleted_at column
            print(f"Adding deleted_at column to {table_name} table...")
            cursor.execute(f"""
                ALTER TABLE {table_name} 
                ADD COLUMN deleted_at TIMESTAMP NULL
            """)
            
            # Add deleted_by column
            print(f"Adding deleted_by column to {table_name} table...")
            cursor.execute(f"""
                ALTER TABLE {table_name} 
                ADD COLUMN deleted_by INTEGER NULL 
                REFERENCES users(id)
            """)
            
            # Create partial index for efficient querying of non-deleted records
            index_name = f"idx_{table_name}_not_deleted"
            print(f"Creating partial index {index_name}...")
            cursor.execute(f"""
                CREATE INDEX {index_name} 
                ON {table_name} ({index_column}) 
                WHERE deleted_at IS NULL
            """)
            
            print(f"SUCCESS: Successfully migrated {table_name} table")
        
        # Commit changes
        conn.commit()
        print("\nSUCCESS: Successfully added soft deletion columns to all related tables")
        
        return True
        
    except sqlite3.Error as e:
        print(f"Database error: {e}")
        if conn:
            conn.rollback()
        return False
    except Exception as e:
        print(f"Unexpected error: {e}")
        if conn:
            conn.rollback()
        return False
    finally:
        if conn:
            conn.close()

def verify_migration():
    """Verify that the migration was successful."""
    
    db_path = Path(__file__).parent / "vfx_project_management.db"
    
    try:
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        # Tables to verify
        tables = [
            ("submissions", "idx_submissions_not_deleted"),
            ("task_attachments", "idx_task_attachments_not_deleted"),
            ("production_notes", "idx_production_notes_not_deleted"),
            ("reviews", "idx_reviews_not_deleted")
        ]
        
        all_verified = True
        
        for table_name, index_name in tables:
            print(f"\n--- Verifying {table_name} table ---")
            
            # Check table structure
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = {column[1]: column[2] for column in cursor.fetchall()}
            
            # Verify columns exist
            if 'deleted_at' not in columns:
                print(f"❌ deleted_at column not found in {table_name}")
                all_verified = False
                continue
            if 'deleted_by' not in columns:
                print(f"❌ deleted_by column not found in {table_name}")
                all_verified = False
                continue
                
            print(f"SUCCESS: Soft deletion columns verified in {table_name} table")
            
            # Check index exists
            cursor.execute(f"PRAGMA index_list({table_name})")
            indexes = [index[1] for index in cursor.fetchall()]
            
            if index_name not in indexes:
                print(f"❌ Partial index {index_name} not found")
                all_verified = False
                continue
                
            print(f"SUCCESS: Partial index {index_name} verified")
        
        return all_verified
        
    except sqlite3.Error as e:
        print(f"Verification error: {e}")
        return False
    finally:
        if conn:
            conn.close()

if __name__ == "__main__":
    print("=== Related Tables Soft Deletion Migration ===")
    
    success = migrate_related_tables_soft_deletion()
    
    if success:
        print("\n=== Verifying Migration ===")
        if verify_migration():
            print("\nSUCCESS: Migration completed successfully!")
            sys.exit(0)
        else:
            print("\n❌ Migration verification failed!")
            sys.exit(1)
    else:
        print("\n❌ Migration failed!")
        sys.exit(1)