"""Run cleanup excluding protected studies."""
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

from optimization_engine.utils.study_cleanup import cleanup_study, get_study_info

m1_dir = Path(r"C:\Users\antoi\Atomizer\studies\M1_Mirror")

# Studies to SKIP (user requested)
skip_patterns = [
    "cost_reduction_V10",
    "cost_reduction_V11",
    "cost_reduction_V12",
    "flat_back",
]

# Parse args
dry_run = "--execute" not in sys.argv
keep_best = 5

total_saved = 0
studies_to_clean = []

print("=" * 75)
print(f"CLEANUP (excluding V10-V12 and flat_back studies)")
print(f"Mode: {'DRY RUN' if dry_run else 'EXECUTE'}")
print("=" * 75)
print(f"{'Study':<45} {'Trials':>7} {'Size':>8} {'Savings':>8}")
print("-" * 75)

for study_path in sorted(m1_dir.iterdir()):
    if not study_path.is_dir():
        continue
    # Check if has iterations
    if not (study_path / "2_iterations").exists():
        continue

    # Skip protected studies
    skip = False
    for pattern in skip_patterns:
        if pattern in study_path.name:
            skip = True
            break

    if skip:
        info = get_study_info(study_path)
        print(f"{study_path.name:<45} {info['trial_count']:>7}   SKIPPED")
        continue

    # This study will be cleaned
    result = cleanup_study(study_path, dry_run=dry_run, keep_best=keep_best)
    saved = result["space_saved_gb"]
    total_saved += saved
    status = "would save" if dry_run else "saved"
    print(f"{study_path.name:<45} {result['trial_count']:>7} {result['total_size_before']/(1024**3):>7.1f}G {saved:>7.1f}G")
    studies_to_clean.append(study_path.name)

print("-" * 75)
print(f"{'TOTAL SAVINGS:':<45} {' '*15} {total_saved:>7.1f}G")

if dry_run:
    print(f"\n[!] This was a dry run. Run with --execute to actually delete files.")
else:
    print(f"\n[OK] Cleanup complete! Freed {total_saved:.1f} GB")