#!/usr/bin/env python3
"""
Final database fix script using raw SQL in a single transaction
"""

import asyncio
import asyncpg
import os

async def fix_database():
    # Get DATABASE_URL from environment
    db_url = os.getenv('DATABASE_URL', 'postgresql+asyncpg://postgres:postgres@sf_postgres:5432/service_finder')
    # Convert to sync URL for asyncpg
    sync_url = db_url.replace('+asyncpg', '').replace('postgresql://', 'postgres://')
    
    print("=== DATABASE SURGERY STARTED ===")
    
    conn = await asyncpg.connect(sync_url)
    
    try:
        # Start transaction
        await conn.execute("BEGIN")
        
        # 1. Fix Private Organization (ID 21)
        print("1. Fixing Private Organization owner_id...")
        result = await conn.execute("""
            UPDATE fleet.organizations 
            SET owner_id = 29, updated_at = NOW()
            WHERE id = 21 AND org_type = 'individual'
        """)
        print(f"   ✓ Updated {result.split()[1]} rows")
        
        # 2. Fix Corporate Organization (ID 15)
        print("2. Fixing Corporate Organization owner_id...")
        result = await conn.execute("""
            UPDATE fleet.organizations 
            SET owner_id = 29, updated_at = NOW()
            WHERE id = 15 AND org_type = 'fleet_owner'
        """)
        print(f"   ✓ Updated {result.split()[1]} rows")
        
        # 3. Create Alpha Organization
        print("3. Creating Alpha Organization...")
        alpha_org = await conn.fetchrow("""
            INSERT INTO fleet.organizations (
                full_name, name, folder_slug, org_type, owner_id, status, is_active,
                first_registered_at, current_lifecycle_started_at, lifecycle_index,
                is_anonymized, default_currency, country_code, language,
                subscription_plan, base_asset_limit, purchased_extra_slots,
                notification_settings, external_integration_config, is_verified,
                created_at, updated_at, is_ownership_transferable
            ) VALUES (
                $1, $2, $3, 'fleet_owner', 29, 'active', true,
                NOW(), NOW(), 1,
                false, 'HUF', 'HU', 'hu',
                'FREE', 1, 0,
                '{"notify_owner": true, "alert_days_before": [30, 15, 7, 1]}'::jsonb,
                '{}'::jsonb, false,
                NOW(), NOW(), true
            ) RETURNING id, folder_slug
        """, 'Test Kft. Alpha', 'Test Kft. Alpha', f'alpha{os.urandom(3).hex()}')
        
        print(f"   ✓ Created Alpha Organization (id={alpha_org['id']}, slug={alpha_org['folder_slug']})")
        
        # Create Alpha Branch
        import uuid
        alpha_branch_id = uuid.uuid4()
        await conn.execute("""
            INSERT INTO fleet.branches (
                id, organization_id, name, is_main,
                branch_rating, opening_hours, status, is_deleted, created_at
            ) VALUES ($1, $2, $3, $4, 0.0, '{}'::jsonb, 'active', false, NOW())
        """, alpha_branch_id, alpha_org['id'], 'Alpha Main Garage', True)
        print(f"   ✓ Created Alpha Branch (id={alpha_branch_id})")
        
        # 4. Create Beta Organization
        print("4. Creating Beta Organization...")
        beta_org = await conn.fetchrow("""
            INSERT INTO fleet.organizations (
                full_name, name, folder_slug, org_type, owner_id, status, is_active,
                first_registered_at, current_lifecycle_started_at, lifecycle_index,
                is_anonymized, default_currency, country_code, language,
                subscription_plan, base_asset_limit, purchased_extra_slots,
                notification_settings, external_integration_config, is_verified,
                created_at, updated_at, is_ownership_transferable
            ) VALUES (
                $1, $2, $3, 'fleet_owner', 29, 'active', true,
                NOW(), NOW(), 1,
                false, 'HUF', 'HU', 'hu',
                'FREE', 1, 0,
                '{"notify_owner": true, "alert_days_before": [30, 15, 7, 1]}'::jsonb,
                '{}'::jsonb, false,
                NOW(), NOW(), true
            ) RETURNING id, folder_slug
        """, 'Test Kft. Beta', 'Test Kft. Beta', f'beta{os.urandom(3).hex()}')
        
        print(f"   ✓ Created Beta Organization (id={beta_org['id']}, slug={beta_org['folder_slug']})")
        
        # Create Beta Branch
        beta_branch_id = uuid.uuid4()
        await conn.execute("""
            INSERT INTO fleet.branches (
                id, organization_id, name, is_main,
                branch_rating, opening_hours, status, is_deleted, created_at
            ) VALUES ($1, $2, $3, $4, 0.0, '{}'::jsonb, 'active', false, NOW())
        """, beta_branch_id, beta_org['id'], 'Beta Main Garage', True)
        print(f"   ✓ Created Beta Branch (id={beta_branch_id})")
        
        # 5. Get first 3 vehicles owned by person_id 29
        print("5. Distributing vehicles...")
        vehicles = await conn.fetch("""
            SELECT id, license_plate
            FROM vehicle.assets 
            WHERE owner_person_id = 29 
            AND license_plate IS NOT NULL
            ORDER BY id
            LIMIT 3
        """)
        
        if len(vehicles) >= 3:
            # Update vehicle 1 to Private Organization (ID 21)
            await conn.execute("""
                UPDATE vehicle.assets 
                SET current_organization_id = 21
                WHERE id = $1
            """, vehicles[0]['id'])
            print(f"   ✓ Vehicle '{vehicles[0]['license_plate']}' -> Private Organization")
            
            # Update vehicle 2 to Alpha Organization
            await conn.execute("""
                UPDATE vehicle.assets 
                SET current_organization_id = $1
                WHERE id = $2
            """, alpha_org['id'], vehicles[1]['id'])
            print(f"   ✓ Vehicle '{vehicles[1]['license_plate']}' -> Alpha Organization")
            
            # Update vehicle 3 to Beta Organization
            await conn.execute("""
                UPDATE vehicle.assets 
                SET current_organization_id = $1
                WHERE id = $2
            """, beta_org['id'], vehicles[2]['id'])
            print(f"   ✓ Vehicle '{vehicles[2]['license_plate']}' -> Beta Organization")
            
            # Update asset assignments
            # Delete existing assignments
            await conn.execute("""
                DELETE FROM fleet.asset_assignments 
                WHERE asset_id IN ($1, $2, $3)
            """, vehicles[0]['id'], vehicles[1]['id'], vehicles[2]['id'])
            
            # Create new assignments
            await conn.execute("""
                INSERT INTO fleet.asset_assignments (asset_id, organization_id)
                VALUES ($1, 21), ($2, $3), ($4, $5)
            """, vehicles[0]['id'], vehicles[1]['id'], alpha_org['id'], vehicles[2]['id'], beta_org['id'])
            
            print(f"   ✓ Updated asset assignments")
        else:
            print(f"   ⚠️ Not enough vehicles (need 3, have {len(vehicles)})")
        
        # Commit transaction
        await conn.execute("COMMIT")
        print("\n=== DATABASE SURGERY COMPLETED SUCCESSFULLY ===")
        
    except Exception as e:
        await conn.execute("ROLLBACK")
        print(f"\n=== ERROR: {e} ===")
        raise
    finally:
        await conn.close()

if __name__ == "__main__":
    asyncio.run(fix_database())