#!/usr/bin/env python3
"""
Migration script to add soft deletion columns to assets table.
Adds deleted_at and deleted_by columns and creates partial index for efficient querying.
"""

import sqlite3
import sys
from pathlib import Path

def migrate_assets_soft_deletion():
    """Add soft deletion columns to assets table and create partial index."""
    
    # 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 assets table soft deletion migration...")
        
        # Check if columns already exist
        cursor.execute("PRAGMA table_info(assets)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'deleted_at' in columns:
            print("Soft deletion columns already exist in assets table")
            return True
        
        # Add deleted_at column
        print("Adding deleted_at column to assets table...")
        cursor.execute("""
            ALTER TABLE assets 
            ADD COLUMN deleted_at TIMESTAMP NULL
        """)
        
        # Add deleted_by column
        print("Adding deleted_by column to assets table...")
        cursor.execute("""
            ALTER TABLE assets 
            ADD COLUMN deleted_by INTEGER NULL 
            REFERENCES users(id)
        """)
        
        # Create partial index for efficient querying of non-deleted records
        print("Creating partial index idx_assets_not_deleted...")
        cursor.execute("""
            CREATE INDEX idx_assets_not_deleted 
            ON assets (id) 
            WHERE deleted_at IS NULL
        """)
        
        # Commit changes
        conn.commit()
        print("SUCCESS: Successfully added soft deletion columns to assets table")
        print("SUCCESS: Created partial index idx_assets_not_deleted")
        
        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()
        
        # Check table structure
        cursor.execute("PRAGMA table_info(assets)")
        columns = {column[1]: column[2] for column in cursor.fetchall()}
        
        # Verify columns exist
        if 'deleted_at' not in columns:
            print("❌ deleted_at column not found")
            return False
        if 'deleted_by' not in columns:
            print("❌ deleted_by column not found")
            return False
            
        print("SUCCESS: Soft deletion columns verified in assets table")
        
        # Check index exists
        cursor.execute("PRAGMA index_list(assets)")
        indexes = [index[1] for index in cursor.fetchall()]
        
        if 'idx_assets_not_deleted' not in indexes:
            print("❌ Partial index idx_assets_not_deleted not found")
            return False
            
        print("SUCCESS: Partial index idx_assets_not_deleted verified")
        
        return True
        
    except sqlite3.Error as e:
        print(f"Verification error: {e}")
        return False
    finally:
        if conn:
            conn.close()

if __name__ == "__main__":
    print("=== Assets Table Soft Deletion Migration ===")
    
    success = migrate_assets_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)