#!/home/lucas/.venv_general/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 os
import sys
import json
import ffmpeg
import re
from pathlib import Path

# ----------------------------
# OUTPUT ENFORCEMENT (clean + centralized)
# ----------------------------
REQUIRES_OUTPUT = {
    "convert-av",
    "convert-audio",
    "convert-video",
    "remove-subs",
    "remove-audio",
    "downsize",
    "filter-audio",
    "cut",
    "set-primary-audio",
}
# ---------- 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 run_ffmpeg(stream, overwrite=True):
    """
    Run FFmpeg and stream output directly to terminal.

    - No log buffering
    - Real-time progress/errors visible
    - Easier debugging than capture_stderr=True
    """
    try:
        ffmpeg.run(
            stream,
            overwrite_output=overwrite
        )
    except ffmpeg.Error as e:
        print("FFmpeg failed.\n")
        if e.stderr:
            print(e.stderr.decode(errors="ignore"))
        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:
        print("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:
        print("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
        print(f"Provided index {provided} is not a valid audio stream.")

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

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

    try:
        choice = int(input("Select stream number (the list index [x]): "))
        return streams[choice][0]
    except (ValueError, IndexError):
        print("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: str):
    ext = output_path.lower().split(".")[-1]

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

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

    return "copy"

def get_vcodec(output_path: str):
    ext = output_path.lower().split(".")[-1]

    if ext == "mp4":
        return "libx264"

    return "copy"

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

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

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,
        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 [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,
        args.output,
        vcodec="copy",
        acodec=args.audio_codec,
        scodec=get_scodec(args.output),
        map='0',
        map_metadata=0
    )
    return [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,
        args.output,
        vcodec=args.video_codec,
        acodec="copy",
        scodec=get_scodec(args.output),
        map='0',
        map_metadata=0
    )
    return [out]

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

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

    Parameters:
    - sn: Disables subtitle recording.
    """
    inp = ffmpeg.input(args.input)
    out = ffmpeg.output(
        inp,
        args.output,
        sn=None,
        c="copy",
        map='0',
        map_metadata=0
    )
    return [out]

def cmd_remove_audio_track(args):
    """
    1. Removes a specific audio track from the video (Main Output).
    2. Saves that specific audio track as a standalone MP3 (Extra Output).
    """
    inp = ffmpeg.input(args.input)
    probe = ffmpeg.probe(args.input)

    remove_idx = select_audio_stream(args.input, args.audio_index)

    streams_to_keep = []
    extracted_audio = None

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

    video_out = ffmpeg.output(
        *streams_to_keep,
        args.output,
        vcodec="copy",
        acodec="copy",
        scodec= get_scodec(args.output),
    )

    outputs = [video_out]

    if args.output_audio:
        audio_out = ffmpeg.output(
            extracted_audio,
            args.output_audio,
            acodec='libmp3lame',
            audio_bitrate='192k'
        )
        outputs.append(audio_out)

    return 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(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,
        args.output,
        vcodec="libx264",
        crf=args.crf,
        acodec="copy",
        scodec=get_scodec(args.output)
    )
    return [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.
    """
    idx_filt = select_audio_stream(args.input, args.audio_index)
    inp = ffmpeg.input(args.input)

    probe = ffmpeg.probe(args.input)

    audio_streams = [
        s["index"]
        for s in probe["streams"]
        if s.get("codec_type") == "audio"
    ]
    filtered_audios = []
    for idx in audio_streams:
        audio = inp[str(idx)]
        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)
        else:

            filtered_audios.append(audio)
    info = ffmpeg.probe(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,
        *filtered_audios,
        *subs,
        args.output,
        vcodec="copy",
        acodec="libvorbis",
        scodec=get_scodec(args.output),
    )

    return [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)

    out = ffmpeg.output(
        inp,
        args.output,
        ss=start,
        t=duration,
        vcodec=get_vcodec(args.output),
        acodec="copy",
        scodec=get_scodec(args.output),
        map="0"
    )

    return [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)

    primary_audio = select_audio_stream(args.input, args.audio_index)
    probe = ffmpeg.probe(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(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,
        args.output,
        vcodec="copy",
        acodec="copy",
        scodec=get_scodec(args.output),
        **out_kwargs
    )

    return [out]

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

    script_dir = os.path.dirname(os.path.abspath(__file__))
    test_files = [ "test_auto.mkv", "test_auto.mp4","test_auto_no_subs.mkv", "test_auto_no_subs.mp4"]
    test_files = [ os.path.join(script_dir,file) for file in test_files ]
    test_dir = "/tmp/"


    RED = "\033[91m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    BLUE = "\033[94m"
    RESET = "\033[0m"

    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, test_dict: list=[]):
        print(f"TEST: {name}")
        # Run FFmpeg for this single command only
        outputs = cmd_func(args)

        if not outputs:
            print("No outputs returned.")
            return 1

        try:
            ffmpeg.run(
                outputs,
                overwrite_output=True,
                quiet=True,
                capture_stderr=True
            )
        except ffmpeg.Error as e:
            print("FFmpeg failed:")
            if e.stderr:
                print(e.stderr.decode(errors="ignore"))
            return 1

        # Probe result file(s)
        out_files = []

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

        if hasattr(args, "output_audio") and args.output_audio:
            out_files.append(args.output_audio)


        for i, f in enumerate(out_files):
            if not os.path.exists(f):
                print(f"{red_tag("FAIL")} Missing output: {f}")
                continue

            try:
                ret_va = validate_ffprobe(f,test_dict[i])
            except ffmpeg.Error as e:
                print(f"{red_tag("FAIL")} Probe error: {f}")
                print(e)
                return 1

        return ret_va

    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"]
        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"]
        if not os.path.exists(test_file):
            print("test file not found in script directory.")
            sys.exit(1)

        return video_streams,audio_streams,subtitle_streams

    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]
        }
        """

        info = ffmpeg.probe(file_path)
        streams = info.get("streams", [])

        video_streams = [s for s in streams if s.get("codec_type") == "video"]
        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"]
        ret = 0

        # ---------- COUNT CHECKS ----------
        if "video_count" in expected:
            if len(video_streams) != expected["video_count"]:
                print(f"{red_tag("VIDEO COUNT FAIL")} expected {expected['video_count']}, got {len(video_streams)}")
                ret = 1

        if "audio_count" in expected:
            if len(audio_streams) != expected["audio_count"]:
                print(f"{red_tag("AUDIO COUNT FAIL")} expected {expected['audio_count']}, got {len(audio_streams)}")
                ret = 1

        if "subtitle_count" in expected:
            if len(subtitle_streams) != expected["subtitle_count"]:
                print(f"{red_tag("SUBTITLE COUNT FAIL")} expected {expected['subtitle_count']}, got {len(subtitle_streams)}")
                ret = 1

        # ---------- 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:
                print(f"{red_tag("AUDIO CODEC FAIL")}\nExpected: {expected_codecs}\nGot: {actual}")
                ret = 1

        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:
                print(f"{red_tag("VIDEO CODEC FAIL")}\nExpected: {expected_codecs}\nGot: {actual}")
                ret = 1

        # ---------- 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:
                print(f"{red_tag("DURATION FAIL")} expected ~{expected_duration} (±{tolerance}), got {actual_duration}s")
                ret = 1

        print(f"{green_tag("PASS")} {file_path}")
        return ret

    class Args:
        pass

    ret = 0
    # ---------- INFO ----------
    #args = Args()
    #args.input = test_file_mp4
    #test_dict = {}
    #run_and_probe("info", cmd_info, args, test_dict)
    #args.input = test_file_mkv
    #ret += run_and_probe("info", cmd_info, args, test_dict)

    # ---------- DOWNSIZE ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_downsize" + os.path.splitext(test_file)[1])
        args.width = 640
        args.height = 360
        args.crf = 28
        test_dict = [{"video_count": len(video_streams), "audio_count": len(audio_streams),
                      "subtitle_count": len(subtitle_streams), "video_codecs": ["h264"]}]
        ret += run_and_probe("downsize", cmd_downsize, args, test_dict)

    # ---------- REMOVE AUDIO TRACK ----------
    args = Args()
    args.audio_index = 1

    for test_file in test_files:
        video_streams,audio_streams,subtitle_streams = get_streams(test_file)
        args.input = test_file
        test_dict = [{"video_count": len(video_streams), "audio_count": 1, "subtitle_count": len(subtitle_streams)},
                     {"video_count": 0, "audio_count": 1, "subtitle_count": 0}]
        args.output = os.path.join(test_dir, "out_remove_audio" + os.path.splitext(test_file)[1])
        args.output_audio = os.path.join(test_dir, "extracted_audio_" + os.path.splitext(test_file)[1][1:] + ".mp3")
        ret += run_and_probe("remove-audio-track", cmd_remove_audio_track, args, test_dict)

    # ---------- CONVERT AV ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_convert_av" + os.path.splitext(test_file)[1])
        args.video_codec = "libx264"
        args.audio_codec = "aac"
        args.preset = "fast"
        args.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 += run_and_probe("convert-av", cmd_convert_av, args, test_dict)


    # ---------- CONVERT AUDIO ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_convert_audio" + os.path.splitext(test_file)[1])
        args.audio_codec = "aac"
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams),
                     "subtitle_count": len(subtitle_streams), "audio_codecs": ["aac"]}]
        ret += run_and_probe("convert-audio", cmd_convert_audio, args, test_dict)

    # ---------- CONVERT VIDEO ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_convert_video" + os.path.splitext(test_file)[1])
        args.video_codec = "libx264"
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams),
                      "subtitle_count": len(subtitle_streams), "video_codecs": ["h264"]}]
        ret += run_and_probe("convert-video", cmd_convert_video, args, test_dict)

    # ---------- REMOVE SUBS ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams), "subtitle_count": 0}]
        args.output = os.path.join(test_dir, "out_no_subs" + os.path.splitext(test_file)[1])
        ret += run_and_probe("remove-subs", cmd_remove_subs, args, test_dict)

    # ---------- FILTER AUDIO ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_filter_audio" + os.path.splitext(test_file)[1])
        args.audio_index = 1
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams),
                      "subtitle_count": len(subtitle_streams), "audio_codecs": ["vorbis"]}]
        ret += run_and_probe("filter-audio", cmd_filter_audio, args, test_dict)

    # ---------- CUT ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_cut" + os.path.splitext(test_file)[1])
        args.start = "0:0:5"
        args.duration = "10s"
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams),
                     "subtitle_count": len(subtitle_streams), "duration": "10s"}]
        ret += run_and_probe("cut", cmd_cut_segment, args, test_dict)

    # ---------- SET PRIMARY AUDIO ----------
    args = Args()
    for test_file in test_files:
        video_streams, audio_streams, subtitle_streams = get_streams(test_file)
        args.input = test_file
        args.output = os.path.join(test_dir, "out_primary_audio" + os.path.splitext(test_file)[1])
        args.audio_index = 2
        test_dict = [ {"video_count":  len(video_streams), "audio_count": len(audio_streams),
                     "subtitle_count": len(subtitle_streams)}]
        ret += run_and_probe("set-primary-audio", cmd_set_primary_audio, args, test_dict)

    if ret == 0:
        print(green_tag("ALL TESTS SUCCEEDED"))
    else:
        print(red_tag("SOME TESTS FAILED"))

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

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

    # ----------------------------
    # BASE ARGS (shared)
    # ----------------------------
    base_p = argparse.ArgumentParser(add_help=False)
    base_p.add_argument("input", nargs="?", help="Path to input media file")
    base_p.add_argument("-o", "--output", help="Path for output file")

    sub = parser.add_subparsers(dest="command")  # command is optional now

    # ----------------------------
    # INFO
    # ----------------------------
    p_info = sub.add_parser("info", aliases=["i"], parents=[base_p], help="Inspect media metadata")
    p_info.set_defaults(func=cmd_info, command="info")

    # TEST
    p_test = sub.add_parser("test", aliases=["t"], help="Run test function")
    p_test.set_defaults(func=cmd_test, command="test")

    # ----------------------------
    # CONVERT AV
    # ----------------------------
    p_av = sub.add_parser("convert-av", aliases=["cav"], parents=[base_p], help="Full re-encode of video/audio")
    p_av.add_argument("--video-codec", default="libx264",
                      help="Video codec to use (default: libx264, e.g. libx265, vp9)")
    p_av.add_argument("--audio-codec", default="aac", help="Audio codec to use (default: aac, e.g. mp3, opus)")
    p_av.add_argument("--preset", default="medium",
                      help="Encoding speed/quality preset (default: medium, e.g. ultrafast, fast, slow, veryslow)")
    p_av.add_argument("--quality", type=int,
                      help="Output quality level (e.g. CRF value or bitrate depending on encoder)")
    p_av.set_defaults(func=cmd_convert_av, command="convert-av")

    # ----------------------------
    # CONVERT AUDIO
    # ----------------------------
    p_a = sub.add_parser("convert-audio", aliases=["ca"], parents=[base_p], help="Change audio format only")
    p_a.add_argument("--audio-codec", default="aac", help="Audio codec to use (default: aac, e.g. mp3, opus)")
    p_a.set_defaults(func=cmd_convert_audio, command="convert-audio")

    # ----------------------------
    # CONVERT VIDEO
    # ----------------------------
    p_v = sub.add_parser("convert-video", aliases=["cv"], parents=[base_p], help="Change video format only")
    p_v.add_argument("--video-codec", default="libx264",
                      help="Video codec to use (default: libx264, e.g. libx265, vp9)")
    p_v.set_defaults(func=cmd_convert_video, command="convert-video")

    # ----------------------------
    # REMOVE SUBS
    # ----------------------------
    p_rs = sub.add_parser("remove-subs", aliases=["rs"], parents=[base_p], help="Strip all subtitle streams (-sn)")
    p_rs.set_defaults(func=cmd_remove_subs, command="remove-subs")

    # ----------------------------
    # REMOVE AUDIO TRACK
    # ----------------------------
    p_ra = sub.add_parser("remove-audio", aliases=["ra"], parents=[base_p], help="Remove one audio track")
    p_ra.add_argument("-oa", "--output-audio")
    p_ra.add_argument("-a", "--audio-index", type=int)
    p_ra.set_defaults(func=cmd_remove_audio_track, command="remove-audio")

    # ----------------------------
    # DOWNSIZE
    # ----------------------------
    p_ds = sub.add_parser("downsize", aliases=["d"], parents=[base_p], help="Scale video dimensions (-vf scale)")
    p_ds.add_argument("-w", "--width", type=int, help="width: eg 1080")
    p_ds.add_argument("-he", "--height", type=int, help="height: eg 720")
    p_ds.add_argument("-c" , "--crf", type=int, help="crf: eg 23")
    p_ds.set_defaults(func=cmd_downsize, command="downsize")

    # ----------------------------
    # FILTER AUDIO
    # ----------------------------
    p_fa = sub.add_parser("filter-audio", aliases=["fa"], parents=[base_p], help="Apply audio enhancement chain")
    p_fa.add_argument("-a", "--audio-index", type=int)
    p_fa.set_defaults(func=cmd_filter_audio, command="filter-audio")

    # ----------------------------
    # CUT
    # ----------------------------
    p_c = sub.add_parser("cut", aliases=["c"], parents=[base_p], help="Cut X seconds out of video")
    p_c.add_argument("--start")
    p_c.add_argument("--duration")
    p_c.set_defaults(func=cmd_cut_segment, command="cut")

    # ----------------------------
    # SET PRIMARY AUDIO
    # ----------------------------
    p_pa = sub.add_parser("set-primary-audio", aliases=["spa"], parents=[base_p], help="Set primary audio track")
    p_pa.add_argument("-a", "--audio-index", type=int)
    p_pa.set_defaults(func=cmd_set_primary_audio, command="set-primary-audio")

    # =========================================================
    # PARSE
    # =========================================================
    args = parser.parse_args()

    # ----------------------------
    # COMMAND VALIDATION
    # ----------------------------
    if not hasattr(args, "func"):
        print("\nAvailable commands:")

        seen = set()
        for name, parser in sub.choices.items():
            if id(parser) in seen:
                continue
            seen.add(id(parser))
            print(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 = ensure(getattr(args, "input", None), f"Input file ({Path.cwd()})")
        if not Path(args.input).is_file():
            parser.error("Input file does not exist!")

    if args.command in REQUIRES_OUTPUT and not getattr(args, "output", None):
        args.output = ensure(getattr(args, "output", None), "Output file")
    # ----------------------------
    # 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", "set-primary-audio"]:
        args.audio_index = ensure(getattr(args, "audio_index", None), "Audio index")

    # ----------------------------
    # EXECUTE COMMAND
    # ----------------------------
    try:
        outputs = args.func(args)

        if outputs is None:
            return

        if not isinstance(outputs, list):
            outputs = [outputs]

        if args.command != "test":
            ffmpeg.run(
                ffmpeg.merge_outputs(*outputs),
                overwrite_output=True
            )

    except ffmpeg.Error as e:
        print("FFmpeg failed.\n")
        if e.stderr:
            print(e.stderr.decode(errors="ignore"))
        sys.exit(1)

    sys.exit(0)

if __name__ == "__main__":
    main()