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

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

from __future__ import annotations

from typing import TYPE_CHECKING, Any, TypeVar

from engine.pipeline.core import Stage

if TYPE_CHECKING:
    from engine.pipeline.core import Stage

T = TypeVar("T")


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

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

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


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,
            PoetryDataSource,
        )

        StageRegistry.register("source", HeadlinesDataSource)
        StageRegistry.register("source", PoetryDataSource)

        StageRegistry._categories["source"]["headlines"] = HeadlinesDataSource
        StageRegistry._categories["source"]["poetry"] = PoetryDataSource
    except ImportError:
        pass

    # Register pipeline introspection source
    try:
        from engine.pipeline_sources.pipeline_introspection import (
            PipelineIntrospectionSource,
        )

        StageRegistry.register("source", PipelineIntrospectionSource)
        StageRegistry._categories["source"]["pipeline-inspect"] = (
            PipelineIntrospectionSource
        )
    except ImportError:
        pass

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

    # Register display stages
    _register_display_stages()

    StageRegistry._discovered = True


def _register_display_stages() -> None:
    """Register display backends as stages."""
    try:
        from engine.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 engine.display import DisplayRegistry
        from engine.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)