Server-scripts/media_editing/ffmpeg_handling.py

#!/home/lucas/scripts/.venv/bin/python3
"""
FFmpeg Python CLI Tool

A wrapper for ffmpeg-python to perform common media tasks.
Parameters are documented with their equivalent FFmpeg command-line flags.
"""

import argparse
import json
import logging
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

import ffmpeg
from mutagen.easyid3 import EasyID3
from serverlibrary import *

logging.setLoggerClass(CustomLogger)
logger: CustomLogger = logging.getLogger(__name__)  # type: ignore[assignment]

VIDEO_EXTS = {".mp4", ".mkv", ".avi", ".mov", ".webm"}
AUDIO_EXTS = {".mp3", ".m4a", ".m4b", ".opus"}

# ----------------------------
# OUTPUT ENFORCEMENT (clean + centralized)
# ----------------------------
commands = {
    "test": {"name": "test", "output": False},
    "info": {"name": "info", "output": False},
    "convert-av": {"name": "convert-av", "output": True},
    "convert-audio": {"name": "convert-audio", "output": True},
    "convert-video": {"name": "convert-video", "output": True},
    "remove-subs": {"name": "remove-subs", "output": True},
    "remove-audio": {"name": "remove-audio", "output": True},
    "extract-audio": {"name": "extract-audio", "output": True},
    "downsize": {"name": "downsize", "output": True},
    "filter-audio": {"name": "filter-audio", "output": True},
    "cut": {"name": "cut", "output": True},
    "set-primary-audio": {"name": "set-primary-audio", "output": True},
    "chapters": {"name": "chapters", "output": True},
    "metadata-print": {"name": "metadata-print", "output": False},
    "metadata-set": {"name": "metadata-set", "output": True},
    "remove-russian-from-db": {"name": "remove-russian-from-db", "output": True},
}
# ------ Russian Dictonary --------------#
AUDIO_TRANSLATIONS = {
    # Additional descriptors
    "Закадровый": "Voice-over",
    "Закадровое озвучивание": "Voice-over",
    "Многоголосый закадровый": "Multi-voice voice-over",
    "Многоголосное закадровое озвучивание": "Multi-voice voice-over",
    # Variants found in metadata
    "Оригинален език": "Original language",
    "Многоголосное озвучивание": "Multi-voice dub",
    "Многоголосый": "Multi-voice",
    "Любительский": "Amateur voice-over",
    "многоголосый": "Multi-voice",
    "Профессиональный": "Professional",
    "Одноголосый": "Single voice",
    "Двухголосый": "Dual voice",
    "Авторский": "Author voice-over",
    "закадровый": "Voice-over",
    "Оригинал": "Original",
    "Дубляж": "Dubbed",
    # Russian groups
    "Кубик в кубе": "Kubik v Kube",
    "LostFilm": "LostFilm",
    "AlexFilm": "AlexFilm",
    "Baibako": "Baibako",
    "NewStudio": "NewStudio",
    "Jaskier": "Jaskier",
    "Kerob": "Kerob",
    "TVShows": "TVShows",
    "HDRezka": "HDRezka",
    # Groups / studios
    "Синема УС": "Cinema US",
    "KinoPub": "KinoPub",
    "Озвучка Rumble": "Rumble",
    "Rumble": "Rumble",
    # Voice actor
    "Юрий Сербин": "Yuriy Serbin",
    "Ю. Сербин": "Y. Serbin",
    # Languages
    "Английски": "English",
    "Русский": "Russian",
    "Русская": "Russian",
    "Русское": "Russian",
    # ISO codes
    "rus": "Russian",
    "ru": "Russian",
    "ger": "German",
    "deu": "German",
    "eng": "English",
    "und": "Unknown",
}
RUSSIAN_RE = re.compile(
    r"\brus\b|"
    r"\bru\b|"
    r"russian|"
    r"рус|"
    r"русский|"
    r"русская|"
    r"русское|"
    r"дубляж|"
    r"оригинал|"
    r"многоголос|"
    r"двухголос|"
    r"одноголос|"
    r"авторский|"
    r"любительский|"
    r"профессиональн|"
    r"закадров|"
    r"юрий\s+сербин|"
    r"ю\.\s*сербин|"
    r"сербин|"
    r"синема\s+ус|"
    r"kinopub|"
    r"rumble|"
    r"озвучка\s+rumble|"
    r"кубик\s+в\s+кубе|"
    r"lostfilm|"
    r"alexfilm|"
    r"baibako|"
    r"newstudio|"
    r"jaskier|"
    r"kerob|"
    r"tvshows|"
    r"hdrezka",
    re.IGNORECASE,
)


def is_russian_track(track: dict) -> bool:
    descr = bool(RUSSIAN_RE.search(track["description"]))
    language = bool(RUSSIAN_RE.search(track["language"]))
    lang = bool(RUSSIAN_RE.search(track["lang"]))
    return descr or lang or language


class ArgsBuffer:
    input: Path | None = None
    output: Path | None = None
    kidflag: bool | None = False
    loud: bool | None = False
    key: str | None = None
    value: str | None = None
    audio_index: list | None = None
    dry_run: bool | None = None
    audio_codec: str | None = None
    video_codec: str | None = None
    preset: str | None = None
    quality: int | None = None
    width: int | None = None
    height: int | None = None
    crf: int | None = None
    num: int | None = None
    start: str | None = None
    duration: str | None = None


CYRILLIC_RE = re.compile(r"[А-Яа-яЁё]")


def translate_audio_tracks(tracks: list[dict]) -> list[dict]:
    translated = []

    for track in tracks:
        new_track = track.copy()

        for key, value in new_track.items():
            if not isinstance(value, str):
                continue

            for src, dst in AUDIO_TRANSLATIONS.items():
                value = re.sub(
                    rf"\b{re.escape(src)}\b",
                    dst,
                    value,
                    flags=re.IGNORECASE,
                )

            if CYRILLIC_RE.search(value):
                raise ValueError(
                    f"Contains unidentified Cyrillic characters in {key}: {value}"
                )

            new_track[key] = value

        translated.append(new_track)

    return translated


# ---------- Utilities ---------- #


# ----------------------------
# HELPERS
# ----------------------------
def ask(prompt, default=None):
    value = input(f"{prompt}" + (f" [{default}]" if default else "") + ": ").strip()
    return value if value else default


def ensure(value, prompt):
    if value is None:
        return ask(prompt)
    return value


def move_overwrite_file(files_move_from, files_move_to, args):

    backup_files = args.backup
    silent = not args.loud
    for file_to, file_from in zip(files_move_to, files_move_from):
        if not file_to.is_file():
            logger.error(f"overwrite Target File does not exist: {file_to}")
            return False

        if not file_from.is_file():
            logger.trace(f"overwrite File does not exist: {file_from}")
            continue
        backup = file_to.parent / f"{file_to.stem}_backup{file_from.suffix}"

        if backup_files:
            # move original -> backup
            logger.trace(f"Move original to backup: {backup}")
            safe_rename(file_to, backup, silent=silent)

        # move overwrite -> original name
        retstr, ret = safe_rename(file_from, file_to, silent=silent, overwrite=True)
        logger.trace(retstr)
        if not ret:
            logger.error(f"Safe rename failed for: {file_from} -> {file_to}")
            return False
    return True


LOG_FILE = Path.home() / "ffmpeg_log.log"


