Mainline/sideline/pipeline/registry.py

"""
Stage registry - Unified registration for all pipeline stages.

Provides a single registry for sources, effects, displays, and cameras.
Supports plugin discovery via entry points and explicit registration.
"""

from __future__ import annotations

import importlib
import importlib.metadata
import logging
from typing import TYPE_CHECKING, Any, TypeVar

from sideline.pipeline.core import Stage

if TYPE_CHECKING:
    from sideline.pipeline.core import Stage

T = TypeVar("T")
logger = logging.getLogger(__name__)


class StageRegistry:
    """Unified registry for all pipeline stage types.

    Supports both explicit registration and automatic discovery via entry points.
    Plugins can be registered manually or discovered automatically.
    """

    _categories: dict[str, dict[str, type[Any]]] = {}
    _discovered: bool = False
    _instances: dict[str, Stage] = {}
    _plugins_discovered: bool = False
    _plugin_modules: set[str] = set()  # Track loaded plugin modules

    @classmethod
    def register(cls, category: str, stage_class: type[Any]) -> None:
        """Register a stage class in a category.

        Args:
            category: Category name (source, effect, display, camera)
            stage_class: Stage subclass to register
        """
        if category not in cls._categories:
            cls._categories[category] = {}

        key = getattr(stage_class, "__name__", stage_class.__class__.__name__)
        cls._categories[category][key] = stage_class

    @classmethod
    def get(cls, category: str, name: str) -> type[Any] | None:
        """Get a stage class by category and name."""
        return cls._categories.get(category, {}).get(name)

    @classmethod
    def list(cls, category: str) -> list[str]:
        """List all stage names in a category."""
        return list(cls._categories.get(category, {}).keys())

    @classmethod
    def list_categories(cls) -> list[str]:
        """List all registered categories."""
        return list(cls._categories.keys())

    @classmethod
    def create(cls, category: str, name: str, **kwargs) -> Stage | None:
        """Create a stage instance by category and name."""
        stage_class = cls.get(category, name)
        if stage_class:
            return stage_class(**kwargs)
        return None

    @classmethod
    def create_instance(cls, stage: Stage | type[Stage], **kwargs) -> Stage:
        """Create an instance from a stage class or return as-is."""
        if isinstance(stage, Stage):
            return stage
        if isinstance(stage, type) and issubclass(stage, Stage):
            return stage(**kwargs)
        raise TypeError(f"Expected Stage class or instance, got {type(stage)}")

    @classmethod
    def register_instance(cls, name: str, stage: Stage) -> None:
        """Register a stage instance by name."""
        cls._instances[name] = stage

    @classmethod
    def get_instance(cls, name: str) -> Stage | None:
        """Get a registered stage instance by name."""
        return cls._instances.get(name)

    @classmethod
    def register_plugin_module(cls, plugin_module: str) -> None:
        """Register stages from an external plugin module.

        The module should define a register_stages(registry) function.

        Args:
            plugin_module: Full module path (e.g., 'engine.plugins')
        """
        if plugin_module in cls._plugin_modules:
            logger.debug(f"Plugin module {plugin_module} already loaded")
            return

        try:
            module = importlib.import_module(plugin_module)
            if hasattr(module, "register_stages"):
                module.register_stages(cls)
                cls._plugin_modules.add(plugin_module)
                logger.info(f"Registered stages from {plugin_module}")
            else:
                logger.warning(
                    f"Plugin module {plugin_module} has no register_stages function"
                )
        except ImportError as e:
            logger.warning(f"Failed to import plugin module {plugin_module}: {e}")

    # Backward compatibility alias
    register_plugin = register_plugin_module

    @classmethod
    def discover_plugins(cls) -> None:
        """Auto-discover and register plugins via entry points.

        Looks for 'sideline.stages' entry points in installed packages.
        Each entry point should point to a register_stages(registry) function.
        """
        if cls._plugins_discovered:
            return

        try:
            # Discover entry points for sideline.stages
            # Python 3.12+ changed the entry_points() API
            try:
                entry_points = importlib.metadata.entry_points()
                if hasattr(entry_points, "get"):
                    # Python < 3.12
                    stages_eps = entry_points.get("sideline.stages", [])
                else:
                    # Python 3.12+
                    stages_eps = entry_points.select(group="sideline.stages")
            except Exception:
                # Fallback: try both approaches
                try:
                    entry_points = importlib.metadata.entry_points()
                    stages_eps = entry_points.get("sideline.stages", [])
                except Exception:
                    stages_eps = []

            for ep in stages_eps:
                try:
                    register_func = ep.load()
                    if callable(register_func):
                        register_func(cls)
                        logger.info(f"Discovered and registered plugin: {ep.name}")
                except Exception as e:
                    logger.warning(f"Failed to load entry point {ep.name}: {e}")

            cls._plugins_discovered = True
        except Exception as e:
            logger.warning(f"Failed to discover plugins: {e}")

    @classmethod
    def get_discovered_modules(cls) -> set[str]:
        """Get set of plugin modules that have been loaded."""
        return cls._plugin_modules.copy()


def discover_stages() -> None:
    """Auto-discover and register all stage implementations.

    This function now only registers framework-level stages (displays, etc.).
    Application-specific stages should be registered via plugins.
    """
    if StageRegistry._discovered:
        return

    # Register display stages (framework-level)
    _register_display_stages()

    # Discover plugins via entry points
    StageRegistry.discover_plugins()

    StageRegistry._discovered = True


def _register_display_stages() -> None:
    """Register display backends as stages."""
    try:
        from sideline.display import DisplayRegistry
    except ImportError:
        return

    DisplayRegistry.initialize()

    for backend_name in DisplayRegistry.list_backends():
        factory = _DisplayStageFactory(backend_name)
        StageRegistry._categories.setdefault("display", {})[backend_name] = factory


class _DisplayStageFactory:
    """Factory that creates DisplayStage instances for a specific backend."""

    def __init__(self, backend_name: str):
        self._backend_name = backend_name

    def __call__(self):
        from sideline.display import DisplayRegistry
        from sideline.pipeline.adapters import DisplayStage

        display = DisplayRegistry.create(self._backend_name)
        if display is None:
            raise RuntimeError(
                f"Failed to create display backend: {self._backend_name}"
            )
        return DisplayStage(display, name=self._backend_name)

    @property
    def __name__(self) -> str:
        return self._backend_name.capitalize() + "Stage"


# Convenience functions
def register_source(stage_class: type[Stage]) -> None:
    """Register a source stage."""
    StageRegistry.register("source", stage_class)


def register_effect(stage_class: type[Stage]) -> None:
    """Register an effect stage."""
    StageRegistry.register("effect", stage_class)


def register_display(stage_class: type[Stage]) -> None:
    """Register a display stage."""
    StageRegistry.register("display", stage_class)


def register_camera(stage_class: type[Stage]) -> None:
    """Register a camera stage."""
    StageRegistry.register("camera", stage_class)