Mainline/scripts/capture_output.py

#!/usr/bin/env python3
"""
Capture output utility for Mainline.

This script captures the output of a Mainline pipeline using NullDisplay
and saves it to a JSON file for comparison with other branches.
"""

import argparse
import json
import time
from pathlib import Path

from engine.display import DisplayRegistry
from engine.pipeline import Pipeline, PipelineConfig, PipelineContext
from engine.pipeline.adapters import create_stage_from_display
from engine.pipeline.presets import get_preset


def capture_pipeline_output(
    preset_name: str,
    output_file: str,
    frames: int = 60,
    width: int = 80,
    height: int = 24,
):
    """Capture pipeline output for a given preset.

    Args:
        preset_name: Name of preset to use
        output_file: Path to save captured output
        frames: Number of frames to capture
        width: Terminal width
        height: Terminal height
    """
    print(f"Capturing output for preset '{preset_name}'...")

    # Get preset
    preset = get_preset(preset_name)
    if not preset:
        print(f"Error: Preset '{preset_name}' not found")
        return False

    # Create NullDisplay with recording
    display = DisplayRegistry.create("null")
    display.init(width, height)
    display.start_recording()

    # Build pipeline
    config = PipelineConfig(
        source=preset.source,
        display="null",  # Use null display
        camera=preset.camera,
        effects=preset.effects,
        enable_metrics=False,
    )

    # Create pipeline context with params
    from engine.pipeline.params import PipelineParams

    params = PipelineParams(
        source=preset.source,
        display="null",
        camera_mode=preset.camera,
        effect_order=preset.effects,
        viewport_width=preset.viewport_width,
        viewport_height=preset.viewport_height,
        camera_speed=preset.camera_speed,
    )

    ctx = PipelineContext()
    ctx.params = params

    pipeline = Pipeline(config=config, context=ctx)

    # Add stages based on preset
    from engine.data_sources.sources import HeadlinesDataSource
    from engine.pipeline.adapters import DataSourceStage

    # Add source stage
    source = HeadlinesDataSource()
    pipeline.add_stage("source", DataSourceStage(source, name="headlines"))

    # Add message overlay if enabled
    if getattr(preset, "enable_message_overlay", False):
        from engine import config as engine_config
        from engine.pipeline.adapters import MessageOverlayConfig, MessageOverlayStage

        overlay_config = MessageOverlayConfig(
            enabled=True,
            display_secs=getattr(engine_config, "MESSAGE_DISPLAY_SECS", 30),
            topic_url=getattr(engine_config, "NTFY_TOPIC", None),
        )
        pipeline.add_stage(
            "message_overlay", MessageOverlayStage(config=overlay_config)
        )

    # Add display stage
    pipeline.add_stage("display", create_stage_from_display(display, "null"))

    # Build and initialize
    pipeline.build()
    if not pipeline.initialize():
        print("Error: Failed to initialize pipeline")
        return False

    # Capture frames
    print(f"Capturing {frames} frames...")
    start_time = time.time()

    for frame in range(frames):
        try:
            pipeline.execute([])
            if frame % 10 == 0:
                print(f"  Frame {frame}/{frames}")
        except Exception as e:
            print(f"Error on frame {frame}: {e}")
            break

    elapsed = time.time() - start_time
    print(f"Captured {frame + 1} frames in {elapsed:.2f}s")

    # Get captured frames
    captured_frames = display.get_frames()
    print(f"Retrieved {len(captured_frames)} frames from display")

    # Save to JSON
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    recording_data = {
        "version": 1,
        "preset": preset_name,
        "display": "null",
        "width": width,
        "height": height,
        "frame_count": len(captured_frames),
        "frames": [
            {
                "frame_number": i,
                "buffer": frame,
                "width": width,
                "height": height,
            }
            for i, frame in enumerate(captured_frames)
        ],
    }

    with open(output_path, "w") as f:
        json.dump(recording_data, f, indent=2)

    print(f"Saved recording to {output_path}")
    return True


def main():
    parser = argparse.ArgumentParser(description="Capture Mainline pipeline output")
    parser.add_argument(
        "--preset",
        default="demo",
        help="Preset name to use (default: demo)",
    )
    parser.add_argument(
        "--output",
        default="output/capture.json",
        help="Output file path (default: output/capture.json)",
    )
    parser.add_argument(
        "--frames",
        type=int,
        default=60,
        help="Number of frames to capture (default: 60)",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=80,
        help="Terminal width (default: 80)",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=24,
        help="Terminal height (default: 24)",
    )

    args = parser.parse_args()

    success = capture_pipeline_output(
        preset_name=args.preset,
        output_file=args.output,
        frames=args.frames,
        width=args.width,
        height=args.height,
    )

    return 0 if success else 1


if __name__ == "__main__":
    exit(main())