def run_ffmpeg(stream, overwrite=True, loud=False, dry_run=False):
    kwargs = {
        "overwrite_output": overwrite,
    }

    kwargs.update(
        capture_stdout=True,
        capture_stderr=True,
        quiet=True,
    )

    if dry_run:
        logger.debug(f"Dry run: {stream}")
        return
    try:
        out, err = ffmpeg.run(stream, **kwargs)
        logger.trace(out)
        if err:
            logger.error(err)
    except ffmpeg.Error as e:
        logger.error(e.stderr)
        logger.error("FFmpeg failed.\n")
        sys.exit(1)


def probe_audio_streams(file):
    """
    Probes metadata to identify audio tracks.

    FFmpeg Equivalent:
    ffprobe -v error -select_streams a -show_entries stream=index:tags=language -of json
    """
    try:
        probe = ffmpeg.probe(file)
    except ffmpeg.Error as e:
        logger.error(f"ffprobe failed:{e}")
        sys.exit(1)

    return [
        (s.get("index"), s.get("tags", {}).get("language", "unknown"))
        for s in probe.get("streams", [])
        if s.get("codec_type") == "audio"
    ]


def select_audio_stream(file, provided=None):
    """
    Logical handler for stream selection. Ensures we target the correct
    absolute stream index when building filter graphs.
    """
    streams = probe_audio_streams(file)

    if not streams:
        logger.info("No audio streams found.")
        sys.exit(1)

    provided = int(provided) if provided else None
    if provided is not None:
        if any(s[0] == provided for s in streams):
            return provided
        logger.info(f"Provided index {provided} is not a valid audio stream.")

    if len(streams) == 1:
        return streams[0][0]

    logger.info(f"Found {len(streams)} audio streams:")
    for i, (idx, lang) in enumerate(streams):
        logger.info(f"[{i}] Index: {idx}, Language: {lang}")

    try:
        choice = int(input("Select stream number (the list index [x]): "))
        return streams[choice][0]
    except (ValueError, IndexError):
        logger.error("Invalid selection.")
        sys.exit(1)


def parse_time_to_seconds(value: str) -> float:
    """
    Convert human-readable time formats into seconds.

    Supported formats:
        - "100s" → 100
        - "1m20s" → 80
        - "2h10m5s" → 7805
        - "0h1m5s" → 65
        - "0:1:5" → 65 (hh:mm:ss or mm:ss)
    """

    if value is None:
        return None

    value = str(value).strip()

    # Case 1: pure seconds (e.g. 100s)
    if value.endswith("s") and value[:-1].isdigit():
        return float(value[:-1])

    # Case 2: h/m/s format (e.g. 1h2m3s)
    pattern = re.compile(r"(?:(?P<h>\d+)h)?(?:(?P<m>\d+)m)?(?:(?P<s>\d+)s)?")
    match = pattern.fullmatch(value)
    if match and any(match.groups()):
        h = int(match.group("h") or 0)
        m = int(match.group("m") or 0)
        s = int(match.group("s") or 0)
        return h * 3600 + m * 60 + s

    # Case 3: colon format (hh:mm:ss or mm:ss)
    if ":" in value:
        parts = list(map(int, value.split(":")))
        if len(parts) == 2:
            m, s = parts
            return m * 60 + s
        elif len(parts) == 3:
            h, m, s = parts
            return h * 3600 + m * 60 + s

    # fallback: raw number
    try:
        return float(value)
    except ValueError:
        raise argparse.ArgumentTypeError(f"Invalid time format: {value}")


def get_scodec(output_path: Path):
    ext = output_path.suffix

    if ext == "mp4":
        return "mov_text"

    if ext == "mkv":
        return "srt"

    return "copy"


def get_vcodec(output_path: Path):
    ext = output_path.suffix
    if ext == "mp4":
        return "libx264"

    return "copy"


def get_chapter_lengths(file_path):
    """
    Returns an array of chapter durations (in seconds).
    """

    cmd = [
        "ffprobe",
        "-v",
        "error",
        "-print_format",
        "json",
        "-show_chapters",
        file_path,
    ]

    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        check=True,
    )
    logger.trace(result.stdout)
    if result.returncode != 0:
        raise RuntimeError(result.stderr)

    data = json.loads(result.stdout)

    chapters = []
    for ch in data.get("chapters", []):
        chapters.append(
            {"start": float(ch["start_time"]), "end": float(ch["end_time"])}
        )

    return chapters


def natural_sort_key(path: Path):
    s = str(path.name)  # force string conversion
    return [
        int(part) if part.isdigit() else part.lower() for part in re.split(r"(\d+)", s)
    ]


# ---------- Commands ---------- #
def cmd_info(args):
    """
    Retrieves and prints metadata.

    FFmpeg Equivalent:
    ffprobe -show_format -show_streams -of json [input]
    """
    try:
        info = ffmpeg.probe(str(args.input))
        logger.info(json.dumps(info, indent=2))
    except ffmpeg.Error as e:
        logger.error("Error probing file:", e)

    return False, "nooutputneeded"


def cmd_convert_av(args):
    """
    Re-encodes both video and audio.

    FFmpeg Equivalent:
    ffmpeg -i [input] -c:v [video_codec] -c:a [audio_codec] -preset [preset] -crf [quality] [output]

    Parameters:
    - vcodec/c:v: Sets the video encoder (e.g., libx264).
    - acodec/c:a: Sets the audio encoder (e.g., aac).
    - preset: Trade-off between compression speed and efficiency (ultrafast to veryslow).
    - qp/crf: Constant Rate Factor. Lower is higher quality (18-28 is standard).
    """
    inp = ffmpeg.input(args.input)
    out = ffmpeg.output(
        inp,
        str(args.output),
        vcodec=args.video_codec,
        acodec=args.audio_codec,
        scodec=get_scodec(args.output),
        preset=args.preset,
        qp=args.quality,
        # This maps all streams from the first input once
        map="0",
        # Prevents FFmpeg from trying to guess/add extra default streams
        map_metadata=0,
    )
    return True, out


def cmd_convert_audio(args):
    """
    Transcodes audio while keeping video intact.

    FFmpeg Equivalent:
    ffmpeg -i [input] -map 0:v -map 0:a -c:v copy -c:a [audio_codec] [output]

    Parameters:
    - vcodec='copy': Enables 'stream copy' mode, bypassing the video decoder/encoder.
    """
    inp = ffmpeg.input(args.input)
    out = ffmpeg.output(
        inp,
        str(args.output),
        vcodec="copy",
        acodec=args.audio_codec,
        scodec=get_scodec(args.output),
        map="0",
        map_metadata=0,
    )
    return True, out


def cmd_convert_video(args):
    """
    Transcodes video while keeping audio intact.

    FFmpeg Equivalent:
    ffmpeg -i [input] -map 0:v -map 0:a -c:v [video_codec] -c:a copy [output]
    """
    inp = ffmpeg.input(args.input)
    out = ffmpeg.output(
        inp,
        str(args.output),
        vcodec=args.video_codec,
        acodec="copy",
        scodec=get_scodec(args.output),
        map="0",
        map_metadata=0,
    )
    return True, out


