"""Memory-mapped frozen semantic features used by Dense UV Parser."""

import json
from pathlib import Path

import numpy as np
import torch


SIGLIP_CACHE_VERSION = 2
SUPPORTED_SIGLIP_CACHE_VERSIONS = (1, 2)


class SigLIPGlobalCache:
    def __init__(
        self,
        cache_dir,
        expected_views=None,
        expected_model=None,
        expected_data_dir=None,
        require_spatial=False,
    ):
        self.cache_dir = Path(cache_dir)
        metadata_path = self.cache_dir / "metadata.json"
        embeddings_path = self.cache_dir / "embeddings.npy"
        if not metadata_path.is_file() or not embeddings_path.is_file():
            raise FileNotFoundError(
                f"Incomplete SigLIP cache in {self.cache_dir}; expected metadata.json "
                "and embeddings.npy."
            )
        self.metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
        if self.metadata.get("version") not in SUPPORTED_SIGLIP_CACHE_VERSIONS:
            raise ValueError(
                f"Unsupported SigLIP cache version {self.metadata.get('version')!r}."
            )
        if expected_views is not None and self.metadata.get("views") != list(
            expected_views
        ):
            raise ValueError(
                f"SigLIP cache views={self.metadata.get('views')!r}, "
                f"expected {list(expected_views)!r}."
            )
        if (
            expected_model is not None
            and self.metadata.get("siglip_model") != expected_model
        ):
            raise ValueError(
                f"SigLIP cache model={self.metadata.get('siglip_model')!r}, "
                f"expected {expected_model!r}."
            )
        if expected_data_dir is not None:
            cached_data_dir = self.metadata.get("data_dir")
            requested_data_dir = Path(expected_data_dir).resolve()
            if cached_data_dir and Path(cached_data_dir).resolve() != requested_data_dir:
                raise ValueError(
                    f"SigLIP cache data_dir={str(Path(cached_data_dir).resolve())!r}, "
                    f"expected {str(requested_data_dir)!r}."
                )

        filenames = self.metadata.get("filenames")
        if not isinstance(filenames, list) or len(set(filenames)) != len(filenames):
            raise ValueError("SigLIP cache filenames must be a unique list.")
        self.filename_to_index = {
            filename: index for index, filename in enumerate(filenames)
        }
        self.embeddings = np.load(embeddings_path, mmap_mode="r")
        expected_shape = (
            len(filenames),
            len(self.metadata["views"]),
            int(self.metadata["feature_dim"]),
        )
        if self.embeddings.shape != expected_shape:
            raise ValueError(
                f"SigLIP cache shape={self.embeddings.shape}, expected {expected_shape}."
            )

        spatial_metadata = (
            self.metadata.get("spatial_feature_dim"),
            self.metadata.get("spatial_height"),
            self.metadata.get("spatial_width"),
        )
        self.has_spatial = all(value is not None for value in spatial_metadata)
        self.spatial_embeddings = None
        if self.has_spatial:
            spatial_path = self.cache_dir / "spatial_embeddings.npy"
            if not spatial_path.is_file():
                raise FileNotFoundError(
                    "SigLIP cache metadata declares spatial features but "
                    f"{spatial_path} is missing."
                )
            self.spatial_embeddings = np.load(spatial_path, mmap_mode="r")
            expected_spatial_shape = (
                len(filenames),
                len(self.metadata["views"]),
                int(self.metadata["spatial_feature_dim"]),
                int(self.metadata["spatial_height"]),
                int(self.metadata["spatial_width"]),
            )
            if self.spatial_embeddings.shape != expected_spatial_shape:
                raise ValueError(
                    "SigLIP spatial cache shape="
                    f"{self.spatial_embeddings.shape}, expected "
                    f"{expected_spatial_shape}."
                )
        elif any(value is not None for value in spatial_metadata):
            raise ValueError("SigLIP spatial cache metadata is incomplete.")
        if require_spatial and not self.has_spatial:
            raise ValueError(
                f"SigLIP cache {self.cache_dir} contains only global features; "
                "rebuild it with spatial caching enabled."
            )

    def get(self, filename):
        try:
            index = self.filename_to_index[filename]
        except KeyError as error:
            raise KeyError(f"{filename} is missing from the SigLIP cache.") from error
        return torch.from_numpy(
            np.array(self.embeddings[index], dtype=np.float32, copy=True)
        )

    def get_spatial(self, filename):
        if not self.has_spatial:
            raise ValueError("This SigLIP cache has no spatial features.")
        try:
            index = self.filename_to_index[filename]
        except KeyError as error:
            raise KeyError(f"{filename} is missing from the SigLIP cache.") from error
        return torch.from_numpy(
            np.array(
                self.spatial_embeddings[index],
                dtype=np.float16,
                copy=True,
            )
        )