# Copyright (c) Metis. All rights reserved.

"""Algorithm registry with auto-discovery via __subclasses__().

Discovers all Algorithm subclasses by importing packages under mantisdk.algorithm,
then walking __subclasses__() recursively. Algorithms with a non-empty ``name``
class attribute are registered.

Usage::

    from mantisdk.platform.registry import discover_algorithms, get_algorithm_by_name

    registry = discover_algorithms()
    # {'gepa': <class GEPA>, 'apo': <class APO>, ...}

    algo_cls = get_algorithm_by_name("gepa")
    algo = algo_cls()
    result = algo.verify({"population_size": 5})
"""

from __future__ import annotations

import importlib
import logging
import pkgutil
from typing import Any, Dict, List, Optional, Type

from mantisdk.algorithm.base import Algorithm

logger = logging.getLogger(__name__)

_registry: Dict[str, Type[Algorithm]] = {}
_discovered = False


def discover_algorithms() -> Dict[str, Type[Algorithm]]:
    """Auto-discover all Algorithm subclasses.

    Imports all modules under mantisdk.algorithm to trigger class registration,
    then walks __subclasses__() recursively to find all algorithms with a name.

    Returns:
        Dict mapping algorithm name to Algorithm class.
    """
    global _discovered

    if _discovered:
        return _registry

    _import_algorithm_packages()

    def _walk(cls: Type[Algorithm]) -> None:
        for sub in cls.__subclasses__():
            if sub.name:
                _registry[sub.name] = sub
                logger.debug(f"Discovered algorithm: {sub.name} " f"({sub.__module__}.{sub.__qualname__})")
            _walk(sub)

    _walk(Algorithm)
    _discovered = True

    logger.info(f"Discovered {len(_registry)} algorithms: {list(_registry.keys())}")
    return _registry


def get_algorithm_by_name(name: str) -> Optional[Type[Algorithm]]:
    """Look up an algorithm class by its name attribute.

    Args:
        name: Algorithm name (e.g., "gepa", "apo").

    Returns:
        Algorithm class or None if not found.
    """
    if not _discovered:
        discover_algorithms()
    return _registry.get(name)


def list_algorithms() -> List[Dict[str, Any]]:
    """List all discovered algorithms with metadata.

    Returns:
        List of dicts with name, display_name, description, requirements.
    """
    if not _discovered:
        discover_algorithms()

    return [
        {
            "name": cls.name,
            "display_name": cls.display_name,
            "description": cls.description,
            "requirements": cls.requirements,
            "config_schema": cls.config_schema,
            "manages_containers": cls.manages_containers,
        }
        for cls in _registry.values()
    ]


def _import_algorithm_packages() -> None:
    """Import all modules under mantisdk.algorithm to trigger class registration.

    Uses try/except per module to tolerate missing optional dependencies
    (e.g., verl requires PyTorch + GPU).
    """
    try:
        import mantisdk.algorithm as algo_pkg

        for _importer, modname, _ispkg in pkgutil.walk_packages(
            algo_pkg.__path__,
            prefix="mantisdk.algorithm.",
        ):
            try:
                importlib.import_module(modname)
            except ImportError as e:
                logger.debug(f"Optional dependency not available for {modname}: {e}")
            except Exception as e:
                logger.warning(f"Failed to import algorithm module {modname}: {e}")
    except Exception as e:
        logger.warning(f"Could not walk mantisdk.algorithm: {e}")
