sideline/engine/pipeline/adapters/frame_capture.py

"""
Frame Capture Stage Adapter

Wraps pipeline stages to capture frames for animation report generation.
"""

from typing import Any

from engine.display.backends.animation_report import AnimationReportDisplay
from engine.pipeline.core import PipelineContext, Stage


class FrameCaptureStage(Stage):
    """
    Wrapper stage that captures frames before and after a wrapped stage.

    This allows generating animation reports showing how each stage
    transforms the data.
    """

    def __init__(
        self,
        wrapped_stage: Stage,
        display: AnimationReportDisplay,
        name: str | None = None,
    ):
        """
        Initialize frame capture stage.

        Args:
            wrapped_stage: The stage to wrap and capture frames from
            display: The animation report display to send frames to
            name: Optional name for this capture stage
        """
        self._wrapped_stage = wrapped_stage
        self._display = display
        self.name = name or f"capture_{wrapped_stage.name}"
        self.category = wrapped_stage.category
        self.optional = wrapped_stage.optional

        # Capture state
        self._captured_input = False
        self._captured_output = False

    @property
    def stage_type(self) -> str:
        return self._wrapped_stage.stage_type

    @property
    def capabilities(self) -> set[str]:
        return self._wrapped_stage.capabilities

    @property
    def dependencies(self) -> set[str]:
        return self._wrapped_stage.dependencies

    @property
    def inlet_types(self) -> set:
        return self._wrapped_stage.inlet_types

    @property
    def outlet_types(self) -> set:
        return self._wrapped_stage.outlet_types

    def init(self, ctx: PipelineContext) -> bool:
        """Initialize the wrapped stage."""
        return self._wrapped_stage.init(ctx)

    def process(self, data: Any, ctx: PipelineContext) -> Any:
        """
        Process data through wrapped stage and capture frames.

        Args:
            data: Input data (typically a text buffer)
            ctx: Pipeline context

        Returns:
            Output data from wrapped stage
        """
        # Capture input frame (before stage processing)
        if isinstance(data, list) and all(isinstance(line, str) for line in data):
            self._display.start_stage(f"{self._wrapped_stage.name}_input")
            self._display.show(data)
            self._captured_input = True

        # Process through wrapped stage
        result = self._wrapped_stage.process(data, ctx)

        # Capture output frame (after stage processing)
        if isinstance(result, list) and all(isinstance(line, str) for line in result):
            self._display.start_stage(f"{self._wrapped_stage.name}_output")
            self._display.show(result)
            self._captured_output = True

        return result

    def cleanup(self) -> None:
        """Cleanup the wrapped stage."""
        self._wrapped_stage.cleanup()


class FrameCaptureController:
    """
    Controller for managing frame capture across the pipeline.

    This class provides an easy way to enable frame capture for
    specific stages or the entire pipeline.
    """

    def __init__(self, display: AnimationReportDisplay):
        """
        Initialize frame capture controller.

        Args:
            display: The animation report display to use for capture
        """
        self._display = display
        self._captured_stages: list[FrameCaptureStage] = []

    def wrap_stage(self, stage: Stage, name: str | None = None) -> FrameCaptureStage:
        """
        Wrap a stage with frame capture.

        Args:
            stage: The stage to wrap
            name: Optional name for the capture stage

        Returns:
            Wrapped stage that captures frames
        """
        capture_stage = FrameCaptureStage(stage, self._display, name)
        self._captured_stages.append(capture_stage)
        return capture_stage

    def wrap_stages(self, stages: dict[str, Stage]) -> dict[str, Stage]:
        """
        Wrap multiple stages with frame capture.

        Args:
            stages: Dictionary of stage names to stages

        Returns:
            Dictionary of stage names to wrapped stages
        """
        wrapped = {}
        for name, stage in stages.items():
            wrapped[name] = self.wrap_stage(stage, name)
        return wrapped

    def get_captured_stages(self) -> list[FrameCaptureStage]:
        """Get list of all captured stages."""
        return self._captured_stages

    def generate_report(self, title: str = "Pipeline Animation Report") -> str:
        """
        Generate the animation report.

        Args:
            title: Title for the report

        Returns:
            Path to the generated HTML file
        """
        report_path = self._display.generate_report(title)
        return str(report_path)