"""
Helper script to generate referral codes for existing contractors without codes
"""
import secrets
import string
from app import create_app
from app.extensions import db
from app.models.user import User


def generate_referral_code():
    """Generate a unique 8-character referral code"""
    chars = string.ascii_uppercase.replace('O', '').replace('I', '') + string.digits.replace('0', '').replace('1', '')
    
    for _ in range(10):
        code = ''.join(secrets.choice(chars) for _ in range(8))
        if not User.query.filter_by(referral_code=code).first():
            return code
    
    # Fallback to 12 characters if collisions
    return ''.join(secrets.choice(chars) for _ in range(12))


def generate_codes_for_contractors():
    """Generate referral codes for all contractors without codes"""
    app = create_app()
    
    with app.app_context():
        # Find contractors without referral codes
        contractors = User.query.filter_by(
            role='contractor',
            referral_code=None,
            is_deleted=False
        ).all()
        
        if not contractors:
            print("No contractors found without referral codes")
            return
        
        print(f"Found {len(contractors)} contractors without referral codes")
        
        updated = 0
        for contractor in contractors:
            referral_code = generate_referral_code()
            contractor.referral_code = referral_code
            print(f"Generated code {referral_code} for {contractor.email}")
            updated += 1
        
        # Commit all changes
        try:
            db.session.commit()
            print(f"\n✅ Successfully generated referral codes for {updated} contractors")
        except Exception as e:
            db.session.rollback()
            print(f"\n❌ Error: {e}")


if __name__ == '__main__':
    generate_codes_for_contractors()