"""Model metadata loaded from exported JSON."""

import json
from dataclasses import dataclass, fields
from pathlib import Path

from sam2_onnx.types import MetadataDict


@dataclass(frozen=True)
class ModelMetadata:
    """Dimensions and sizes from the exported SAM 2 model."""

    image_size: int  # e.g. 1024
    embed_dim: int  # e.g. 256
    backbone_stride: int  # e.g. 16
    num_multimask_outputs: int  # e.g. 3
    mask_threshold: float  # e.g. 0.0
    quantization: str | None = None  # None, "uint8", "int8" (for tolerance hints)
    sha256: str | None = None  # SHA-256 hash for H1 (stored, not verified)

    def __post_init__(self) -> None:
        """Validate metadata fields."""
        if (
            not isinstance(self.image_size, int)
            or self.image_size <= 0
            or self.image_size > 16384
        ):
            raise ValueError(
                f"image_size must be int in (0, 16384], got {self.image_size}"
            )
        if not isinstance(self.embed_dim, int) or self.embed_dim <= 0:
            raise ValueError(f"embed_dim must be int > 0, got {self.embed_dim}")
        if not isinstance(self.backbone_stride, int) or self.backbone_stride <= 0:
            raise ValueError(
                f"backbone_stride must be int > 0, got {self.backbone_stride}"
            )
        if (
            not isinstance(self.num_multimask_outputs, int)
            or self.num_multimask_outputs <= 0
        ):
            raise ValueError(
                "num_multimask_outputs must be int > 0, "
                f"got {self.num_multimask_outputs}"
            )
        if not isinstance(self.mask_threshold, (int, float)):
            raise ValueError(
                "mask_threshold must be numeric, "
                f"got {type(self.mask_threshold).__name__}"
            )
        if self.quantization not in {None, "uint8", "int8"}:
            raise ValueError(
                "quantization must be None, 'uint8', or 'int8', "
                f"got {self.quantization!r}"
            )

    @property
    def image_embedding_size(self) -> int:
        return self.image_size // self.backbone_stride

    @property
    def mask_input_size(self) -> int:
        return 4 * self.image_embedding_size

    @classmethod
    def from_json(cls, path: str | Path) -> "ModelMetadata":
        with open(path) as f:
            data = json.load(f)
        if not isinstance(data, dict):
            raise TypeError(f"Expected JSON object, got {type(data).__name__}")
        # Filter unknown keys
        known = {f.name for f in fields(cls)}
        data = {k: v for k, v in data.items() if k in known}
        return cls(**data)

    def to_json(self, path: str | Path) -> None:
        with open(path, "w") as f:
            json.dump(self.__dict__, f, indent=2)  # pragma: no mutate

    def to_dict(self) -> MetadataDict:
        """Convert to dictionary for JSON serialization.

        Always includes all fields (matching to_json behavior).
        """
        return {
            "image_size": self.image_size,
            "embed_dim": self.embed_dim,
            "backbone_stride": self.backbone_stride,
            "num_multimask_outputs": self.num_multimask_outputs,
            "mask_threshold": self.mask_threshold,
            "quantization": self.quantization,
            "sha256": self.sha256,
        }
