core/task_cache.py

"""
core/task_cache.py — deterministic task cache for repeated ZIP+prompt runs.

The cache is intentionally conservative:
- fingerprint = uploaded ZIP bytes hash + normalized prompt + selected model
- default TTL = 24 hours
- stores both the pipeline state and the exported ZIP artifact (when present)
"""

from __future__ import annotations

import copy
import hashlib
import json
import os
import shutil
import time
from pathlib import Path
from typing import Any, Dict, Optional


CACHE_TTL_SECONDS = 24 * 60 * 60


def _normalize_prompt(prompt: str) -> str:
    return " ".join((prompt or "").lower().split())


def _safe_model_name(model_name: Optional[str]) -> str:
    return (model_name or "auto").strip().lower() or "auto"


def task_fingerprint(
    zip_path: Optional[str],
    prompt: str,
    model_name: Optional[str] = None,
) -> str:
    zip_hash = "nozip"
    if zip_path and os.path.exists(zip_path):
        h = hashlib.sha256()
        with open(zip_path, "rb") as fh:
            for chunk in iter(lambda: fh.read(1024 * 1024), b""):
                h.update(chunk)
        zip_hash = h.hexdigest()[:16]
    prompt_norm = _normalize_prompt(prompt)
    prompt_hash = hashlib.md5(prompt_norm.encode("utf-8")).hexdigest()[:12]
    model_hash = hashlib.md5(_safe_model_name(model_name).encode("utf-8")).hexdigest()[:8]
    return f"{zip_hash}-{prompt_hash}-{model_hash}"


class TaskCache:
    def __init__(self, workspace: str, ttl_seconds: int = CACHE_TTL_SECONDS) -> None:
        self.root = Path(workspace).resolve()
        self.entries_dir = self.root / "entries"
        self.artifacts_dir = self.root / "artifacts"
        self.ttl_seconds = int(ttl_seconds)
        self.entries_dir.mkdir(parents=True, exist_ok=True)
        self.artifacts_dir.mkdir(parents=True, exist_ok=True)

    def _entry_path(self, fingerprint: str) -> Path:
        return self.entries_dir / f"{fingerprint}.json"

    def _artifact_path(self, fingerprint: str, original_name: str = "cached_result.zip") -> Path:
        base = Path(original_name or "cached_result.zip").name
        return self.artifacts_dir / f"{fingerprint}__{base}"

    def get(self, fingerprint: str) -> Optional[Dict[str, Any]]:
        entry_path = self._entry_path(fingerprint)
        if not entry_path.exists():
            return None
        age = time.time() - entry_path.stat().st_mtime
        if age > self.ttl_seconds:
            try:
                entry_path.unlink(missing_ok=True)
            except Exception:
                pass
            return None
        try:
            payload = json.loads(entry_path.read_text(encoding="utf-8"))
        except Exception:
            return None
        if not isinstance(payload, dict):
            return None
        return payload

    def store(
        self,
        fingerprint: str,
        state: Dict[str, Any],
        out_zip_path: Optional[str] = None,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "fingerprint": fingerprint,
            "cached_at": int(time.time()),
            "state": copy.deepcopy(state or {}),
            "out_zip_path": None,
        }
        if out_zip_path and os.path.exists(out_zip_path):
            artifact_path = self._artifact_path(fingerprint, os.path.basename(out_zip_path))
            shutil.copyfile(out_zip_path, artifact_path)
            payload["out_zip_path"] = str(artifact_path)
        self._entry_path(fingerprint).write_text(
            json.dumps(payload, indent=2, ensure_ascii=False, default=str),
            encoding="utf-8",
        )
        return payload

    def restore_result(self, fingerprint: str) -> Optional[tuple[Dict[str, Any], Optional[str]]]:
        payload = self.get(fingerprint)
        if not payload:
            return None
        state = copy.deepcopy(payload.get("state") or {})
        out_zip_path = payload.get("out_zip_path") or None
        if out_zip_path and not os.path.exists(out_zip_path):
            out_zip_path = None
        state.setdefault("cache", {})
        state["cache"].update({
            "hit": True,
            "fingerprint": fingerprint,
            "cached_at": payload.get("cached_at"),
        })
        state.setdefault("timeline", []).append(
            {
                "timestamp": time.strftime("%H:%M:%S"),
                "agent": "TaskCache",
                "status": "OK",
                "message": "cache hit — reused prior pipeline result",
            }
        )
        return state, out_zip_path


__all__ = ["TaskCache", "task_fingerprint", "CACHE_TTL_SECONDS"]