#!/usr/bin/env python3
"""
Migration script to assign existing vehicles to their organization's default garage (branch).
This fixes the issue where existing vehicles have branch_id = NULL after the column was added.

Logic:
1. For each Asset with owner_org_id or operator_org_id
2. Find the default Branch (Garage) for that Organization (is_main = True)
3. Update Asset.branch_id to that Branch's UUID
4. If no default branch exists, create one
"""

import asyncio
import logging
import sys
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import AsyncSessionLocal
from app.models.vehicle.asset import Asset
from app.models.marketplace.organization import Branch, Organization
from app.models.identity import User

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def get_or_create_default_branch(session: AsyncSession, organization_id: int) -> Branch:
    """Get the default branch (is_main = True) for an organization, create if doesn't exist."""
    # Try to find existing default branch
    stmt = select(Branch).where(
        Branch.organization_id == organization_id,
        Branch.is_main == True,
        Branch.is_deleted == False
    )
    result = await session.execute(stmt)
    branch = result.scalar_one_or_none()
    
    if branch:
        logger.info(f"Found default branch {branch.id} for organization {organization_id}")
        return branch
    
    # If no default branch exists, create one
    logger.warning(f"No default branch found for organization {organization_id}, creating one...")
    
    # Get organization name for branch naming
    org_stmt = select(Organization).where(Organization.id == organization_id)
    org_result = await session.execute(org_stmt)
    organization = org_result.scalar_one_or_none()
    
    org_name = organization.name if organization else f"Organization {organization_id}"
    
    # Create default branch
    import uuid
    from datetime import datetime
    from sqlalchemy.sql import func
    
    new_branch = Branch(
        id=uuid.uuid4(),
        organization_id=organization_id,
        name=f"{org_name} - Main Garage",
        is_main=True,
        status="active",
        is_deleted=False,
        created_at=datetime.utcnow()
    )
    
    session.add(new_branch)
    await session.flush()
    logger.info(f"Created default branch {new_branch.id} for organization {organization_id}")
    
    return new_branch

async def migrate_vehicles_to_garages():
    """Main migration function."""
    async with AsyncSessionLocal() as session:
        try:
            # Get all assets that have organization ownership but no branch_id
            stmt = select(Asset).where(
                (Asset.owner_org_id.is_not(None) | Asset.operator_org_id.is_not(None)),
                Asset.branch_id.is_(None)
            )
            result = await session.execute(stmt)
            assets = result.scalars().all()
            
            logger.info(f"Found {len(assets)} assets without branch assignment")
            
            updated_count = 0
            skipped_count = 0
            
            for asset in assets:
                # Determine which organization to use (prefer owner, fallback to operator)
                org_id = asset.owner_org_id or asset.operator_org_id
                
                if not org_id:
                    logger.warning(f"Asset {asset.id} has no organization reference, skipping")
                    skipped_count += 1
                    continue
                
                # Get or create default branch for the organization
                branch = await get_or_create_default_branch(session, org_id)
                
                # Update the asset
                update_stmt = (
                    update(Asset)
                    .where(Asset.id == asset.id)
                    .values(branch_id=branch.id, relocation_performed=True)
                )
                await session.execute(update_stmt)
                
                logger.info(f"Updated asset {asset.id} with branch {branch.id} (org {org_id})")
                updated_count += 1
            
            # Commit all changes
            await session.commit()
            
            logger.info(f"Migration completed: {updated_count} assets updated, {skipped_count} skipped")
            
            # Also update assets that already have branch_id but need relocation_performed flag
            if updated_count > 0:
                stmt = select(Asset).where(
                    Asset.branch_id.is_not(None),
                    Asset.relocation_performed == False
                )
                result = await session.execute(stmt)
                assets_without_flag = result.scalars().all()
                
                for asset in assets_without_flag:
                    update_stmt = (
                        update(Asset)
                        .where(Asset.id == asset.id)
                        .values(relocation_performed=True)
                    )
                    await session.execute(update_stmt)
                
                await session.commit()
                logger.info(f"Updated relocation_performed flag for {len(assets_without_flag)} assets")
            
            return updated_count
            
        except Exception as e:
            await session.rollback()
            logger.error(f"Migration failed: {e}")
            raise

async def verify_migration():
    """Verify the migration results."""
    async with AsyncSessionLocal() as session:
        # Count assets with branch_id
        stmt = select(Asset).where(Asset.branch_id.is_not(None))
        result = await session.execute(stmt)
        assets_with_branch = result.scalars().all()
        
        # Count assets without branch_id but with organizations
        stmt = select(Asset).where(
            (Asset.owner_org_id.is_not(None) | Asset.operator_org_id.is_not(None)),
            Asset.branch_id.is_(None)
        )
        result = await session.execute(stmt)
        assets_still_missing = result.scalars().all()
        
        logger.info(f"Verification:")
        logger.info(f"  - Assets with branch_id: {len(assets_with_branch)}")
        logger.info(f"  - Assets still missing branch_id: {len(assets_still_missing)}")
        
        if assets_still_missing:
            logger.warning("Some assets still missing branch_id:")
            for asset in assets_still_missing[:5]:  # Show first 5
                logger.warning(f"    Asset {asset.id}: owner_org={asset.owner_org_id}, operator_org={asset.operator_org_id}")
        
        return len(assets_with_branch), len(assets_still_missing)

if __name__ == "__main__":
    logger.info("Starting vehicle-to-garage migration...")
    
    try:
        # Run migration
        updated = asyncio.run(migrate_vehicles_to_garages())
        
        # Verify
        with_branch, missing = asyncio.run(verify_migration())
        
        if missing == 0:
            logger.info("✅ Migration successful! All organizational vehicles now have branch assignments.")
        else:
            logger.warning(f"⚠️  Migration incomplete: {missing} assets still lack branch_id")
            sys.exit(1)
            
    except Exception as e:
        logger.error(f"❌ Migration failed: {e}")
        sys.exit(1)