def cmd_remove_subs(args):
    """
    Strips all subtitle data.

    FFmpeg Equivalent:
    ffmpeg -i [input] -sn -c copy [output]

    Parameters:
    - sn: Disables subtitle recording.
    """
    probe = ffmpeg.probe(str(args.input))

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(args.input),
    ]

    if args.num is None:
        # remove all subtitle streams
        for stream in probe["streams"]:
            if stream["codec_type"] != "subtitle":
                cmd.extend(["-map", f"0:{stream['index']}"])
    else:
        # remove only specified subtitle streams
        for stream in probe["streams"]:
            if stream["codec_type"] == "subtitle" and stream["index"] in args.num:
                continue

            cmd.extend(["-map", f"0:{stream['index']}"])

    cmd.extend(
        [
            "-c",
            "copy",
            "-map_metadata",
            "0",
            str(args.output),
        ]
    )

    logger.debug("%s", " ".join(cmd))
    if args.dry_run:
        logger.debug(f"Dry run: {' '.join(cmd)}")
        return True, "nooutputneeded"
    try:
        subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
    except subprocess.CalledProcessError as e:
        logger.error("Command: %s", " ".join(e.cmd))
        logger.error("Return code: %s", e.returncode)
        logger.error("STDOUT:\n%s", e.stdout)
        logger.error("STDERR:\n%s", e.stderr)
        raise
    return True, "nooutputneeded"


def cmd_remove_audio_track(args):
    """
    1. Removes a specific audio track from the video (Main Output).
    """

    probe = ffmpeg.probe(str(args.input))
    remove_indices = []
    for audio_index in args.audio_index:
        remove_indices.append(select_audio_stream(args.input, audio_index))

    remaining_audio = [
        s
        for s in probe["streams"]
        if s["codec_type"] == "audio" and s["index"] not in remove_indices
    ]
    if not remaining_audio:
        logger.error("No audiotracks remaining")
        return False, "nooutputneeded"
    remaining_audio = translate_audio_tracks(remaining_audio)

    removed_was_default = any(
        s["index"] in remove_indices and s.get("disposition", {}).get("default", 0) == 1
        for s in probe["streams"]
    )

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(args.input),
    ]

    for stream in probe["streams"]:
        if stream["index"] not in remove_indices:
            cmd.extend(["-map", f"0:{stream['index']}"])

    cmd.extend(
        [
            "-c:v",
            "copy",
            "-c:a",
            "copy",
            "-c:s",
            get_scodec(args.output),
        ]
    )

    if removed_was_default and remaining_audio:
        # First remaining audio stream in the output becomes default
        cmd.extend(
            [
                "-disposition:a",
                "0",
                "-disposition:a:0",
                "default",
            ]
        )

    cmd.append(str(args.output))
    logger.trace(" ".join(cmd))
    if args.dry_run:
        logger.debug(f"Dry run: {' '.join(cmd)}")
        return True, "nooutputneeded"
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        logger.debug("FFmpeg subprocess successful")
        logger.trace(result.stdout)
        return True, "nooutputneeded"

    except subprocess.CalledProcessError as e:
        logger.error(
            "Command failed (%s):\n%s\n\nstderr:\n%s",
            e.returncode,
            " ".join(map(str, e.cmd)),
            e.stderr,
        )
        raise RuntimeError(f"Command failed: {' '.join(e.cmd)}") from e


def cmd_extract_audio_track(args):
    """
    1. Saves that specific audio track as a standalone MP3 (Extra Output).
    """
    inp = ffmpeg.input(args.input)
    probe = ffmpeg.probe(str(args.input))

    output_audio = Path(args.output)
    output_audio = str(output_audio.with_name(f"{output_audio.stem}.mp3"))

    outputs = []
    for index in args.audio_index:
        remove_idx = select_audio_stream(args.input, index)

        extracted_audio = None

        for s in probe["streams"]:
            idx = s["index"]
            if idx == remove_idx:
                extracted_audio = inp[str(idx)]

        outputs.append(
            ffmpeg.output(
                extracted_audio, output_audio, acodec="libmp3lame", audio_bitrate="192k"
            )
        )

    return True, outputs


def cmd_downsize(args):
    """
    Scales resolution and compresses.

    FFmpeg Equivalent:
    ffmpeg -i [input] -vf scale=[w]:[h] -c:v libx264 -crf [crf] -c:a copy [output]

    Parameters:
    - filter('scale'): Invokes the video filter (vf) for resizing.
    - crf: Set to 23 (default) for balanced quality.
    """
    inp = ffmpeg.input(args.input)
    video = inp.video.filter("scale", args.width, args.height)
    info = ffmpeg.probe(str(args.input))
    streams = info.get("streams", [])
    subtitle_streams = [s for s in streams if s.get("codec_type") == "subtitle"]

    subs = {}
    if len(subtitle_streams) > 0:
        subs = {inp["s"]}

    out = ffmpeg.output(
        video,
        inp.audio,
        *subs,
        str(args.output),
        vcodec="libx264",
        crf=args.crf,
        acodec="copy",
        scodec=get_scodec(args.output),
    )
    return True, out


def cmd_filter_audio(args):
    """
    Multi-stage audio mastering chain.

    FFmpeg Equivalent:
    ffmpeg -i [input] -filter_complex "[0:a:X]highpass...,acompressor...,equalizer...out" \
           -map 0:v -map out -c:v copy -c:a libvorbis [output]

    Filter Descriptions:
    1. highpass: Cuts frequencies below 100Hz (removes low-end rumble).
    2. acompressor: Reduces volume of loud peaks to increase overall clarity.
    3. equalizer: Adjusts specific frequency bands (2.5kHz and 4kHz for presence).
    4. deesser: Dampens harsh 'S' sounds.
    5. alimiter: Hard limit at 0.95 to prevent digital clipping (distortion).
    6. dynaudnorm: Smooths out loudness levels over time.
    """
    if len(args.audio_index) > 1:
        logger.warning(
            f"Filter Audio only filters the first element of list audio_index: {args.audio_index}"
        )

    index = args.audio_index[0]

    idx_filt = select_audio_stream(args.input, index)

    inp = ffmpeg.input(args.input)
    probe = ffmpeg.probe(str(args.input))

    audio_streams = []
    audio_titles = {}
    audio_languages = {}

    for s in probe["streams"]:
        if s.get("codec_type") == "audio":
            idx = s["index"]
            audio_streams.append(idx)

            tags = s.get("tags", {})
            audio_titles[idx] = tags.get("title", f"Audio {idx}")
            audio_languages[idx] = tags.get("language", "und")

    filtered_audios = []
    metadata_args = []

    out_audio_idx = 0

    for idx in audio_streams:
        audio = inp[str(idx)]

        # -------------------
        # ORIGINAL STREAM
        # -------------------
        filtered_audios.append(audio)

        metadata_args += [
            f"-metadata:s:a:{out_audio_idx}",
            f"title={audio_titles[idx]}",
            f"-metadata:s:a:{out_audio_idx}",
            f"language={audio_languages[idx]}",
        ]

        out_audio_idx += 1

        # -------------------
        # FILTERED STREAM
        # -------------------
        if idx == idx_filt:
            audio_filtered = (
                audio.filter("highpass", f=100)
                .filter(
                    "acompressor", threshold="-22dB", ratio=4, attack=10, release=150
                )
                .filter("equalizer", f=2500, t="q", w=1, g=5)
                .filter("equalizer", f=4000, t="q", w=1, g=4)
                .filter("deesser", i=0.4, f=0.25)
                .filter("alimiter", limit=0.95)
                .filter("dynaudnorm", f=200, g=15)
            )

            filtered_audios.append(audio_filtered)

            base_title = audio_titles[idx]
            filtered_title = base_title + " + filtered"
            lang = audio_languages[idx]

            metadata_args += [
                f"-metadata:s:a:{out_audio_idx}",
                f"title={filtered_title}",
                f"-metadata:s:a:{out_audio_idx}",
                f"language={lang}",
                f"-metadata:s:a:{out_audio_idx}",
                f"handler_name={filtered_title}",
            ]

            out_audio_idx += 1

    # -------------------
    # SUBTITLES
    # -------------------
    subs = []
    if any(s.get("codec_type") == "subtitle" for s in probe["streams"]):
        subs.append(inp["s"])

    # -------------------
    # OUTPUT (IMPORTANT FIX HERE)
    # -------------------
    input_args = [
        inp.video,
        *filtered_audios,
        *subs,
    ]

    out = ffmpeg.output(
        *input_args,
        str(args.output),
        vcodec="copy",
        acodec="libvorbis",
        scodec=get_scodec(args.output),
    ).global_args(*metadata_args)

    return True, out


