import sqlite3
import os

# Correct Path for Docker Container
DB_PATH = "/app/companies_v3_fixed_2.db"

def migrate():
    if not os.path.exists(DB_PATH):
        print(f"❌ Database not found at {DB_PATH}. Please ensure the volume is mounted correctly.")
        # Fallback for local testing (optional)
        if os.path.exists("companies_v3_fixed_2.db"):
            print("⚠️  Found DB in current directory, using that instead.")
            db_to_use = "companies_v3_fixed_2.db"
        else:
            return
    else:
        db_to_use = DB_PATH

    print(f"Migrating database at {db_to_use}...")
    conn = sqlite3.connect(db_to_use)
    cursor = conn.cursor()

    columns_to_add = [
        # Industries (Existing List)
        ("industries", "pains", "TEXT"),
        ("industries", "gains", "TEXT"),
        ("industries", "notes", "TEXT"),
        ("industries", "priority", "TEXT"),
        ("industries", "ops_focus_secondary", "BOOLEAN DEFAULT 0"),
        ("industries", "secondary_category_id", "INTEGER"),
        
        # Companies (New List for CRM Data)
        ("companies", "crm_name", "TEXT"),
        ("companies", "crm_website", "TEXT"),
        ("companies", "crm_address", "TEXT"),
        ("companies", "crm_vat", "TEXT"),
        
        # Companies (Status & Quality)
        ("companies", "confidence_score", "FLOAT DEFAULT 0.0"),
        ("companies", "data_mismatch_score", "FLOAT DEFAULT 0.0"),
        ("companies", "website_scrape_status", "TEXT DEFAULT 'PENDING'"),
        ("companies", "wiki_search_status", "TEXT DEFAULT 'PENDING'"),
    ]

    for table, col_name, col_type in columns_to_add:
        try:
            # Check if column exists first to avoid error log spam
            cursor.execute(f"PRAGMA table_info({table})")
            existing_cols = [row[1] for row in cursor.fetchall()]
            
            if col_name in existing_cols:
                print(f"  - Column '{col_name}' already exists in '{table}'.")
            else:
                print(f"  + Adding column '{col_name}' to '{table}'...")
                cursor.execute(f"ALTER TABLE {table} ADD COLUMN {col_name} {col_type}")
                
        except sqlite3.OperationalError as e:
            print(f"❌ Error adding '{col_name}' to '{table}': {e}")
    
    conn.commit()
    conn.close()
    print("✅ Migration complete.")

if __name__ == "__main__":
    migrate()
