#!/usr/bin/env python3
"""
Database Surgery Script for Organization and Vehicle Fix
Fixes:
1. Update Private Organization owner_id from None to tester_pro's person_id
2. Update Corporate Organization owner_id from user_id to person_id
3. Create 2 new Corporate Orgs (Alpha, Beta) with Branches
4. Distribute vehicles: 1 to Private, 1 to Alpha, 1 to Beta
"""

import asyncio
import uuid
from datetime import datetime
from sqlalchemy import select, update, insert
from app.database import AsyncSessionLocal
from app.models.identity.identity import User, Person
from app.models.marketplace.organization import Organization, Branch, OrgType
from app.models.vehicle.asset import Asset, AssetAssignment


async def fix_database():
    async with AsyncSessionLocal() as session:
        print("=== DATABASE SURGERY STARTED ===")
        
        # 1. Find tester_pro user and person
        result = await session.execute(
            select(User).where(User.email == 'tester_pro@profibot.hu')
        )
        user = result.scalar_one_or_none()
        
        if not user:
            print("ERROR: tester_pro user not found!")
            return
        
        print(f"Found user: {user.email} (id={user.id}, person_id={user.person_id})")
        
        # Get person
        result = await session.execute(
            select(Person).where(Person.id == user.person_id)
        )
        person = result.scalar_one_or_none()
        
        if not person:
            print("ERROR: tester_pro person not found!")
            return
        
        print(f"Found person: id={person.id}")
        
        # 2. Fix Private Organization (ID 21) - set owner_id to person.id
        result = await session.execute(
            select(Organization).where(
                Organization.id == 21,
                Organization.org_type == OrgType.individual
            )
        )
        private_org = result.scalar_one_or_none()
        
        if private_org:
            print(f"\n1. Fixing Private Organization: {private_org.full_name}")
            print(f"   Current owner_id: {private_org.owner_id} (should be {person.id})")
            
            if private_org.owner_id != person.id:
                await session.execute(
                    update(Organization)
                    .where(Organization.id == 21)
                    .values(owner_id=person.id)
                )
                print(f"   ✓ Updated owner_id to {person.id}")
            else:
                print(f"   ✓ Already correct")
        else:
            print("\n1. Private Organization (ID 21) not found, creating...")
            # Create private org if it doesn't exist
            private_org = Organization(
                full_name=f"Private Organization for {user.email}",
                org_type=OrgType.individual,
                owner_id=person.id,
                status="active",
                is_active=True
            )
            session.add(private_org)
            await session.flush()
            print(f"   ✓ Created Private Organization with id={private_org.id}")
        
        # 3. Fix Corporate Organization (ID 15) - update owner_id from user.id to person.id
        result = await session.execute(
            select(Organization).where(
                Organization.id == 15,
                Organization.org_type == OrgType.fleet_owner
            )
        )
        corp_org = result.scalar_one_or_none()
        
        if corp_org:
            print(f"\n2. Fixing Corporate Organization: {corp_org.full_name}")
            print(f"   Current owner_id: {corp_org.owner_id} (user_id, should be person_id={person.id})")
            
            if corp_org.owner_id != person.id:
                await session.execute(
                    update(Organization)
                    .where(Organization.id == 15)
                    .values(owner_id=person.id)
                )
                print(f"   ✓ Updated owner_id from {corp_org.owner_id} to {person.id}")
            else:
                print(f"   ✓ Already correct")
        
        # 4. Create 2 new Corporate Orgs: Alpha and Beta
        print(f"\n3. Creating 2 new Corporate Organizations...")
        
        # Alpha Organization
        import uuid
        from datetime import datetime
        alpha_slug = f"a{uuid.uuid4().hex[:6]}"  # 7 characters total
        now = datetime.utcnow()
        alpha_org = Organization(
            full_name="Test Kft. Alpha",
            name="Test Kft. Alpha",  # Required field
            folder_slug=alpha_slug,  # Required unique field, max 12 chars
            org_type=OrgType.fleet_owner,
            owner_id=person.id,
            status="active",
            is_active=True,
            first_registered_at=now,
            current_lifecycle_started_at=now,
            subscription_plan="FREE"  # Required field with default
        )
        session.add(alpha_org)
        await session.flush()
        
        # Alpha Branch
        alpha_branch = Branch(
            organization_id=alpha_org.id,
            name="Alpha Main Garage",
            is_main=True
        )
        session.add(alpha_branch)
        
        print(f"   ✓ Created Alpha Organization: {alpha_org.full_name} (id={alpha_org.id}, slug={alpha_slug})")
        print(f"   ✓ Created Alpha Branch: {alpha_branch.name}")
        
        # Beta Organization
        beta_slug = f"b{uuid.uuid4().hex[:6]}"  # 7 characters total
        beta_org = Organization(
            full_name="Test Kft. Beta",
            name="Test Kft. Beta",  # Required field
            folder_slug=beta_slug,  # Required unique field, max 12 chars
            org_type=OrgType.fleet_owner,
            owner_id=person.id,
            status="active",
            is_active=True,
            first_registered_at=now,
            current_lifecycle_started_at=now,
            subscription_plan="FREE"  # Required field with default
        )
        session.add(beta_org)
        await session.flush()
        
        # Beta Branch
        beta_branch = Branch(
            organization_id=beta_org.id,
            name="Beta Main Garage",
            is_main=True
        )
        session.add(beta_branch)
        
        print(f"   ✓ Created Beta Organization: {beta_org.full_name} (id={beta_org.id}, slug={beta_slug})")
        print(f"   ✓ Created Beta Branch: {beta_branch.name}")
        
        # 5. Get vehicles owned by tester_pro
        result = await session.execute(
            select(Asset).where(Asset.owner_person_id == person.id)
        )
        vehicles = result.scalars().all()
        
        print(f"\n4. Distributing {len(vehicles)} vehicles...")
        
        if len(vehicles) >= 3:
            # Distribute: vehicle 0 -> Private, vehicle 1 -> Alpha, vehicle 2 -> Beta
            private_vehicle = vehicles[0]
            alpha_vehicle = vehicles[1]
            beta_vehicle = vehicles[2]
            
            # Update private vehicle to Private Organization
            await session.execute(
                update(Asset)
                .where(Asset.id == private_vehicle.id)
                .values(current_organization_id=private_org.id)
            )
            print(f"   ✓ Vehicle '{private_vehicle.license_plate}' -> Private Organization")
            
            # Update alpha vehicle to Alpha Organization
            await session.execute(
                update(Asset)
                .where(Asset.id == alpha_vehicle.id)
                .values(current_organization_id=alpha_org.id)
            )
            print(f"   ✓ Vehicle '{alpha_vehicle.license_plate}' -> Alpha Organization")
            
            # Update beta vehicle to Beta Organization
            await session.execute(
                update(Asset)
                .where(Asset.id == beta_vehicle.id)
                .values(current_organization_id=beta_org.id)
            )
            print(f"   ✓ Vehicle '{beta_vehicle.license_plate}' -> Beta Organization")
            
            # Update asset assignments as well
            # First delete existing assignments for these vehicles
            from sqlalchemy import delete
            await session.execute(
                delete(AssetAssignment).where(
                    AssetAssignment.asset_id.in_([private_vehicle.id, alpha_vehicle.id, beta_vehicle.id])
                )
            )
            
            # Create new assignments
            assignments = [
                AssetAssignment(asset_id=private_vehicle.id, organization_id=private_org.id),
                AssetAssignment(asset_id=alpha_vehicle.id, organization_id=alpha_org.id),
                AssetAssignment(asset_id=beta_vehicle.id, organization_id=beta_org.id)
            ]
            
            for assignment in assignments:
                session.add(assignment)
            
            print(f"   ✓ Updated asset assignments")
        else:
            print(f"   ⚠️ Not enough vehicles (need 3, have {len(vehicles)})")
        
        # Commit all changes
        await session.commit()
        
        print(f"\n=== DATABASE SURGERY COMPLETED ===")
        print(f"Summary:")
        print(f"  - Fixed Private Organization owner_id")
        print(f"  - Fixed Corporate Organization owner_id")
        print(f"  - Created 2 new Corporate Orgs: Alpha and Beta")
        print(f"  - Created Branches for each org")
        print(f"  - Distributed 3 vehicles to Private, Alpha, Beta garages")


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