def cmd_cut_segment(args):
    """
    Cut a segment from a video or audio file.

    FFmpeg parameters:
        ss (start time): where to begin cutting (seconds or timestamp)
        t (duration): length of segment to extract
        c=copy: stream copy (fast, no re-encoding)

    Notes:
        - Using ss before input enables fast seek
        - Copy mode is fast but may be less frame-accurate
    """

    start = parse_time_to_seconds(args.start)
    duration = parse_time_to_seconds(args.duration)

    inp = ffmpeg.input(args.input, ss=start)

    out = ffmpeg.output(
        inp,
        str(args.output),
        t=duration,
        map="0",
        vcodec=get_vcodec(args.output),
        acodec="copy",
        scodec=get_scodec(args.output),
    )
    logger.trace(f"Command: {out}")
    return True, out


def cmd_set_primary_audio(args):
    """
    Set a chosen audio stream as the PRIMARY (first/default) audio track.

    Uses select_audio_stream() to ensure correct stream selection.

    FFmpeg behavior:
        - We explicitly reorder audio streams
        - Selected stream becomes first audio track
        - All other audio streams are preserved
        - Video and subtitles are kept

    Important:
        FFmpeg does NOT support "set default audio track" directly in MP4.
        So we simulate it by stream reordering.
    """

    inp = ffmpeg.input(args.input)

    if len(args.audio_index) > 1:
        logger.warning(
            f"Set primary can only set the first element of list audio_index: {args.audio_index}"
        )
    index = args.audio_index[0]

    primary_audio = select_audio_stream(args.input, index)
    probe = ffmpeg.probe(str(args.input))

    audio_streams = [
        s["index"] for s in probe["streams"] if s.get("codec_type") == "audio"
    ]

    reordered_audio = [primary_audio] + [a for a in audio_streams if a != primary_audio]

    audio_inputs = [inp[f"a:{i - 1}"] for i in reordered_audio]

    out_kwargs = {}
    for i, stream in enumerate(reordered_audio):
        if stream == primary_audio:
            out_kwargs[f"disposition:a:{i}"] = "default"
        else:
            out_kwargs[f"disposition:a:{i}"] = "0"

    info = ffmpeg.probe(str(args.input))
    streams = info.get("streams", [])
    subtitle_streams = [s for s in streams if s.get("codec_type") == "subtitle"]

    subs = {}
    if len(subtitle_streams) > 0:
        subs = {inp["s"]}

    out = ffmpeg.output(
        inp.video,
        *audio_inputs,
        *subs,
        str(args.output),
        vcodec="copy",
        acodec="copy",
        scodec=get_scodec(args.output),
        **out_kwargs,
    )

    return True, out


def cmd_show_metadata(args, show_metadata=False):
    """
    Prints metadata directly from ffprobe output.
    """

    probe = ffmpeg.probe(str(args.input))
    format_tags = probe.get("format", {}).get("tags", {})

    if show_metadata:
        logger.info("=== FORMAT METADATA ===")
        for k, v in format_tags.items():
            logger.info(f"{k}: {v}")

    # logger.info(f"\n=== STREAM METADATA ===")
    # for i, stream in enumerate(probe.get("streams", [])):
    #    logger.info(f"\n--- Stream {i} ({stream.get('codec_type')}) ---")

    #    tags = stream.get("tags", {})
    #    if not tags:
    #        logger.info(f"No tags")
    #    else:
    #        for k, v in tags.items():
    #            logger.info(f"{k}: {v}")
    return False, format_tags


def cmd_set_metadata(args):
    """
    Set metadata for media files.

    - Default: uses ffmpeg (stream copy)
    - If file is .opus: uses kid3-cli (Vorbis tags)
    """

    ext = args.input.suffix.lower()

    # ─────────────────────────────────────────────
    # OPUS → use Kid3 CLI (Vorbis comments)
    # ─────────────────────────────────────────────
    if ext == ".opus" or args.kidflag:
        cmd = ["kid3-cli", "-c", f"set {args.key} '{args.value}'", args.input]
        if args.dry_run:
            logger.debug(f"Dry run: {' '.join(cmd)}")
            return True, "nooutputneeded"
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        logger.trace(result.stdout)
        return False, "nooutputneeded"  # Kid3 writes in-place

    # ─────────────────────────────────────────────
    # Default → ffmpeg (copy stream)
    # ─────────────────────────────────────────────
    inp = ffmpeg.input(args.input)

    out = ffmpeg.output(
        inp, str(args.output), c="copy", metadata=f"{args.key}={args.value}"
    )
    return True, out


def cmd_apply_chapters(args):
    """
    Apply chapter titles to an M4B using a simple title-only text file.

    Chapters file format:
        One title per line:
            Introduction
            Acknowledgment
            The War Begins

    Behavior:
        - Evenly distributes chapters across media duration
        - Generates ffmetadata internally
        - Applies chapters using ffmpeg-python
    """
    # ----------------------------
    # read chapter titles
    # ----------------------------
    with open(args.chapter_file, "r", encoding="utf-8") as f:
        titles = [line.strip() for line in f if line.strip()]

    if not titles:
        raise ValueError("Chapter file is empty")

    if Path(args.input).suffix == ".m4b":
        # ----------------------------
        # build ffmetadata chapters
        # ----------------------------
        chapters = get_chapter_lengths(args.input)

        if len(chapters) != len(titles):
            logger.error(
                f"Chapter lengths do not match: {len(chapters)}, {len(titles)}"
            )
        for idx, title in enumerate(titles):
            chapters[idx].update({"title": title})

        metadata_text = ";FFMETADATA1\n\n"
        for ch in chapters:
            start_ms = int(float(ch["start"]) * 1000)
            end_ms = int(float(ch["end"]) * 1000)
            metadata_text += (
                "[CHAPTER]\n"
                "TIMEBASE=1/1000\n"
                f"START={start_ms}\n"
                f"END={end_ms}\n"
                f"title={ch['title']}\n\n"
            )

        # ----------------------------
        # run ffmpeg
        # ----------------------------
        with tempfile.NamedTemporaryFile("w+", suffix=".txt", delete=False) as f:
            f.write(metadata_text)
            meta_path = f.name

        cmd = [
            "ffmpeg",
            "-y",
            "-i",
            args.input,
            "-f",
            "ffmetadata",
            "-i",
            meta_path,
            "-map_chapters",
            "1",
            "-map_metadata",
            "0",
            "-c",
            "copy",
            args.output,
        ]
        if args.dry_run:
            logger.debug(f"Dry run: {' '.join(cmd)}")
            return True, "nooutputneeded"
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        logger.trace(result.stdout)

    elif Path(args.input).suffix == ".mp3":
        filename = Path(args.input)
        logger.info(f"{args.file_counter}:{filename}")
        title = titles[args.file_counter]
        safe_title = title.replace("/", "-").replace(":", "-")
        new_name = f"{args.file_counter + 1:03d} - {safe_title}.mp3"

        new_path = filename.with_name(new_name)

        logger.info(f"{filename.name} -> {new_name}")
        audio = EasyID3(filename)
        audio["title"] = title
        audio["tracknumber"] = str(args.file_counter + 1)
        if args.dry_run:
            logger.debug("Dry run")
            return True, "nooutputneeded"
        audio.save()
        safe_rename(filename, new_path)

    return True, "nooutputneeded"


