import os
import sys
import psycopg2
from pathlib import Path

def get_db_url():
    """Get database URL from environment variables."""
    db_user = os.getenv('DB_USER', 'mold_user')
    db_password = os.getenv('DB_PASSWORD', 'mold_password')
    db_host = os.getenv('DB_HOST', 'localhost')
    db_port = os.getenv('DB_PORT', '5434')
    db_name = os.getenv('DB_NAME', 'mold_cost_calculator_dev')
    
    return f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"

def run_migration():
    # Get database URL from environment
    db_url = get_db_url()
    
    try:
        # Connect to the database
        conn = psycopg2.connect(db_url)
        conn.autocommit = True
        cur = conn.cursor()
        
        # Read and execute the migration SQL
        migration_file = Path(__file__).parent / 'update_mold_types.sql'
        with open(migration_file, 'r') as f:
            sql = f.read()
            cur.execute(sql)
        
        print("Migration completed successfully!")
        
    except Exception as e:
        print(f"Error running migration: {str(e)}")
        sys.exit(1)
    finally:
        if 'cur' in locals():
            cur.close()
        if 'conn' in locals():
            conn.close()

if __name__ == '__main__':
    run_migration()