cache.py

"""Small atomic JSON cache used for answer resumability."""

from __future__ import annotations

import hashlib
import json
import os
import threading
from pathlib import Path
from typing import Any


class AnswerCache:
    def __init__(self, path: Path, enabled: bool = True):
        self.path = path
        self.enabled = enabled
        self._lock = threading.RLock()

    @staticmethod
    def key(task: dict[str, Any], file_path: Path | None = None) -> str:
        digest = hashlib.sha256()
        digest.update(str(task.get("task_id", "")).encode())
        digest.update(b"\0")
        digest.update(str(task.get("question", "")).encode())
        if file_path and file_path.exists():
            digest.update(b"\0")
            with file_path.open("rb") as stream:
                for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                    digest.update(chunk)
        return digest.hexdigest()

    def _read(self) -> dict[str, Any]:
        if not self.enabled or not self.path.exists():
            return {}
        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
            return data if isinstance(data, dict) else {}
        except (OSError, json.JSONDecodeError):
            return {}

    def _write(self, data: dict[str, Any]) -> None:
        if not self.enabled:
            return
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_suffix(self.path.suffix + ".tmp")
        temporary.write_text(
            json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True),
            encoding="utf-8",
        )
        os.replace(temporary, self.path)

    def get(self, key: str) -> dict[str, Any] | None:
        with self._lock:
            value = self._read().get(key)
            return value if isinstance(value, dict) else None

    def put(self, key: str, value: dict[str, Any]) -> None:
        with self._lock:
            data = self._read()
            data[key] = value
            self._write(data)

    def save_run(self, results: list[dict[str, Any]]) -> None:
        with self._lock:
            self._write({"results": results})

    def load_run(self) -> list[dict[str, Any]]:
        with self._lock:
            value = self._read().get("results", [])
            return value if isinstance(value, list) else []


class ResultStore:
    """Atomic task-id-keyed checkpoints for evaluation and reruns."""

    def __init__(self, path: Path):
        self.path = path
        self._cache = AnswerCache(path, enabled=True)

    def load(self) -> dict[str, dict[str, Any]]:
        with self._cache._lock:
            data = self._cache._read()
            results = data.get("results", {})
            return results if isinstance(results, dict) else {}

    def save_result(self, task_id: str, result: dict[str, Any]) -> None:
        with self._cache._lock:
            data = self._cache._read()
            results = data.get("results", {})
            if not isinstance(results, dict):
                results = {}
            results[str(task_id)] = result
            data["results"] = results
            self._cache._write(data)

    def remove(self, task_id: str) -> None:
        with self._cache._lock:
            data = self._cache._read()
            results = data.get("results", {})
            if isinstance(results, dict):
                results.pop(str(task_id), None)
            data["results"] = results
            self._cache._write(data)

    def ordered(self, task_ids: list[str] | None = None) -> list[dict[str, Any]]:
        results = self.load()
        if task_ids is None:
            return list(results.values())
        return [results[task_id] for task_id in task_ids if task_id in results]