def cmd_remove_russian_from_db(args_in):
    """
    Find all media with russian audio
    remove russian audio tracks
    """
    file = args_in.input
    logger.trace("Run remove russian cmd")
    try:
        audio_tracks = translate_audio_tracks(get_audio_tracks(file))
    except ValueError as e:
        logger.warning("%s", e)
        return False, "nooutputneeded"

    logger.trace(f"Remove Russian from File: {file.name}")
    logger.trace(f"With Audio tracks: {audio_tracks}")
    remove_tracks = []
    for track in audio_tracks:
        if is_russian_track(track):
            remove_tracks.append(track["number"])
            logger.debug(f"Removing russian track {track['number']} from {file}")
    args = args_in
    args.audio_index = remove_tracks
    if remove_tracks:
        cmd_remove_audio_track(args)
        return True, "nooutputneeded"
    else:
        return False, "nooutputneeded"


def cmd_test(args):
    """
    TEST MODE:
    Runs each command individually on test.mp4 and probes outputs.
    """

    test_file_dir = Path("/home/lucas/scripts/media_editing/testfiles")
    test_files = [
        "test_auto.mkv",
        "test_auto.mp4",
        "test_auto_only_russian.mkv",
        "test_auto_no_subs.mkv",
        "test_auto_no_subs.mp4",
    ]
    test_files = [test_file_dir / file for file in test_files]
    test_dir = Path(tempfile.mkdtemp(prefix="ffmpeg_test_"))
    logger.info(f"Using test dir: {test_dir}")

    def red_tag(tag: str) -> str:
        return f"{red}[{tag}]{reset}"

    def green_tag(tag: str) -> str:
        return f"{green}[{tag}]{reset}"

    def run_and_probe(name, cmd_func, args_probe, _test_dict=None):
        if _test_dict is None:
            _test_dict = []
        logger.info(f"TEST: {name}")
        # Run FFmpeg for this single command only
        return_value ,outputs = cmd_func(args_probe)
        if not outputs:
            logger.error("No outputs returned.")
            return False
        if outputs != "nooutputneeded":
            try:
                ffmpeg.run(
                    outputs, overwrite_output=True, quiet=True, capture_stderr=True
                )
            except ffmpeg.Error as e:
                logger.error("FFmpeg failed:")
                if e.stderr:
                    logger.info(e.stderr.decode(errors="ignore"))
                return False
        logger.info("FFmpeg succeeded")
        # Probe result file(s)
        out_files = []

        if hasattr(args_probe, "output") and args_probe.output and return_value:
            out_files.append(args_probe.output)

        for idx, file in enumerate(out_files):
            if not file.exists():
                logger.info(f"{red_tag('FAIL')} Missing output: {file}")
                continue
            try:
                ret_va = validate_ffprobe(file, _test_dict[idx])
                return ret_va
            except ffmpeg.Error as e:
                logger.error(f"{red_tag('FAIL')} Probe error: {file}")
                logger.error(e)
                return False
        return True

    def is_cover_art(stream):
        if stream["codec_type"] != "video":
            return False

        if stream.get("disposition", {}).get("attached_pic") == 1:
            return True

        if stream.get("nb_frames") and int(stream.get("nb_frames")) <= 1:
            return True
        tags = stream.get("tags", {})
        if "DURATION" in tags:
            h, m, s = tags["DURATION"].split(":")
            if int(h) * 3600 + int(m) * 60 + float(s) < 1:
                return True
        if "MIMETYPE" in tags and tags["MIMETYPE"] == "image/jpg":
            return True
        if "FILENAME" in tags and tags["FILENAME"] == "cover.jpg":
            return True

        return False

    def get_streams(_test_file):
        info = ffmpeg.probe(_test_file)
        streams = info.get("streams", [])

        _video_streams = [
            s for s in streams if s.get("codec_type") == "video" and not is_cover_art(s)
        ]
        _audio_streams = [s for s in streams if s.get("codec_type") == "audio"]
        _subtitle_streams = [s for s in streams if s.get("codec_type") == "subtitle"]
        _cover_image = [
            s for s in streams if s.get("codec_type") == "video" and is_cover_art(s)
        ]
        if not os.path.exists(_test_file):
            logger.info("test file not found in script directory.")
            sys.exit(1)

        return _video_streams, _audio_streams, _subtitle_streams, _cover_image

    def validate_ffprobe(file_path, expected: dict):
        """
        Validate ffprobe output using a dictionary-based expectation.

        Expected dict format:
        {
            "video_count": int,
            "audio_count": int,
            "subtitle_count": int,
            "audio_codecs": list[str],
            "video_codecs": list[str]
        }
        """
        if not file_path.is_file():
            logger.error(f"Test file {file_path} is not a file")
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_chapters",
                "-show_format",
                file_path,
            ],
            capture_output=True,
            text=True,
            check=True,
        )

        info = json.loads(result.stdout)
        video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
            file_path
        )
        args = ArgsBuffer()
        args.input = file_path
        tags = cmd_show_metadata(args)

        return_value = True

        # ---------- COUNT CHECKS ----------
        if "video_count" in expected:
            if len(video_streams) != expected["video_count"]:
                logger.error(
                    f"{red_tag('VIDEO COUNT FAIL')} expected {expected['video_count']}, got {len(video_streams)}"
                )
                return_value = return_value and False

        if "audio_count" in expected:
            if len(audio_streams) != expected["audio_count"]:
                logger.error(
                    f"{red_tag('AUDIO COUNT FAIL')} expected {expected['audio_count']}, got {len(audio_streams)}"
                )
                return_value = return_value and False

        if "subtitle_count" in expected:
            if len(subtitle_streams) != expected["subtitle_count"]:
                logger.error(
                    f"{red_tag('SUBTITLE COUNT FAIL')} expected {expected['subtitle_count']}, got {len(subtitle_streams)}"
                )
                return_value = return_value and False

        # ---------- CODEC CHECKS ----------
        if "audio_codecs" in expected:
            actual = sorted([s.get("codec_name") for s in audio_streams])
            expected_codecs = sorted(expected["audio_codecs"])
            if len(expected_codecs) < len(actual):
                last = expected_codecs[-1] if expected_codecs else None
                expected_codecs = expected_codecs + [last] * (
                    len(actual) - len(expected_codecs)
                )

            if actual != expected_codecs:
                logger.error(
                    f"{red_tag('AUDIO CODEC FAIL')}\nExpected: {expected_codecs}\nGot: {actual}"
                )
                return_value = return_value and False

        if "video_codecs" in expected:
            actual = sorted([s.get("codec_name") for s in video_streams])
            expected_codecs = sorted(expected["video_codecs"])
            if len(expected_codecs) < len(actual):
                last = expected_codecs[-1] if expected_codecs else None
                expected_codecs = expected_codecs + [last] * (
                    len(actual) - len(expected_codecs)
                )
            if actual != expected_codecs:
                logger.error(
                    f"{red_tag('VIDEO CODEC FAIL')}\nExpected: {expected_codecs}\nGot: {actual}"
                )
                return_value = return_value and False

        # ---------- DURATION CHECKS ----------
        if "duration" in expected:
            format_info = info.get("format", {})
            actual_duration = float(format_info.get("duration", 0))
            expected_duration = parse_time_to_seconds(expected["duration"])

            # allow small ffmpeg timing variance (important!)
            tolerance = expected.get("duration_tolerance", 0.2)

            if abs(actual_duration - expected_duration) <= tolerance:
                logger.error(
                    f"{red_tag('DURATION FAIL')} expected ~{expected_duration} (±{tolerance}), got {actual_duration}s"
                )
                return_value = return_value and False

        # ---------- CHAPTER CHECKS ----------
        if "chapter_name" in expected:
            chapter_name = [
                chapter.get("tags", {}).get("title", "")
                for chapter in info.get("chapters", [])
            ]
            expected_chapter_name = expected["chapter_name"]
            if expected_chapter_name != chapter_name:
                logger.error(
                    f"Expected chapter_name: {expected_chapter_name}, got {chapter_name}"
                )
                return_value = return_value and False

        if "chapter_count" in expected:
            chapter_num = len(info.get("chapters", []))
            expected_chapter_count = expected["chapter_count"]
            if expected_chapter_count != chapter_num:
                logger.error(
                    f"Expected chapters: {expected_chapter_count}, got {chapter_num}"
                )
                return_value = return_value and False

        # ---------- METADATA CHECKS ----------

        if "metadata" in expected:
            keys = expected["metadata"].keys()
            values = expected["metadata"].values()
            for key, value in zip(keys, values):
                if key in tags.keys() and tags[key] == value:
                    logger.debug(f"{key}: {value}")
                else:
                    logger.error(f"Key {key}: {value} not in expected tags")
                    return_value = return_value and False

        if "no_russian_dir" in expected:
            audio_tracks = translate_audio_tracks(get_audio_tracks(file_path))
            logger.debug(f"Audio tracks: {audio_tracks}")
            for audio_track in audio_tracks:
                if is_russian_track(audio_track["description"]):
                    logger.error(
                        f"Found russian track: {audio_track} in file {file_path}"
                    )
                    return return_value and False

            # logger.debug(f"Subtitle tracks: {subtitle_tracks}")
            # for subtitle_track in subtitle_tracks:
            #    if is_russian_track(audio_track["description"]):
            #        logger.error(f"Found russian track: {subtitle_track} in file {file}")
            #        return return_value and False

        logger.info(f"{green_tag('PASS')} {file_path}")
        return return_value

    ret = True
    # ---------- INFO ----------
    # args = ArgsBuffer()
    # args.input = test_file_mp4
    # test_dict = {}
    # run_and_probe(cmd["name"], cmd_info, args, test_dict)
    # args.input = test_file_mkv
    # ret = ret and  run_and_probe(cmd["name"], cmd_info, args, test_dict)

    # ---------- DOWNSIZE ----------
    args_test = ArgsBuffer()

    cmd = commands["downsize"]
    if cmd["name"] in args.choices:
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.width = 640
            args_test.height = 360
            args_test.crf = 28
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                    "video_codecs": ["h264"],
                }
            ]
            ret = ret and run_and_probe(cmd["name"], cmd_downsize, args_test, test_dict)

    # ---------- REMOVE AUDIO TRACK ----------
    cmd = commands["remove-audio"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        args_test.audio_index = [1]

        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": 1,
                    "subtitle_count": len(subtitle_streams),
                }
            ]
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            
            ret = ret and run_and_probe(
                cmd["name"], cmd_remove_audio_track, args_test, test_dict
            )

    # ---------- EXTRACT AUDIO TRACK ----------
    cmd = commands["extract-audio"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        args_test.audio_index = [1]

        for test_file in test_files:
            args_test.input = test_file
            test_dict = [{"video_count": 0, "audio_count": 1, "subtitle_count": 0}]
            args_test.output = (
                Path(test_dir) / f"extracted_audio_{Path(test_file).suffix[1:]}.mp3"
            )
            ret = ret and run_and_probe(
                cmd["name"], cmd_extract_audio_track, args_test, test_dict
            )

    # ---------- CONVERT AV ----------
    cmd = commands["convert-av"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.video_codec = "libx264"
            args_test.audio_codec = "aac"
            args_test.preset = "fast"
            args_test.quality = 23
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                    "video_codecs": ["h264"],
                    "audio_codecs": ["aac"],
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_convert_av, args_test, test_dict
            )

    # ---------- CONVERT AUDIO ----------
    cmd = commands["convert-audio"]

    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.audio_codec = "aac"
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                    "audio_codecs": ["aac"],
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_convert_audio, args_test, test_dict
            )

    # ---------- CONVERT VIDEO ----------
    cmd = commands["convert-video"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.video_codec = "libx264"
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                    "video_codecs": ["h264"],
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_convert_video, args_test, test_dict
            )

    # ---------- REMOVE SUBS ----------
    cmd = commands["remove-subs"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.num = None
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": 0,
                }
            ]
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            ret = ret and run_and_probe(
                cmd["name"], cmd_remove_subs, args_test, test_dict
            )

    # ---------- FILTER AUDIO ----------
    cmd = commands["filter-audio"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.audio_index = [1]
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams) + 1,
                    "subtitle_count": len(subtitle_streams),
                    "audio_codecs": ["vorbis"],
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_filter_audio, args_test, test_dict
            )

    # ---------- CUT ----------
    cmd = commands["cut"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.start = "0:0:5"
            args_test.duration = "10s"
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                    "duration": "10s",
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_cut_segment, args_test, test_dict
            )

    # ---------- SET PRIMARY AUDIO ----------
    cmd = commands["set-primary-audio"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        for test_file in test_files:
            video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
                test_file
            )
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.audio_index = [2]
            test_dict = [
                {
                    "video_count": len(video_streams),
                    "audio_count": len(audio_streams),
                    "subtitle_count": len(subtitle_streams),
                }
            ]
            ret = ret and run_and_probe(
                cmd["name"], cmd_set_primary_audio, args_test, test_dict
            )

    # ---------- APPLY CHAPTERS-----
    cmd = commands["chapters"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        test_file = test_file_dir / "test_auto.m4b"
        video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
            test_file
        )
        args_test.input = test_file
        args_test.output = (
            Path(test_dir)
            / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
        )
        args_test.chapter_file = test_file_dir / "test_auto_chapters.txt"
        lines = Path(args_test.chapter_file).read_text(encoding="utf-8").splitlines()
        test_dict = [
            {
                "video_count": len(video_streams),
                "audio_count": len(audio_streams),
                "chapter_name": lines,
                "chapter_count": len(lines),
            }
        ]
        ret = ret and run_and_probe(
            cmd["name"], cmd_apply_chapters, args_test, test_dict
        )

    # ---------- WRITE METADATA-----
    cmd = commands["metadata-set"]
    if cmd["name"] in args.choices:
        args_test = ArgsBuffer()
        test_file = test_file_dir / "test_auto.m4b"
        video_streams, audio_streams, subtitle_streams, cover_image = get_streams(
            test_file
        )
        args_test.input = test_file
        args_test.output = (
            Path(test_dir)
            / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
        )
        args_test.key = "artist"
        args_test.value = "TEST"
        args_test.kidflag = False
        test_dict = [
            {
                "video_count": len(video_streams),
                "audio_count": len(audio_streams),
                "metadata": {"artist": "TEST"},
            }
        ]
        ret = ret and run_and_probe(cmd["name"], cmd_set_metadata, args_test, test_dict)

    # ---------- REMOVE RUSSIAN AUDIO -----
    cmd = commands["remove-russian-from-db"]
    if cmd["name"] in args.choices:
        for test_file in test_files:
            args_test = ArgsBuffer()
            args_test.input = test_file
            args_test.output = (
                Path(test_dir)
                / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
            )
            args_test.loud = False
            test_dict = [{"no_russian_dir": test_dir}]
            ret = ret and run_and_probe(
                cmd["name"], cmd_remove_russian_from_db, args_test, test_dict
            )

    ## -----------FOLDER TEST---------------
    ## ---------- REMOVE RUSSIAN AUDIO -----
    # cmd = commands["remove-russian-from-db"]
    # if cmd["name"] in args.choices:
    #    args_test = ArgsBuffer()
    #    args_test.input = test_file_dir
    #    args_test.output = (
    #        Path(test_dir)
    #        / f"out_{str(cmd['name']).replace('-', '_')}_{Path(test_file).stem}{Path(test_file).suffix}"
    #    )
    #    args_test.loud = False
    #    test_dict = [{"no_russian_dir": test_dir}]
    #    ret = ret and run_and_probe(
    #        cmd["name"], cmd_remove_russian_from_db, args_test, test_dict
    #    )

    if ret:
        logger.info(green_tag("ALL TESTS SUCCEEDED"))
    else:
        logger.info(red_tag("SOME TESTS FAILED"))


