"""Versioned cache for processed graph samples."""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import torch

from .graph_dataset import GraphSample
from .metadata import CACHE_SCHEMA_VERSION, FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION


@dataclass(frozen=True)
class CacheMetadata:
    feature_schema_version: int = FEATURE_SCHEMA_VERSION
    graph_schema_version: int = GRAPH_SCHEMA_VERSION
    cache_schema_version: int = CACHE_SCHEMA_VERSION
    source: str | None = None
    preprocessing: dict[str, Any] = field(default_factory=dict)


class GraphSampleCache:
    """Save/load graph samples with schema metadata checked before use."""

    def __init__(self, path: str | Path, metadata: CacheMetadata | None = None) -> None:
        self.path = Path(path)
        self.metadata = metadata or CacheMetadata()

    def save(self, samples: list[GraphSample] | tuple[GraphSample, ...]) -> None:
        payload = {
            "cache_metadata": asdict(self.metadata),
            "samples": tuple(samples),
        }
        self.path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(payload, self.path)

    def load(self) -> tuple[GraphSample, ...]:
        payload = torch.load(self.path, map_location="cpu", weights_only=False)
        if not isinstance(payload, dict):
            raise ValueError("cache payload must be a mapping")
        actual = payload.get("cache_metadata", {})
        if not isinstance(actual, dict):
            raise ValueError("cache metadata is missing or malformed")
        expected = asdict(self.metadata)
        for key in (
            "feature_schema_version",
            "graph_schema_version",
            "cache_schema_version",
        ):
            if actual.get(key) != expected[key]:
                raise ValueError(
                    f"cache schema mismatch for {key}: "
                    f"found {actual.get(key)!r}, expected {expected[key]!r}"
                )
        if actual.get("preprocessing", {}) != expected["preprocessing"]:
            raise ValueError("cache preprocessing fingerprint mismatch")
        if "samples" not in payload:
            raise ValueError("cache samples are missing")
        samples = payload["samples"]
        if not isinstance(samples, (list, tuple)):
            raise ValueError("cache samples are malformed")
        if any(not isinstance(sample, GraphSample) for sample in samples):
            raise ValueError("cache contains invalid graph samples")
        return tuple(samples)

    def exists(self) -> bool:
        return self.path.is_file()