#!/usr/bin/env python3
"""
Complete database surgery for tester_pro organizations and vehicles.
Creates branches for Alpha/Beta organizations, distributes vehicles, and sets up asset assignments.
"""
import asyncio
import asyncpg
import os
import uuid

async def fix_database():
    # Use the correct connection string for asyncpg
    DATABASE_URL = "postgresql://service_finder_app:AppSafePass_2026@db:5432/service_finder"
    
    print("=== DATABASE SURGERY STARTED ===")
    
    conn = await asyncpg.connect(DATABASE_URL)
    
    try:
        # Start transaction
        await conn.execute("BEGIN")
        
        print("1. Checking existing organizations...")
        
        # Get tester_pro's person_id (should be 29)
        person = await conn.fetchrow("""
            SELECT id FROM identity.persons WHERE email = 'tester_pro@profibot.hu'
        """)
        if not person:
            print("   ❌ tester_pro not found!")
            return
        person_id = person['id']
        print(f"   ✓ Found tester_pro with person_id={person_id}")
        
        # Check existing organizations
        orgs = await conn.fetch("""
            SELECT id, full_name, org_type FROM fleet.organizations 
            WHERE owner_id = $1 ORDER BY id
        """, person_id)
        
        print(f"   ✓ Found {len(orgs)} organizations for tester_pro")
        for org in orgs:
            print(f"     - ID {org['id']}: {org['full_name']} ({org['org_type']})")
        
        # 1. Ensure Private Organization has correct owner_id (should already be fixed)
        print("\n2. Fixing Private Organization owner_id...")
        private_org = await conn.fetchrow("""
            SELECT id FROM fleet.organizations 
            WHERE owner_id = $1 AND org_type = 'individual'
        """, person_id)
        
        if private_org:
            print(f"   ✓ Private Organization exists (ID: {private_org['id']})")
        else:
            print("   ⚠️ No Private Organization found")
        
        # 2. Fix Corporate Organization owner_id (should already be fixed)
        print("\n3. Fixing Corporate Organization owner_id...")
        corp_org = await conn.fetchrow("""
            SELECT id FROM fleet.organizations 
            WHERE owner_id = $1 AND org_type = 'fleet_owner' AND full_name LIKE '%Profibot Test Fleet%'
        """, person_id)
        
        if corp_org:
            print(f"   ✓ Corporate Organization exists (ID: {corp_org['id']})")
        else:
            print("   ⚠️ No Corporate Organization found")
        
        # 3. Create branches for Alpha and Beta organizations if they don't exist
        print("\n4. Creating branches for Alpha and Beta organizations...")
        
        # Check Alpha organization
        alpha_org = await conn.fetchrow("""
            SELECT id, full_name FROM fleet.organizations 
            WHERE owner_id = $1 AND full_name = 'Test Kft. Alpha'
        """, person_id)
        
        if alpha_org:
            print(f"   ✓ Alpha Organization exists (ID: {alpha_org['id']})")
            # Check if Alpha has a branch
            alpha_branch = await conn.fetchrow("""
                SELECT id FROM fleet.branches WHERE organization_id = $1
            """, alpha_org['id'])
            
            if not alpha_branch:
                print("   Creating Alpha Branch...")
                alpha_branch_id = str(uuid.uuid4())
                await conn.execute("""
                    INSERT INTO fleet.branches (
                        id, name, organization_id, branch_rating, opening_hours, 
                        status, is_deleted, created_at
                    ) VALUES (
                        $1, $2, $3, 0.0, '{}'::jsonb, 'active', false, NOW()
                    )
                """, alpha_branch_id, f"{alpha_org['full_name']} - Main Branch", alpha_org['id'])
                print(f"   ✓ Created Alpha Branch (ID: {alpha_branch_id})")
            else:
                print(f"   ✓ Alpha Branch already exists (ID: {alpha_branch['id']})")
        
        # Check Beta organization
        beta_org = await conn.fetchrow("""
            SELECT id, full_name FROM fleet.organizations 
            WHERE owner_id = $1 AND full_name = 'Test Kft. Beta'
        """, person_id)
        
        if beta_org:
            print(f"   ✓ Beta Organization exists (ID: {beta_org['id']})")
            # Check if Beta has a branch
            beta_branch = await conn.fetchrow("""
                SELECT id FROM fleet.branches WHERE organization_id = $1
            """, beta_org['id'])
            
            if not beta_branch:
                print("   Creating Beta Branch...")
                beta_branch_id = str(uuid.uuid4())
                await conn.execute("""
                    INSERT INTO fleet.branches (
                        id, name, organization_id, branch_rating, opening_hours, 
                        status, is_deleted, created_at
                    ) VALUES (
                        $1, $2, $3, 0.0, '{}'::jsonb, 'active', false, NOW()
                    )
                """, beta_branch_id, f"{beta_org['full_name']} - Main Branch", beta_org['id'])
                print(f"   ✓ Created Beta Branch (ID: {beta_branch_id})")
            else:
                print(f"   ✓ Beta Branch already exists (ID: {beta_branch['id']})")
        
        # 4. Distribute vehicles
        print("\n5. Distributing vehicles...")
        
        # Get 3 vehicles owned by tester_pro
        vehicles = await conn.fetch("""
            SELECT id, license_plate, current_organization_id
            FROM vehicle.assets 
            WHERE owner_person_id = $1
            ORDER BY license_plate
            LIMIT 3
        """, person_id)
        
        if len(vehicles) >= 3:
            # Vehicle 1: Keep in Private Organization (ID 21)
            private_org_id = 21  # From earlier check
            await conn.execute("""
                UPDATE vehicle.assets 
                SET current_organization_id = $1
                WHERE id = $2
            """, private_org_id, vehicles[0]['id'])
            print(f"   ✓ Vehicle '{vehicles[0]['license_plate']}' -> Private Organization")
            
            # Vehicle 2: Move to Alpha Organization
            if alpha_org:
                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")
            
            # Vehicle 3: Move to Beta Organization
            if beta_org:
                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")
            
            # 5. Update asset assignments with proper UUIDs and status
            print("\n6. Updating asset assignments...")
            
            # Delete existing assignments for these vehicles
            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 with UUIDs and status
            assignments = [
                (str(uuid.uuid4()), vehicles[0]['id'], private_org_id, 'active'),
                (str(uuid.uuid4()), vehicles[1]['id'], alpha_org['id'] if alpha_org else None, 'active'),
                (str(uuid.uuid4()), vehicles[2]['id'], beta_org['id'] if beta_org else None, 'active')
            ]
            
            for assignment_id, asset_id, org_id, status in assignments:
                if org_id:  # Skip if organization doesn't exist
                    await conn.execute("""
                        INSERT INTO fleet.asset_assignments (id, asset_id, organization_id, status)
                        VALUES ($1, $2, $3, $4)
                    """, assignment_id, asset_id, org_id, status)
            
            print(f"   ✓ Created {len([a for a in assignments if a[2]])} 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())