scripts/populate_org_codes.py

"""
Populate SwiftOps codes for existing organizations
Run this after applying the migration: 20250128000000_add_swiftops_code.sql

Usage:
    python scripts/populate_org_codes.py
"""
import sys
import os
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from app.models.client import Client
from app.models.contractor import Contractor
from app.utils.org_code_generator import generate_org_code


def get_database_url():
    """Get database URL from environment"""
    database_url = os.getenv("DATABASE_URL")
    if not database_url:
        raise ValueError("DATABASE_URL environment variable not set")
    return database_url


def populate_codes():
    """Populate swiftops_code for all existing organizations"""
    # Create database connection
    database_url = get_database_url()
    engine = create_engine(database_url)
    SessionLocal = sessionmaker(bind=engine)
    db = SessionLocal()
    
    try:
        print("Starting organization code population...")
        print("-" * 60)
        
        # Process clients
        clients = db.query(Client).filter(Client.swiftops_code.is_(None)).all()
        print(f"\nFound {len(clients)} clients without codes")
        
        for client in clients:
            try:
                code = generate_org_code(client.name, db)
                client.swiftops_code = code
                db.flush()  # Flush to make code available for uniqueness checks
                print(f"✓ Client: {client.name[:40]:<40} -> {code}")
            except Exception as e:
                print(f"✗ Failed for client '{client.name}': {e}")
                db.rollback()
                continue
        
        # Commit all client codes
        db.commit()
        print(f"\n✓ Successfully generated codes for {len(clients)} clients")
        
        # Process contractors
        contractors = db.query(Contractor).filter(Contractor.swiftops_code.is_(None)).all()
        print(f"\nFound {len(contractors)} contractors without codes")
        
        for contractor in contractors:
            try:
                code = generate_org_code(contractor.name, db)
                contractor.swiftops_code = code
                db.flush()  # Flush to make code available for uniqueness checks
                print(f"✓ Contractor: {contractor.name[:40]:<40} -> {code}")
            except Exception as e:
                print(f"✗ Failed for contractor '{contractor.name}': {e}")
                db.rollback()
                continue
        
        # Commit all contractor codes
        db.commit()
        print(f"\n✓ Successfully generated codes for {len(contractors)} contractors")
        
        print("-" * 60)
        print("\n✅ Code population completed successfully!")
        print("\nNext steps:")
        print("1. Verify all organizations have codes:")
        print("   SELECT name, swiftops_code FROM clients WHERE swiftops_code IS NULL;")
        print("   SELECT name, swiftops_code FROM contractors WHERE swiftops_code IS NULL;")
        print("\n2. If all look good, make the field NOT NULL:")
        print("   ALTER TABLE clients ALTER COLUMN swiftops_code SET NOT NULL;")
        print("   ALTER TABLE contractors ALTER COLUMN swiftops_code SET NOT NULL;")
        
    except Exception as e:
        print(f"\n❌ Error during population: {e}")
        db.rollback()
        raise
    finally:
        db.close()


if __name__ == "__main__":
    populate_codes()