import hashlib
import json
from collections.abc import Mapping, Sequence
from datetime import UTC, datetime
from typing import ClassVar

from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator


def canonical_hash(payload: object) -> str:
    canonical = json.dumps(
        payload,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")
    return f"sha256:{hashlib.sha256(canonical).hexdigest()}"


class FrozenModel(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(frozen=True, extra="forbid")


class Example(FrozenModel):
    id: str
    inputs: dict[str, str]
    expected: str


# Nested mappings remain mutable. Validation checks the stored hash;
# serialization and live mutation do not run that check.
class DatasetVersion(FrozenModel):
    dataset_id: str
    version: str
    provenance: Mapping[str, str]
    examples: tuple[Example, ...]
    created_at: AwareDatetime = Field(default_factory=lambda: datetime.now(UTC))

    @staticmethod
    def compute_version(
        provenance: Mapping[str, str],
        examples: Sequence[Example],
    ) -> str:
        return canonical_hash(
            {
                "provenance": dict(provenance),
                "examples": [e.model_dump(mode="json") for e in examples],
            }
        )

    @classmethod
    def create(
        cls,
        *,
        dataset_id: str,
        provenance: Mapping[str, str],
        examples: Sequence[Example],
    ) -> DatasetVersion:
        return cls(
            dataset_id=dataset_id,
            version=cls.compute_version(provenance, examples),
            provenance=provenance,
            examples=examples,
        )

    @model_validator(mode="after")
    def check_version(self) -> DatasetVersion:
        expected = self.compute_version(self.provenance, self.examples)
        if self.version != expected:
            raise ValueError(
                f"version mismatch: expected {expected}, got {self.version}"
            )
        return self