#!/usr/bin/env python3
"""
Database migration script to fix TLDraw ID issues
"""

import os
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

def migrate_database():
    """Migrate the database to fix TLDraw ID issues"""
    
    # Database connection details
    db_host = os.getenv('POSTGRES_HOST', 'localhost')
    db_user = os.getenv('POSTGRES_USER', 'pipeline_admin')
    db_password = os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
    db_port = os.getenv('POSTGRES_PORT', '5433')
    db_name = os.getenv('POSTGRES_DB', 'dev_pipeline')
    
    try:
        conn = psycopg2.connect(
            host=db_host,
            user=db_user,
            password=db_password,
            port=db_port,
            database=db_name
        )
        
        with conn.cursor() as cur:
            print("🔄 Migrating database to fix TLDraw ID issues...")
            
            # Check if tldraw_id column exists
            cur.execute("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = 'wireframe_elements' 
                AND column_name = 'tldraw_id'
            """)
            
            if not cur.fetchone():
                print("   Adding tldraw_id column...")
                cur.execute("ALTER TABLE wireframe_elements ADD COLUMN tldraw_id VARCHAR(255)")
            
            # Check if parent_id is already VARCHAR
            cur.execute("""
                SELECT data_type 
                FROM information_schema.columns 
                WHERE table_name = 'wireframe_elements' 
                AND column_name = 'parent_id'
            """)
            
            column_info = cur.fetchone()
            if column_info and column_info[0] == 'uuid':
                print("   Converting parent_id from UUID to VARCHAR...")
                # Drop the foreign key constraint first
                cur.execute("""
                    ALTER TABLE wireframe_elements 
                    DROP CONSTRAINT IF EXISTS wireframe_elements_parent_id_fkey
                """)
                
                # Change the column type
                cur.execute("""
                    ALTER TABLE wireframe_elements 
                    ALTER COLUMN parent_id TYPE VARCHAR(255)
                """)
            
            # Update the function
            print("   Updating get_wireframe_with_elements function...")
            cur.execute("""
                CREATE OR REPLACE FUNCTION get_wireframe_with_elements(p_wireframe_id UUID)
                RETURNS TABLE(
                  wireframe_data JSONB,
                  elements_data JSONB
                ) AS $$
                BEGIN
                  RETURN QUERY
                  SELECT 
                    to_jsonb(w.*) as wireframe_data,
                    COALESCE(
                      jsonb_agg(
                        jsonb_build_object(
                          'id', we.id,
                          'tldraw_id', we.tldraw_id,
                          'type', we.element_type,
                          'data', we.element_data,
                          'position', we.position,
                          'size', we.size,
                          'style', we.style,
                          'parent_id', we.parent_id,
                          'z_index', we.z_index
                        ) ORDER BY we.z_index, we.created_at
                      ) FILTER (WHERE we.id IS NOT NULL),
                      '[]'::jsonb
                    ) as elements_data
                  FROM wireframes w
                  LEFT JOIN wireframe_elements we ON w.id = we.wireframe_id
                  WHERE w.id = p_wireframe_id
                  GROUP BY w.id, w.user_id, w.project_id, w.name, w.description, 
                           w.device_type, w.dimensions, w.metadata, w.is_active, 
                           w.created_at, w.updated_at;
                END;
                $$ LANGUAGE plpgsql;
            """)
            
            conn.commit()
            print("✅ Database migration completed successfully!")
            
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ Database migration failed: {e}")
        return False

if __name__ == "__main__":
    print("🚀 Starting database migration...")
    success = migrate_database()
    
    if success:
        print("\n✅ Migration completed successfully!")
        print("The database now supports TLDraw string IDs properly.")
    else:
        print("\n❌ Migration failed!")
        print("Please check the database connection and try again.")