#!/usr/bin/env python3
"""
Migration script to add new fields to the projects table:
- code_name: Unique project code identifier
- client_name: Client or studio name  
- project_type: TV, Cinema, or Game
"""

import sqlite3
import sys
from pathlib import Path

def migrate_project_fields():
    """Add new fields to the projects table"""
    
    # Database path
    db_path = Path("vfx_project_management.db")
    
    if not db_path.exists():
        print("Database file not found. Creating new database with updated schema.")
        return True
    
    try:
        # Connect to database
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # Check if new columns already exist
        cursor.execute("PRAGMA table_info(projects)")
        columns = [column[1] for column in cursor.fetchall()]
        
        new_columns_needed = []
        if 'code_name' not in columns:
            new_columns_needed.append('code_name')
        if 'client_name' not in columns:
            new_columns_needed.append('client_name')
        if 'project_type' not in columns:
            new_columns_needed.append('project_type')
            
        if not new_columns_needed:
            print("All new columns already exist. No migration needed.")
            return True
            
        print(f"Adding columns: {', '.join(new_columns_needed)}")
        
        # Add new columns
        if 'code_name' in new_columns_needed:
            cursor.execute("ALTER TABLE projects ADD COLUMN code_name VARCHAR")
            print("Added code_name column")
            
        if 'client_name' in new_columns_needed:
            cursor.execute("ALTER TABLE projects ADD COLUMN client_name VARCHAR")
            print("Added client_name column")
            
        if 'project_type' in new_columns_needed:
            cursor.execute("ALTER TABLE projects ADD COLUMN project_type VARCHAR")
            print("Added project_type column")
        
        # Update existing projects with default values
        cursor.execute("SELECT id, name FROM projects WHERE code_name IS NULL OR client_name IS NULL OR project_type IS NULL")
        existing_projects = cursor.fetchall()
        
        for project_id, project_name in existing_projects:
            # Generate default code_name from project name
            code_name = project_name.upper().replace(' ', '_')[:10] + f"_{project_id:03d}"
            
            cursor.execute("""
                UPDATE projects 
                SET code_name = COALESCE(code_name, ?),
                    client_name = COALESCE(client_name, 'Default Client'),
                    project_type = COALESCE(project_type, 'tv')
                WHERE id = ?
            """, (code_name, project_id))
            
        print(f"Updated {len(existing_projects)} existing projects with default values")
        
        # Create unique index on code_name
        try:
            cursor.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_projects_code_name ON projects(code_name)")
            print("Created unique index on code_name")
        except sqlite3.IntegrityError as e:
            print(f"Warning: Could not create unique index on code_name: {e}")
            print("You may need to manually ensure code_name values are unique")
        
        # Commit changes
        conn.commit()
        print("Migration completed successfully!")
        
        return True
        
    except Exception as e:
        print(f"Migration failed: {e}")
        if conn:
            conn.rollback()
        return False
        
    finally:
        if conn:
            conn.close()

if __name__ == "__main__":
    success = migrate_project_fields()
    sys.exit(0 if success else 1)