from __future__ import annotations

"""Disk cache for drone-query embeddings used by DynamicSimilaritySampler.

Caches the [N, D] tensor produced by `_embed_drone_queries` keyed by epoch
number, so `--resume` doesn't need to recompute and ablations that rerun
from a snapshot are reproducible.

Entries are plain `torch.save` files; there's no atomicity beyond what the
OS provides. Safe enough for single-writer usage within the training loop.
"""

import logging
from pathlib import Path

import torch

LOGGER = logging.getLogger("caption_test.emb_cache")


class EmbeddingCache:
    """File-based cache of epoch → [N, D] CPU tensor."""

    def __init__(self, cache_dir: str | Path) -> None:
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def path_for(self, epoch: int) -> Path:
        return self.cache_dir / f"queries_epoch{epoch:03d}.pt"

    def load(self, epoch: int) -> torch.Tensor | None:
        """Return cached embeddings for `epoch` or None if absent."""
        p = self.path_for(epoch)
        if not p.exists():
            return None
        try:
            emb = torch.load(p, map_location="cpu", weights_only=False)
        except Exception as exc:
            LOGGER.warning("Failed to load embedding cache %s: %s", p, exc)
            return None
        LOGGER.info("Loaded cached embeddings from %s (shape=%s)", p.name, tuple(emb.shape))
        return emb

    def save(self, epoch: int, embeddings: torch.Tensor) -> None:
        p = self.path_for(epoch)
        tmp = p.with_suffix(p.suffix + ".tmp")
        torch.save(embeddings.detach().cpu(), tmp)
        tmp.replace(p)
        LOGGER.info("Saved embedding cache to %s (shape=%s)", p.name, tuple(embeddings.shape))
