"""
hegm.backends -- Backend registry and abstract interface.

Each supported ML framework (PyTorch, TensorFlow, ...) is represented by a
concrete ``AbstractBackend`` subclass.  Backends register themselves at
import time via :func:`register`.
"""

import abc
from typing import Any, Dict, List, Optional

# ---------------------------------------------------------------------------
# Abstract base
# ---------------------------------------------------------------------------


class AbstractBackend(abc.ABC):
    """Contract that every framework backend must fulfil."""

    @property
    @abc.abstractmethod
    def name(self) -> str:
        """Human-readable backend name, e.g. ``'pytorch'``."""

    @property
    @abc.abstractmethod
    def target_module(self) -> str:
        """Top-level module to intercept, e.g. ``'torch'``."""

    @abc.abstractmethod
    def apply_patches(self, module: Any) -> None:
        """Monkey-patch the framework after it has been imported."""

    @abc.abstractmethod
    def save_checkpoint(self, path: str) -> None:
        """Serialize all tracked state to *path*."""

    @abc.abstractmethod
    def load_checkpoint(self, path: str) -> None:
        """Restore tracked state from *path*."""

    @abc.abstractmethod
    def release_gpu(self) -> None:
        """Best-effort GPU teardown before process exit."""

    @abc.abstractmethod
    def global_step(self, optimizer: Any = None) -> int:
        """Return the current global training step."""

    @abc.abstractmethod
    def is_active(self) -> bool:
        """Return ``True`` if this backend's framework has been imported."""


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

_registry: List[AbstractBackend] = []


def register(backend: AbstractBackend) -> None:
    """Add *backend* to the global registry (called at import time)."""
    _registry.append(backend)


def get_backends() -> List[AbstractBackend]:
    """Return all registered backends."""
    return list(_registry)


def get_active() -> Optional[AbstractBackend]:
    """Return the first backend whose framework has been imported."""
    for b in _registry:
        if b.is_active():
            return b
    return None
