import sqlite3

def migrate_sales_schema():
    conn = sqlite3.connect('data/hiffi_catalog.db')
    cur = conn.cursor()
    
    # 1. Rename provenance to provenance_flag
    try:
        cur.execute("ALTER TABLE commercial_sales RENAME COLUMN provenance TO provenance_flag")
        print("✅ Renamed 'provenance' to 'provenance_flag'")
    except sqlite3.OperationalError as e:
        print(f"⚠️ Column rename skipped (likely already renamed): {e}")

    # 2. Add missing granular sales columns
    new_columns = [
        ("pure_album_sales", "REAL DEFAULT 0"),
        ("physical_singles_sales", "REAL DEFAULT 0"),
        ("digital_singles_sales", "REAL DEFAULT 0"),
        ("streaming_eas", "REAL DEFAULT 0")
    ]
    
    for col_name, col_type in new_columns:
        try:
            cur.execute(f"ALTER TABLE commercial_sales ADD COLUMN {col_name} {col_type}")
            print(f"✅ Added column: {col_name}")
        except sqlite3.OperationalError:
            print(f"⚠️ Column {col_name} already exists.")

    conn.commit()

    # 3. Print the current state of the data
    print("\n--- Migration Complete. Current Data State ---")
    cur.execute("SELECT provenance_flag, COUNT(*) FROM commercial_sales GROUP BY provenance_flag")
    for flag, count in cur.fetchall():
        print(f" - {flag}: {count} records")

    cur.execute("SELECT COUNT(*) FROM commercial_sales")
    total = cur.fetchone()[0]
    print(f"\nTotal Artists with Sales Records: {total} / 530 expected")
    
    conn.close()

if __name__ == "__main__":
    migrate_sales_schema()