# ---------- CLI Setup ---------- #


def add_command(
    subparsers,
    name,
    func,
    help_text,
    aliases=None,
    parents=None,
):
    parser = subparsers.add_parser(
        name,
        aliases=aliases or [],
        parents=parents or [],
        help=help_text,
    )
    parser.set_defaults(
        func=func,
        command=name,
    )
    return parser


def add_audio_index(parser):
    parser.add_argument(
        "-a",
        "--audio-index",
        type=int,
        nargs="+",
        help="Audio track index",
    )


def main():
    parser = argparse.ArgumentParser(
        description="Media processing tool using ffmpeg-python",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # =========================================================
    # SHARED ARGS
    # =========================================================

    base_p = argparse.ArgumentParser(add_help=False)

    base_p.add_argument("input", nargs="?", type=Path, help="Input file or directory")
    base_p.add_argument("-o", "--output", type=Path, help="Output file")
    base_p.add_argument("--overwrite", action="store_true", help="Overwrite input")
    base_p.add_argument("--backup", action="store_true", help="Backup original")
    base_p.add_argument("--loud", action="store_true", help="Show ffmpeg output")
    base_p.add_argument(
        "--kidflag", action="store_true", help="Use kid3 instead of ffmpeg"
    )
    base_p.add_argument(
        "-r", "--recursive", action="store_true", help="Process recursively"
    )

    base_p.add_argument(
        "-ll",
        "--log-level",
        type=str.upper,
        default="INFO",
        choices=["TRACE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
    )

    base_p.add_argument(
        "--dry-run",
        action="store_true",
        help="Do not execute actions",
    )

    sub = parser.add_subparsers(
        dest="command",
        required=True,
    )

    # =========================================================
    # INFO
    # =========================================================

    add_command(
        sub,
        "info",
        cmd_info,
        "Inspect media metadata",
        aliases=["i"],
        parents=[base_p],
    )

    # =========================================================
    # TEST
    # =========================================================

    p = add_command(
        sub,
        "test",
        cmd_test,
        "Run tests",
        aliases=["t"],
        parents=[base_p],
    )

    p.add_argument(
        "--choices",
        nargs="+",
        choices=list(commands.keys()),
        default=list(commands.keys()),
    )

    # =========================================================
    # CONVERT AV
    # =========================================================

    p = add_command(
        sub,
        "convert-av",
        cmd_convert_av,
        "Full video/audio re-encode",
        aliases=["cav"],
        parents=[base_p],
    )

    p.add_argument("--video-codec", default="libx264")
    p.add_argument("--audio-codec", default="aac")
    p.add_argument("--preset", default="medium")
    p.add_argument("--quality", type=int)

    # =========================================================
    # CONVERT AUDIO
    # =========================================================

    p = add_command(
        sub,
        "convert-audio",
        cmd_convert_audio,
        "Convert audio only",
        aliases=["ca"],
        parents=[base_p],
    )

    p.add_argument("--audio-codec", default="aac")

    # =========================================================
    # CONVERT VIDEO
    # =========================================================

    p = add_command(
        sub,
        "convert-video",
        cmd_convert_video,
        "Convert video only",
        aliases=["cv"],
        parents=[base_p],
    )

    p.add_argument("--video-codec", default="libx264")

    # =========================================================
    # REMOVE SUBS
    # =========================================================

    add_command(
        sub,
        "remove-subs",
        cmd_remove_subs,
        "Strip subtitle streams",
        aliases=["rs"],
        parents=[base_p],
    )

    # =========================================================
    # REMOVE AUDIO
    # =========================================================

    p = add_command(
        sub,
        "remove-audio",
        cmd_remove_audio_track,
        "Remove audio track",
        aliases=["ra"],
        parents=[base_p],
    )
    add_audio_index(p)

    # =========================================================
    # EXTRACT AUDIO
    # =========================================================

    p = add_command(
        sub,
        "extract-audio",
        cmd_extract_audio_track,
        "Extract audio track",
        aliases=["ea"],
        parents=[base_p],
    )
    add_audio_index(p)

    # =========================================================
    # DOWNSIZE
    # =========================================================

    p = add_command(
        sub,
        "downsize",
        cmd_downsize,
        "Scale video",
        aliases=["d"],
        parents=[base_p],
    )

    p.add_argument("-w", "--width", type=int)
    p.add_argument("-he", "--height", type=int)
    p.add_argument("-c", "--crf", type=int)

    # =========================================================
    # FILTER AUDIO
    # =========================================================

    p = add_command(
        sub,
        "filter-audio",
        cmd_filter_audio,
        "Apply audio filters",
        aliases=["fa"],
        parents=[base_p],
    )
    add_audio_index(p)

    # =========================================================
    # CUT
    # =========================================================

    p = add_command(
        sub,
        "cut",
        cmd_cut_segment,
        "Cut media segment",
        aliases=["c"],
        parents=[base_p],
    )

    p.add_argument("--start")
    p.add_argument("--duration")

    # =========================================================
    # SET PRIMARY AUDIO
    # =========================================================

    p = add_command(
        sub,
        "set-primary-audio",
        cmd_set_primary_audio,
        "Set default audio track",
        aliases=["spa"],
        parents=[base_p],
    )
    add_audio_index(p)

    # =========================================================
    # REMOVE RUSSIAN AUDIO
    # =========================================================

    add_command(
        sub,
        "remove-russian-from-db",
        cmd_remove_russian_from_db,
        "Remove russian audio tracks",
        aliases=["rrfd"],
        parents=[base_p],
    )

    # =========================================================
    # CHAPTERS
    # =========================================================

    p = add_command(
        sub,
        "chapters",
        cmd_apply_chapters,
        "Apply chapter metadata",
        aliases=["ch"],
        parents=[base_p],
    )

    p.add_argument(
        "-cf",
        "--chapter-file",
        help="ffmetadata chapter file",
    )

    # =========================================================
    # METADATA PRINT
    # =========================================================

    add_command(
        sub,
        "metadata-print",
        cmd_show_metadata,
        "Show metadata",
        aliases=["mp"],
        parents=[base_p],
    )

    # =========================================================
    # METADATA SET
    # =========================================================

    p = add_command(
        sub,
        "metadata-set",
        cmd_set_metadata,
        "Set metadata value",
        aliases=["ms"],
        parents=[base_p],
    )

    p.add_argument("-k", "--key")
    p.add_argument("-v", "--value")

    # =========================================================
    # PARSE
    # =========================================================

    args = parser.parse_args()

    setup_logging(args.log_level, "ffmpeg_handling")
    # ----------------------------
    # COMMAND VALIDATION
    # ----------------------------
    if not hasattr(args, "func"):
        logger.info("\nAvailable commands:")

        seen = set()
        for name, parser in sub.choices.items():
            if id(parser) in seen:
                continue
            seen.add(id(parser))
            logger.info(f"  - {name}")

        args.command = ensure(None, "Command")

        if args.command not in sub.choices:
            parser.error("Unknown command")

        # re-parse subcommand properly
        sub_args = sub.choices[args.command].parse_args([])
        args = argparse.Namespace(**vars(sub_args))

    # ----------------------------
    # INTERACTIVE FALLBACK (ensure ALL missing values)
    # ----------------------------
    if args.command != "test":
        args.input = Path(
            ensure(getattr(args, "input", None), f"Input file ({Path.cwd()})")
        )
        args.input = Path(args.input).expanduser().resolve()
        if not Path(args.input).exists():
            parser.error(f"Input file {args.input} does not exist!")
    if commands[args.command]["output"] and not getattr(args, "output", None):
        input_path = Path(args.input)
        if args.overwrite:
            args.output = input_path.with_name(
                f"{input_path.stem}_overwrite{input_path.suffix}"
            )
        else:
            args.output = ensure(
                getattr(args, "output", None),
                "Output file (enter for default file)",
            )
        if args.output == "" or args.output is None:
            input_path = Path(args.input)
            args.output = input_path.with_name(
                f"{input_path.stem}_{args.command}{input_path.suffix}"
            )
        else:
            args.output = Path(args.output)
    # ----------------------------
    # COMMAND-SPECIFIC REQUIRED FIELDS
    # ----------------------------

    if args.command == "convert-av":
        args.preset = ensure(
            getattr(args, "preset", None),
            "Encoding speed/quality preset (default: medium, e.g. ultrafast, fast, slow, veryslow)",
        )
        args.quality = ensure(
            getattr(args, "quality", None),
            "Output quality level (e.g. CRF value or bitrate depending on encoder)",
        )
    if args.command == "convert-av" or args.command == "convert-audio":
        args.audio_codec = ensure(
            getattr(args, "audio_codec", None),
            "Audio codec to use (default: aac, e.g. mp3, opus)",
        )
    if args.command == "convert-av" or args.command == "convert-video":
        args.video_codec = ensure(
            getattr(args, "video_codec", None),
            "Video codec to use (default: libx264, e.g. libx265, vp9)",
        )

    if args.command == "cut":
        args.start = ensure(getattr(args, "start", None), "Start time (e.g. 10s, 1:00)")
        args.duration = ensure(getattr(args, "duration", None), "Duration (e.g. 30s)")

    if args.command == "downsize":
        args.width = ensure(getattr(args, "width", None), "width: eg 1080")
        args.height = ensure(getattr(args, "height", None), "height: eg 720")
        args.crf = ensure(getattr(args, "crf", None), "crf: eg 23")

    if args.command in [
        "filter-audio",
        "remove-audio",
        "extract-audio",
        "set-primary-audio",
    ]:
        index = str(ensure(getattr(args, "audio_index", None), "Audio index"))
        args.audio_index = [int(x) for x in re.split(r"[,\s]+", index.strip()) if x]
        if args.audio_index is None:
            args.audio_index = [1]

    if args.command == "chapters":
        args.chapter_file = ensure(
            getattr(args, "chapter_file", None), "File with chapternames"
        )

    if args.command == "metadata-set":
        args.chapter_file = ensure(getattr(args, "key", None), "Key to change")
        args.chapter_file = ensure(getattr(args, "value", None), "Key value")

    # ----------------------------
    # EXECUTE COMMAND
    # ----------------------------
    try:
        if args.command == "test":
            args.func(args)
            sys.exit(0)
        outputs = []
        input_path = Path(args.input)
        files_move_from = []
        files_move_to = []
        # ---- DIRECTORY ----
        if input_path.is_dir():
            counter = 0
            logger.info(f"Processing dir: {input_path}")
            if args.recursive:
                sorted_files = sorted(
                    (
                        p
                        for p in input_path.rglob("*")
                        if p.is_file() and p.suffix.lower() in VIDEO_EXTS
                    ),
                    key=natural_sort_key,
                )
            else:
                sorted_files = sorted(
                    (
                        p
                        for p in input_path.iterdir()
                        if p.is_file() and p.suffix.lower() in VIDEO_EXTS
                    ),
                    key=natural_sort_key,
                )
            for file in sorted_files:
                if (
                    file.suffix.lower() not in VIDEO_EXTS
                    and file.suffix.lower() not in AUDIO_EXTS
                ):
                    continue
                logger.trace(f"Processing File: {file}")
                args.input = Path(file)
                if args.overwrite:
                    out_path = file.with_name(f"{file.stem}_overwrite{file.suffix}")
                else:
                    out_path = file.with_name(
                        f"{file.stem}_{args.command}{file.suffix}"
                    )
                args.output = out_path
                args.file_counter = counter
                return_value, output_str = args.func(args)
                if output_str != "nooutputneeded":
                    outputs.append(output_str)
                if args.overwrite and return_value:
                    files_move_from.append(out_path)
                    files_move_to.append(file)
                counter += 1
        # ---- SINGLE FILE ----
        elif input_path.is_file():
            return_value, output_str = args.func(args)
            if output_str != "nooutputneeded":
                outputs.append(output_str)
        else:
            raise ValueError(f"Invalid input path: {args.input}")

        if any(item is None for item in outputs) or len(outputs) == 0:
            logger.info("Output is NONE")
        else:
            if not isinstance(outputs, list):
                outputs = [outputs]

            if args.command != "test" and outputs is not None:
                run_ffmpeg(
                    ffmpeg.merge_outputs(*outputs), overwrite=True, loud=args.loud
                )
        if args.overwrite and not move_overwrite_file(
            files_move_from, files_move_to, args
        ):
            logger.error("move_overwrite failed")
            sys.exit(1)
    except ffmpeg.Error as e:
        logger.error("FFmpeg failed.\n")
        if e.stderr:
            logger.error(e.stderr.decode(errors="ignore"))
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()