from argparse import ArgumentParser
import os
import shutil

def clear_cache():
    parser = ArgumentParser()
    parser.add_argument("--cache_dir", type=str, required=True,
                        help="Relative path under finetune_model, like 'finetune_model/early'")
    args = parser.parse_args()
    target_subpath = args.cache_dir.strip("/")

    if target_subpath == "all":
        print("Removing entire 'cache', 'gif', and 'finetune_model' folders.")
        shutil.rmtree("cache", ignore_errors=True)
        shutil.rmtree("gif", ignore_errors=True)
        shutil.rmtree("finetune_model", ignore_errors=True)
        return

    paths_to_remove = [
        os.path.join("cache/evaluations", target_subpath),
        os.path.join("cache/grads", target_subpath),
        os.path.join("gif", target_subpath),
        target_subpath,  # The actual target directory itself (e.g., finetune_model/early)
    ]

    for path in paths_to_remove:
        print(f"Attempting to remove: {path}")
        shutil.rmtree(path, ignore_errors=True)

if __name__ == "__main__":
    clear_cache()
