from __future__ import annotations

import argparse
import os
import pickle
from pathlib import Path
from typing import Any


CURRENT_CACHE_VERSIONS = {
    "next_step": 3,
    "all_future": 5,
}


def infer_cache_kind(path: Path) -> str | None:
    name = path.name
    for kind in CURRENT_CACHE_VERSIONS:
        marker = f"_{kind}_cache_"
        if marker in name:
            return kind
    return None


def read_cache_version(path: Path) -> int | None:
    try:
        with path.open("rb") as f:
            payload: Any = pickle.load(f)
    except Exception:
        return None
    if not isinstance(payload, dict):
        return None
    version = payload.get("_cache_version")
    try:
        return int(version)
    except (TypeError, ValueError):
        return None


def should_remove(path: Path, remove_all: bool) -> tuple[bool, str]:
    kind = infer_cache_kind(path)
    if kind is None:
        return False, "not a DeepHealth dataset cache"

    if remove_all:
        return True, "remove all dataset caches"

    version = read_cache_version(path)
    expected = CURRENT_CACHE_VERSIONS[kind]
    if version is None:
        return True, f"{kind} cache is unreadable or missing _cache_version"
    if version != expected:
        return True, f"{kind} cache version {version} != current {expected}"
    return False, f"{kind} cache version {version} is current"


def iter_cache_files(root: Path, recursive: bool):
    pattern = "**/*_cache_*.pkl" if recursive else "*_cache_*.pkl"
    yield from root.glob(pattern)


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Remove obsolete DeepHealth dataset cache files. Defaults to dry-run; "
            "pass --apply to delete."
        )
    )
    parser.add_argument(
        "--data_dir",
        type=Path,
        default=Path("."),
        help="Directory containing dataset cache files, usually the data_prefix directory.",
    )
    parser.add_argument(
        "--recursive",
        action="store_true",
        help="Search recursively under data_dir.",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Remove all recognized DeepHealth dataset caches, including current ones.",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Actually delete files. Without this flag, only prints what would be removed.",
    )
    args = parser.parse_args()

    root = args.data_dir.expanduser().resolve()
    if not root.exists() or not root.is_dir():
        raise SystemExit(f"data_dir is not a directory: {root}")

    files = sorted(p for p in iter_cache_files(root, args.recursive) if p.is_file())
    if not files:
        print(f"No *_cache_*.pkl files found under {root}")
        return

    kept = 0
    removed = 0
    candidates = 0
    for path in files:
        remove, reason = should_remove(path, remove_all=bool(args.all))
        rel = os.path.relpath(path, root)
        if remove:
            candidates += 1
            action = "DELETE" if args.apply else "WOULD_DELETE"
            print(f"{action}\t{rel}\t{reason}")
            if args.apply:
                path.unlink()
                removed += 1
        else:
            kept += 1
            print(f"KEEP\t{rel}\t{reason}")

    if args.apply:
        print(f"Removed {removed} cache file(s); kept {kept}.")
    else:
        print(
            f"Dry run: {candidates} cache file(s) would be removed; kept {kept}. "
            "Re-run with --apply to delete."
        )


if __name__ == "__main__":
    main()