from sqlalchemy.orm import Session
from db.models import InvestorTable
from db.db import get_db

def update_stage_focus_values():
    """Update existing stage_focus values from lowercase to uppercase"""
    db = next(get_db())
    
    try:
        # Mapping of old lowercase values to new uppercase values
        stage_mappings = {
            'seed': 'SEED',
            'series_a': 'SERIES_A', 
            'series_b': 'SERIES_B',
            'series_c': 'SERIES_C',
            'growth': 'GROWTH',
            'late_stage': 'LATE_STAGE'
        }
        
        updated_count = 0
        
        for old_value, new_value in stage_mappings.items():
            # Update records with the old value
            result = db.query(InvestorTable).filter(
                InvestorTable.stage_focus == old_value
            ).update(
                {InvestorTable.stage_focus: new_value},
                synchronize_session=False
            )
            
            updated_count += result
            print(f"Updated {result} records from '{old_value}' to '{new_value}'")
        
        db.commit()
        print(f"Successfully updated {updated_count} total records")
        
    except Exception as e:
        db.rollback()
        print(f"Error updating stage_focus values: {e}")
        raise
    finally:
        db.close()

# Run the update
if __name__ == "__main__":
    update_stage_focus_values()