src/stemsplitter/cli.py

"""Command-line interface for StemSplitter."""

from __future__ import annotations

import sys
from dataclasses import replace
from pathlib import Path

import click

from stemsplitter.config import get_settings
from stemsplitter.separator import OutputFormat, StemMode, StemSplitter


@click.command()
@click.argument("input_file", type=click.Path(exists=True, dir_okay=False))
@click.option(
    "-m",
    "--mode",
    type=click.Choice(["2stem", "4stem"], case_sensitive=False),
    default="2stem",
    show_default=True,
    help="Separation mode: 2-stem (vocals/instrumental) or 4-stem.",
)
@click.option(
    "-f",
    "--format",
    "output_format",
    type=click.Choice(["WAV", "MP3", "FLAC"], case_sensitive=False),
    default=None,
    help="Output audio format. Defaults to value in .env or WAV.",
)
@click.option(
    "-o",
    "--output-dir",
    type=click.Path(file_okay=False),
    default=None,
    help="Output directory. Defaults to value in .env or ./output.",
)
@click.option(
    "--model",
    default=None,
    help="Override the model filename (e.g., htdemucs.yaml).",
)
def main(
    input_file: str,
    mode: str,
    output_format: str | None,
    output_dir: str | None,
    model: str | None,
) -> None:
    """Separate audio stems from INPUT_FILE.

    Splits an audio file into vocal and instrumental stems (2-stem mode)
    or into vocals, drums, bass, and other (4-stem mode).

    \b
    Examples:
        stemsplitter song.mp3
        stemsplitter song.wav -m 4stem -f FLAC
        stemsplitter song.flac -m 2stem -f MP3 -o ./stems/
    """
    settings = get_settings()

    if output_dir:
        settings = replace(settings, output_dir=output_dir)

    Path(settings.output_dir).mkdir(parents=True, exist_ok=True)

    splitter = StemSplitter(settings=settings)

    stem_mode = StemMode(mode)
    fmt = OutputFormat(output_format.upper()) if output_format else None

    click.echo(
        f"Processing: {input_file}\n"
        f"Mode: {stem_mode.value} | "
        f"Format: {(fmt or OutputFormat(settings.output_format)).value}"
    )

    try:
        result = splitter.separate(
            input_path=input_file,
            mode=stem_mode,
            output_format=fmt,
            model_override=model,
        )
    except FileNotFoundError as exc:
        click.secho(str(exc), fg="red", err=True)
        sys.exit(1)
    except RuntimeError as exc:
        click.secho(f"Separation failed: {exc}", fg="red", err=True)
        sys.exit(1)

    click.secho("Separation complete!", fg="green")
    for f in result.output_files:
        click.echo(f"  -> {f}")