import sqlite3

DB_PATH = "companies_v3_fixed_2.db"

UNIT_MAPPING = {
    "Logistics - Warehouse": "m²",
    "Healthcare - Hospital": "Betten",
    "Infrastructure - Transport": "Passagiere", 
    "Leisure - Indoor Active": "m²",
    "Retail - Food": "m²",
    "Retail - Shopping Center": "m²",
    "Hospitality - Gastronomy": "Sitzplätze",
    "Leisure - Outdoor Park": "Besucher",
    "Leisure - Wet & Spa": "Besucher",
    "Infrastructure - Public": "Kapazität",
    "Retail - Non-Food": "m²",
    "Hospitality - Hotel": "Zimmer",
    "Leisure - Entertainment": "Besucher",
    "Healthcare - Care Home": "Plätze",
    "Industry - Manufacturing": "Mitarbeiter", 
    "Energy - Grid & Utilities": "Kunden",
    "Leisure - Fitness": "Mitglieder",
    "Corporate - Campus": "Mitarbeiter",
    "Energy - Solar/Wind": "MWp", 
    "Tech - Data Center": "Racks", 
    "Automotive - Dealer": "Fahrzeuge",
    "Infrastructure Parking": "Stellplätze",
    "Reinigungsdienstleister": "Mitarbeiter",
    "Infrastructure - Communities": "Einwohner"
}

def fix_units():
    print(f"Connecting to {DB_PATH}...")
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    try:
        cursor.execute("SELECT id, name, scraper_search_term, metric_type FROM industries")
        rows = cursor.fetchall()
        
        updated_count = 0
        
        for row in rows:
            ind_id, name, current_term, m_type = row
            
            new_term = UNIT_MAPPING.get(name)
            
            # Fallback Logic
            if not new_term:
                if m_type in ["AREA_IN", "AREA_OUT"]:
                    new_term = "m²"
                else:
                    new_term = "Anzahl" # Generic fallback
            
            if current_term != new_term:
                print(f"Updating '{name}': '{current_term}' -> '{new_term}'")
                cursor.execute("UPDATE industries SET scraper_search_term = ? WHERE id = ?", (new_term, ind_id))
                updated_count += 1
        
        conn.commit()
        print(f"\n✅ Updated {updated_count} industries with correct units.")
        
    except Exception as e:
        print(f"❌ Error: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    fix_units()