import sqlite3
import os

# DB Path relative to this script
# script is in company-explorer/backend/scripts/
# DB is in company-explorer/companies_v3_fixed_2.db
DB_PATH = os.path.join(os.path.dirname(__file__), "../../companies_v3_fixed_2.db")

def add_column(cursor, table, column, type_def):
    try:
        print(f"Adding column {column} to {table}...")
        cursor.execute(f"ALTER TABLE {table} ADD COLUMN {column} {type_def}")
        print("Done.")
    except sqlite3.OperationalError as e:
        if "duplicate column name" in str(e):
            print(f"Column {column} already exists.")
        else:
            print(f"Error: {e}")

def migrate():
    if not os.path.exists(DB_PATH):
        print(f"Database not found at {DB_PATH}")
        return

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    # Industries columns
    add_column(cursor, "industries", "notion_id", "TEXT")
    add_column(cursor, "industries", "industry_group", "TEXT")
    add_column(cursor, "industries", "status_notion", "TEXT")
    add_column(cursor, "industries", "whale_threshold", "REAL")
    add_column(cursor, "industries", "min_requirement", "REAL")
    add_column(cursor, "industries", "scraper_keywords", "TEXT")
    add_column(cursor, "industries", "core_unit", "TEXT")
    add_column(cursor, "industries", "proxy_factor", "TEXT")

    # Robotics Categories columns
    add_column(cursor, "robotics_categories", "notion_id", "TEXT")
    
    # Enrichment Data columns (just in case)
    add_column(cursor, "enrichment_data", "wiki_verified_empty", "BOOLEAN DEFAULT 0")

    conn.commit()
    conn.close()
    print("Migration completed.")

if __name__ == "__main__":
    migrate()