#!/usr/bin/env python3
"""
Add Multi-Asset Support to Database

This script modifies the database to support multiple assets of the same type
per report (e.g., 3 audio files, 3 video files).
"""
import os
import sys
import psycopg2
from psycopg2.extras import RealDictCursor

def main():
    print("=" * 80)
    print("ADD MULTI-ASSET SUPPORT TO DATABASE")
    print("=" * 80)

    try:
        conn = psycopg2.connect(
            host="127.0.0.1",
            port=5432,
            database="jadwa",
            user="postgres",
        )

        cursor = conn.cursor()

        print("\n🔍 Step 1: Check current database schema...")

        # Check if asset_index column exists
        cursor.execute("""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = 'notebook_assets'
              AND column_name = 'asset_index';
        """)

        has_asset_index = cursor.fetchone()

        if has_asset_index:
            print("  ✅ asset_index column already exists")
        else:
            print("  📋 Adding asset_index column...")

            # Add asset_index column
            cursor.execute("""
                ALTER TABLE notebook_assets
                ADD COLUMN asset_index INTEGER DEFAULT 1;
            """)

            # Update existing assets to have index 1
            cursor.execute("""
                UPDATE notebook_assets
                SET asset_index = 1
                WHERE asset_index IS NULL;
            """)

            print("  ✅ asset_index column added")

        print("\n🔍 Step 2: Check constraints...")

        # Check if unique constraint exists
        cursor.execute("""
            SELECT constraint_name
            FROM information_schema.table_constraints
            WHERE table_name = 'notebook_assets'
              AND constraint_type = 'UNIQUE';
        """)

        constraints = [row[0] for row in cursor.fetchall()]
        print(f"  Current constraints: {constraints}")

        # Check if we need to add the new unique constraint
        if 'notebook_assets_report_asset_type_index' in constraints:
            print("  ✅ Multi-asset constraint already exists")
        else:
            print("  📋 Adding multi-asset unique constraint...")

            # Drop old unique constraint if it exists
            if 'notebook_assets_report_id_asset_type_key' in constraints:
                cursor.execute("""
                    ALTER TABLE notebook_assets
                    DROP CONSTRAINT notebook_assets_report_id_asset_type_key;
                """)
                print("  ✅ Dropped old unique constraint")

            # Add new unique constraint for (report_id, asset_type, asset_index)
            cursor.execute("""
                ALTER TABLE notebook_assets
                ADD CONSTRAINT notebook_assets_report_asset_type_index
                UNIQUE (report_id, asset_type, asset_index);
            """)
            print("  ✅ Added multi-asset unique constraint")

        conn.commit()

        print("\n🔍 Step 3: Add additional audio and video assets for report 35...")

        # Check current assets for report 35
        cursor.execute("""
            SELECT asset_type, COUNT(*) as count
            FROM notebook_assets
            WHERE report_id = 35
            GROUP BY asset_type
            ORDER BY asset_type;
        """)

        current_assets = {row[0]: row[1] for row in cursor.fetchall()}
        print(f"  Current assets: {current_assets}")

        # Add additional audio assets
        audio_count = current_assets.get('audio', 0)
        if audio_count < 3:
            print(f"  📋 Adding {3 - audio_count} additional audio assets...")
            for i in range(audio_count + 1, 4):
                cursor.execute("""
                    INSERT INTO notebook_assets (report_id, asset_type, status, asset_index, created_at)
                    VALUES (35, 'audio', 'pending', %s, NOW())
                    ON CONFLICT (report_id, asset_type, asset_index)
                    DO NOTHING;
                """, (i,))
                print(f"    ✅ Added audio_{i}")

        # Add additional video assets
        video_count = current_assets.get('video', 0)
        if video_count < 3:
            print(f"  📋 Adding {3 - video_count} additional video assets...")
            for i in range(video_count + 1, 4):
                cursor.execute("""
                    INSERT INTO notebook_assets (report_id, asset_type, status, asset_index, created_at)
                    VALUES (35, 'video', 'pending', %s, NOW())
                    ON CONFLICT (report_id, asset_type, asset_index)
                    DO NOTHING;
                """, (i,))
                print(f"    ✅ Added video_{i}")

        # Clear error messages for assets with valid files
        print("\n🔍 Step 4: Clear error messages for assets with valid files...")
        cursor.execute("""
            UPDATE notebook_assets
            SET error_message = NULL
            WHERE report_id = 35
              AND status = 'ready'
              AND error_message IS NOT NULL
              AND file_path IS NOT NULL;
        """)
        cleared = cursor.rowcount
        if cleared > 0:
            print(f"  ✅ Cleared {cleared} error messages for ready assets")

        conn.commit()

        print("\n🔍 Step 5: Verify the changes...")

        cursor.execute("""
            SELECT asset_type, asset_index, status,
                   CASE WHEN file_path IS NOT NULL THEN 'Yes' ELSE 'No' END as has_file
            FROM notebook_assets
            WHERE report_id = 35
            ORDER BY asset_type, asset_index;
        """)

        assets = cursor.fetchall()

        print(f"\n📊 Current assets for report 35:")
        print(f"{'Type':<12} {'Index':<6} {'Status':<12} {'File':<6}")
        print("-" * 40)
        for asset in assets:
            status_icon = "✓" if asset[2] == 'ready' else ("⏳" if asset[2] == 'generating' else "✗")
            print(f"{asset[0]:<12} {asset[1]:<6} {asset[2]:<10} {status_icon} {asset[3]:<6}")

        cursor.close()
        conn.close()

        print("\n" + "=" * 80)
        print("✅ MULTI-ASSET SUPPORT ADDED SUCCESSFULLY")
        print("=" * 80)
        print("Database now supports multiple assets per type per report.")
        print("Report 35 has been configured for 3 audio + 3 video files.")

    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()