import sqlite3
import logging
import os

from logging_utils import setup_logging

def cleanup_market_cap_tables():
    """
    Scans the database and drops all tables related to market cap data
    to allow for a clean refresh.
    """
    setup_logging('normal', 'DBCleanup')
    db_path = os.path.join("_data", "market_data.db")

    if not os.path.exists(db_path):
        logging.error(f"Database file not found at '{db_path}'. Nothing to clean.")
        return

    logging.info(f"Connecting to database at '{db_path}'...")
    try:
        with sqlite3.connect(db_path) as conn:
            cursor = conn.cursor()
            
            # Find all tables that were created by the market cap fetcher
            cursor.execute("""
                SELECT name FROM sqlite_master 
                WHERE type='table' 
                AND (name LIKE '%_market_cap' OR name LIKE 'TOTAL_%')
            """)
            
            tables_to_drop = cursor.fetchall()
            
            if not tables_to_drop:
                logging.info("No market cap tables found to clean up. Database is already clean.")
                return

            logging.warning(f"Found {len(tables_to_drop)} market cap tables to remove...")
            
            for table in tables_to_drop:
                table_name = table[0]
                try:
                    logging.info(f"Dropping table: {table_name}...")
                    conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
                except Exception as e:
                    logging.error(f"Failed to drop table {table_name}: {e}")
            
            conn.commit()
            logging.info("--- Database cleanup complete ---")

    except sqlite3.Error as e:
        logging.error(f"A database error occurred: {e}")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")

if __name__ == "__main__":
    cleanup_market_cap_tables()