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


@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:
        # In-memory, versioned registry with simple thread-safety for MVP
        self._store: Dict[str, Dict[str, ContractDefinition]] = {}
        self._lock = threading.RLock()

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

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

    def list_contracts(self) -> Dict[str, Dict[str, ContractDefinition]]:
        with self._lock:
            # Return a shallow copy to avoid external mutation
            return {k: dict(v) for k, v in self._store.items()}

    # Convenience helpers for MVP tooling
    def has_contract(self, name: str, version: str) -> bool:
        with self._lock:
            return version in self._store.get(name, {})

    def remove_contract(self, name: str, version: str) -> bool:
        with self._lock:
            if name in self._store and version in self._store[name]:
                del self._store[name][version]
                # Clean up empty dicts
                if not self._store[name]:
                    del self._store[name]
                return True
            return False