audio_splitter/audio_splitter/handlers.py

# audio_splitter/handlers.py
"""Handler system for cover image attachment.

This module provides a registry of handlers for different
codec + image combinations. Each handler is responsible for
attaching the cover image to the output file in the appropriate way.
"""

import os
import subprocess
from typing import Dict, Optional


def _run_opustags(input_file: str, cover_image_path: str, output_path: Optional[str] = None) -> bool:
    """
    Run opustags to attach a cover image to an Opus file.

    Args:
        input_file: Path to the input Opus file.
        cover_image_path: Path to the cover image.
        output_path: Optional output path (if None, overwrites input).

    Returns:
        True if successful, False otherwise.
    """
    cmd = ['opustags', '--set-cover', cover_image_path, input_file]

    if output_path:
        cmd.extend(['-o', output_path, '-y'])
    else:
        cmd.extend(['-i'])

    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    if result.returncode != 0:
        print(f"opustags failed: {result.stderr}")
        return False
    return True


def _handler_default(
    input_file: str,
    cover_image_path: str,
    output_path: str,
    stream_info: Dict,
    audio_codec: str,
    video_codec: str,
    video_quality: Optional[int],
    drop_video: bool,
    drop_subs: bool,
) -> bool:
    """
    Default handler: use ffmpeg with -disposition attached_pic.

    This is the standard approach for formats like MP3, FLAC, etc.
    """
    # This handler doesn't do anything - the cover is already handled
    # by the ffmpeg command in core.py via build_ffmpeg_command()
    # We just return True to indicate success.
    return True


def _handler_opus_with_cover(
    input_file: str,
    cover_image_path: str,
    output_path: str,
    stream_info: Dict,
    audio_codec: str,
    video_codec: str,
    video_quality: Optional[int],
    drop_video: bool,
    drop_subs: bool,
) -> bool:
    """
    Handler for Opus files with cover image.

    Uses opustags to attach the cover image after the file is processed.
    """
    return _run_opustags(output_path, cover_image_path)


def _handler_mp3_with_cover(
    input_file: str,
    cover_image_path: str,
    output_path: str,
    stream_info: Dict,
    audio_codec: str,
    video_codec: str,
    video_quality: Optional[int],
    drop_video: bool,
    drop_subs: bool,
) -> bool:
    """
    Handler for MP3 files with cover image.

    Uses ffmpeg with -disposition attached_pic.
    """
    # MP3 handler - the cover is already attached by ffmpeg in core.py
    return True


def _handler_flac_with_cover(
    input_file: str,
    cover_image_path: str,
    output_path: str,
    stream_info: Dict,
    audio_codec: str,
    video_codec: str,
    video_quality: Optional[int],
    drop_video: bool,
    drop_subs: bool,
) -> bool:
    """
    Handler for FLAC files with cover image.

    Uses ffmpeg with -disposition attached_pic.
    """
    # FLAC handler - the cover is already attached by ffmpeg in core.py
    return True


# Handler registry: maps (output_container, input_codec) to handler function
HANDLER_REGISTRY: Dict[tuple, callable] = {
    ('opus', 'opus'): _handler_opus_with_cover,
    ('ogg', 'opus'): _handler_opus_with_cover,
    ('mp3', 'mp3'): _handler_mp3_with_cover,
    ('flac', 'flac'): _handler_flac_with_cover,
}


def get_handler(output_container: str, input_audio_codec: Optional[str] = None):
    """
    Get the appropriate handler for the given container and input codec.

    Args:
        output_container: Output container format (e.g., 'opus', 'mp3', 'flac').
        input_audio_codec: Input audio codec (e.g., 'opus', 'mp3', 'flac').

    Returns:
        Handler function, or the default handler if no specific handler is found.
    """
    if input_audio_codec:
        key = (output_container, input_audio_codec)
        if key in HANDLER_REGISTRY:
            return HANDLER_REGISTRY[key]

    # Fall back to container-only key
    if (output_container, None) in HANDLER_REGISTRY:
        return HANDLER_REGISTRY[(output_container, None)]

    # Default handler
    return _handler_default


def needs_drop_video(output_container: str, input_audio_codec: Optional[str] = None) -> bool:
    """
    Check if the handler requires dropping video even if user didn't specify --drop-video.

    Args:
        output_container: Output container format.
        input_audio_codec: Input audio codec.

    Returns:
        True if video should be dropped, False otherwise.
    """
    handler = get_handler(output_container, input_audio_codec)
    # The opus handler requires dropping video
    return handler == _handler_opus_with_cover