from __future__ import annotations

from dataclasses import dataclass
from typing import Any, TypeVar

from pydantic import BaseModel, ValidationError


T = TypeVar("T", bound=BaseModel)


@dataclass(frozen=True)
class SchemaKey:
    name: str
    version: str


class SchemaRegistry:
    def __init__(self) -> None:
        self._schemas: dict[SchemaKey, type[BaseModel]] = {}

    def register(self, name: str, version: str, model: type[T]) -> None:
        self._schemas[SchemaKey(name, version)] = model

    def get(self, name: str, version: str) -> type[BaseModel]:
        key = SchemaKey(name, version)
        try:
            return self._schemas[key]
        except KeyError as exc:
            raise KeyError(f"Unknown schema {name!r} version {version!r}") from exc

    def validate(self, name: str, version: str, payload: Any) -> BaseModel:
        model = self.get(name, version)
        return model.model_validate(payload)

    def versions(self, name: str) -> list[str]:
        return sorted(key.version for key in self._schemas if key.name == name)

    def manifest(self) -> dict[str, list[str]]:
        result: dict[str, list[str]] = {}
        for key in self._schemas:
            result.setdefault(key.name, []).append(key.version)
        return {name: sorted(versions) for name, versions in sorted(result.items())}

    @classmethod
    def with_defaults(cls) -> "SchemaRegistry":
        from .models import Artifact, BuildContract, BuildEnvironment, BuildLog, BuildObject, DependencyGraph, DependencyNode, LedgerEntry, ReproManifest

        registry = cls()
        registry.register("BuildEnvironment", "1.0", BuildEnvironment)
        registry.register("BuildObject", "1.0", BuildObject)
        registry.register("DependencyNode", "1.0", DependencyNode)
        registry.register("DependencyGraph", "1.0", DependencyGraph)
        registry.register("Artifact", "1.0", Artifact)
        registry.register("BuildLog", "1.0", BuildLog)
        registry.register("ReproManifest", "1.0", ReproManifest)
        registry.register("BuildContract", "1.0", BuildContract)
        registry.register("LedgerEntry", "1.0", LedgerEntry)
        return registry