from typing import Dict, Any


# Lightweight default schemas for core MLTrail contracts. This enables a
# canonical, language-agnostic registry that adapters can reference for
# conformance checks and schema validation in MVP deployments.
_DEFAULT_CONTRACT_SCHEMAS: Dict[str, Dict[str, Any]] = {
    "Experiment": {"fields": ["id", "name", "version", "description", "metadata"]},
    "Run": {"fields": ["id", "experiment_id", "parameters", "metrics", "environment_hash"]},
    "Dataset": {"fields": ["id", "name", "version", "metadata"]},
    "Model": {"fields": ["id", "architecture", "fingerprint", "metadata"]},
    "Environment": {"fields": ["id", "language", "version", "dependencies", "container_hash"]},
    "EvaluationMetric": {"fields": ["name", "value", "unit"]},
    "Policy": {"fields": ["id", "rules", "metadata"]},
}

class ContractRegistry:
    def __init__(self) -> None:
        self._contracts: Dict[str, Dict[str, Any]] = {}

    def register_contract(self, name: str, schema: Dict[str, Any], version: str = "1.0.0") -> None:
        self._contracts[name] = {"schema": schema, "version": version}

    def get_contract(self, name: str) -> Dict[str, Any] | None:
        return self._contracts.get(name)

    def all_contracts(self) -> Dict[str, Any]:
        return self._contracts

    def register_default_contracts(self) -> None:
        """Register a canonical set of core MLTrail contract schemas.

        This helps adapters and tooling validate payloads against expected
        fields, enabling consistent interoperability across languages.
        """
        for name, schema in _DEFAULT_CONTRACT_SCHEMAS.items():
            # Use a stable default version for defaults; allow user overrides if needed
            self.register_contract(name, schema, version="1.0.0")