caching.py

from __future__ import annotations

import functools
import hashlib
from collections.abc import Callable
from pathlib import Path
from typing import Any, ParamSpec, TypeVar

from adagraphrag.config.loader import find_project_root
from adagraphrag.utils.io_utils import read_pickle, write_pickle

_P = ParamSpec("_P")
_R = TypeVar("_R")


def _default_cache_dir() -> Path:
    return find_project_root() / "data" / "processed" / ".cache"


def _make_cache_key(func_name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
    raw_key = f"{func_name}|{args!r}|{sorted(kwargs.items())!r}"
    return hashlib.sha256(raw_key.encode("utf-8")).hexdigest()


def disk_cache(cache_dir: Path | None = None) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:

    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
        @functools.wraps(func)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            resolved_cache_dir = cache_dir or _default_cache_dir()
            resolved_cache_dir.mkdir(parents=True, exist_ok=True)

            cache_key = _make_cache_key(func.__qualname__, args, kwargs)
            cache_file = resolved_cache_dir / f"{cache_key}.pkl"

            if cache_file.is_file():
                return read_pickle(cache_file)  # type: ignore[no-any-return]

            result = func(*args, **kwargs)
            write_pickle(result, cache_file)
            return result

        return wrapper

    return decorator


def make_cache_key(*parts: Any) -> str:
    """Build a stable cache key from plain values (no object identity/repr)."""
    raw_key = "|".join(repr(part) for part in parts)
    return hashlib.sha256(raw_key.encode("utf-8")).hexdigest()


def cache_lookup(cache_dir: Path, key: str) -> tuple[bool, Any]:
    """Return (hit, value). Never raises -- a missing/corrupt entry is a miss."""
    cache_file = cache_dir / f"{key}.pkl"
    if not cache_file.is_file():
        return False, None
    try:
        return True, read_pickle(cache_file)
    except Exception:  # noqa: BLE001 - corrupt cache entry should never break a run
        return False, None


def cache_store(cache_dir: Path, key: str, value: Any) -> None:
    cache_dir.mkdir(parents=True, exist_ok=True)
    write_pickle(value, cache_dir / f"{key}.pkl")


def clear_cache(cache_dir: Path | None = None) -> int:
    resolved_cache_dir = cache_dir or _default_cache_dir()
    if not resolved_cache_dir.is_dir():
        return 0

    deleted_count = 0
    for cache_file in resolved_cache_dir.rglob("*.pkl"):
        cache_file.unlink()
        deleted_count += 1
    return deleted_count