import os
import shutil

ROOT_DIR = r"C:\Users\LENOVO\Desktop\senti_ai"

paths_to_delete = [
    # 1. notebooks checkpoints (redundant training logs/weights)
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-600"),
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-1200"),
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-1800"),
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-2400"),
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-3000"),
    os.path.join(ROOT_DIR, "notebooks", "output", "senti_shujaa_v3", "checkpoint-3600"),
    
    # 2. adapter checkpoints and optimizers (redundant training states)
    os.path.join(ROOT_DIR, "senti", "sentillm", "adapters", "senti_shujaa_1b_v2", "checkpoint-1000"),
    os.path.join(ROOT_DIR, "senti", "sentillm", "adapters", "senti_shujaa_1b_v2", "checkpoint-1782"),
    os.path.join(ROOT_DIR, "senti", "sentillm", "adapters", "senti_shujaa_7b", "optimizer.pt"),
    os.path.join(ROOT_DIR, "senti", "sentillm", "adapters", "shujaa_3b", "optimizer.pt"),
    
    # 3. postgres setup zip
    os.path.join(ROOT_DIR, "postgres.zip"),
    
    # 4. legacy BM25 index
    os.path.join(ROOT_DIR, "senti", "knowledge", "bm25_legacy", "index.pkl"),
    
    # 5. Outdated merged LLM weights (v1/v2 merged weights)
    os.path.join(ROOT_DIR, "senti", "sentillm", "adapters", "senti_shujaa_merged"),
]

def clean_paths():
    print("=== STARTING SENTI AI DISK SPACE CLEANUP ===")
    total_reclaimed = 0

    # Clean the explicit paths list
    for path in paths_to_delete:
        if os.path.exists(path):
            try:
                if os.path.isdir(path):
                    # Calculate directory size
                    size = sum(os.path.getsize(os.path.join(dirpath, filename)) for dirpath, _, filenames in os.walk(path) 
                               for filename in filenames)
                    shutil.rmtree(path)
                    print(f"Deleted directory: {path} (Reclaimed: {size / 1024 / 1024:.2f} MB)")
                    total_reclaimed += size
                else:
                    size = os.path.getsize(path)
                    os.remove(path)
                    print(f"Deleted file: {path} (Reclaimed: {size / 1024 / 1024:.2f} MB)")
                    total_reclaimed += size
            except Exception as e:
                print(f"Error deleting {path}: {str(e)}")

    # Clean SEC EDGAR files from sentianalysis/cleaned
    cleaned_dir = os.path.join(ROOT_DIR, "sentianalysis", "cleaned")
    if os.path.exists(cleaned_dir):
        print("\nScanning sentianalysis/cleaned for SEC EDGAR datasets...")
        files_deleted = 0
        for f in os.listdir(cleaned_dir):
            file_path = os.path.join(cleaned_dir, f)
            if os.path.isfile(file_path) and f.endswith(".json"):
                try:
                    # SEC EDGAR files are all larger than 10MB in this directory
                    size = os.path.getsize(file_path)
                    if size > 10 * 1024 * 1024:
                        os.remove(file_path)
                        print(f"Deleted SEC EDGAR file: {f} (Reclaimed: {size / 1024 / 1024:.2f} MB)")
                        total_reclaimed += size
                        files_deleted += 1
                except Exception as e:
                    print(f"Error checking/deleting file {f}: {str(e)}")
        print(f"Deleted {files_deleted} SEC EDGAR JSON files.")

    print(f"\n=== CLEANUP COMPLETED ===")
    print(f"Total space reclaimed: {total_reclaimed / 1024 / 1024 / 1024:.2f} GB")

if __name__ == "__main__":
    clean_paths()