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

Provides a single registry for sources, effects, displays, and cameras.
"""

from __future__ import annotations

from engine.pipeline.core import Stage


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

    _categories: dict[str, dict[str, type[Stage]]] = {}
    _discovered: bool = False
    _instances: dict[str, Stage] = {}

    @classmethod
    def register(cls, category: str, stage_class: type[Stage]) -> 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] = {}

        # Use class name as key
        key = stage_class.__name__
        cls._categories[category][key] = stage_class

    @classmethod
    def get(cls, category: str, name: str) -> type[Stage] | 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)


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

    # Import and register all stage implementations
    try:
        from engine.sources_v2 import (
            HeadlinesDataSource,
            PipelineDataSource,
            PoetryDataSource,
        )

        StageRegistry.register("source", HeadlinesDataSource)
        StageRegistry.register("source", PoetryDataSource)
        StageRegistry.register("source", PipelineDataSource)
    except ImportError:
        pass

    try:
        from engine.effects.types import EffectPlugin  # noqa: F401
    except ImportError:
        pass

    try:
        from engine.display import Display  # noqa: F401
    except ImportError:
        pass

    StageRegistry._discovered = True


# 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)