from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Dict


@dataclass
class LocalProblem:
    asset_id: str
    payload: Dict[str, Any]


@dataclass
class SharedVariables:
    iter_id: int
    values: Dict[str, Any]


@dataclass
class DualVariables:
    iter_id: int
    values: Dict[str, Any]


@dataclass
class PlanDelta:
    iter_id: int
    delta: Dict[str, Any]


@dataclass
class PrivacyBudget:
    budget_id: str
    limits: Dict[str, Any]


@dataclass
class AuditLog:
    entry_id: str
    payload: Dict[str, Any]


@dataclass
class ContractDefinition:
    name: str
    version: str
    schema: Dict[str, Any]


class ContractRegistry:
    """Lightweight, in-memory, versioned contract registry.

    - Contracts are registered by name and version.
    - Each contract has a schema (dict) describing LocalProblem/SharedVariables/etc.
    - Exposes simple get_contract(name, version) and add_contract(name, version, schema).
    """

    def __init__(self) -> None:
        self._store: Dict[str, Dict[str, ContractDefinition]] = {}

    def add_contract(self, name: str, version: str, schema: Dict[str, Any]) -> None:
        self._store.setdefault(name, {})[version] = ContractDefinition(name, version, schema)

    def get_contract(self, name: str, version: str) -> ContractDefinition | None:
        return self._store.get(name, {}).get(version)

    def list_contracts(self) -> Dict[str, Dict[str, ContractDefinition]]:
        return self._store