"""
Migration script to add custom task type columns to projects table
and convert task_type from enum to string.
"""
import sqlite3
import json
from pathlib import Path

# Database path
DB_PATH = Path(__file__).parent / "vfx_project_management.db"


def migrate_database():
    """Add custom task type columns and migrate task_type to string"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    try:
        print("Starting migration for custom task types...")
        
        # Check if columns already exist
        cursor.execute("PRAGMA table_info(projects)")
        columns = [col[1] for col in cursor.fetchall()]
        
        # Add custom_asset_task_types column if it doesn't exist
        if 'custom_asset_task_types' not in columns:
            print("Adding custom_asset_task_types column...")
            cursor.execute("""
                ALTER TABLE projects 
                ADD COLUMN custom_asset_task_types TEXT
            """)
            # Initialize with empty arrays
            cursor.execute("""
                UPDATE projects 
                SET custom_asset_task_types = '[]'
                WHERE custom_asset_task_types IS NULL
            """)
            print("✓ Added custom_asset_task_types column")
        else:
            print("✓ custom_asset_task_types column already exists")
        
        # Add custom_shot_task_types column if it doesn't exist
        if 'custom_shot_task_types' not in columns:
            print("Adding custom_shot_task_types column...")
            cursor.execute("""
                ALTER TABLE projects 
                ADD COLUMN custom_shot_task_types TEXT
            """)
            # Initialize with empty arrays
            cursor.execute("""
                UPDATE projects 
                SET custom_shot_task_types = '[]'
                WHERE custom_shot_task_types IS NULL
            """)
            print("✓ Added custom_shot_task_types column")
        else:
            print("✓ custom_shot_task_types column already exists")
        
        # Note: SQLite doesn't support changing column types directly
        # The task_type column will remain as TEXT in the database
        # The Enum constraint was only enforced at the application level
        print("✓ Task type column already supports string values")
        
        conn.commit()
        print("\n✅ Migration completed successfully!")
        
    except Exception as e:
        conn.rollback()
        print(f"\n❌ Migration failed: {e}")
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    migrate_database()
