#!/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
import subprocess
import tempfile
from pathlib import Path
from mutagen.easyid3 import EasyID3

from serverlibary import safe_rename
VIDEO_EXTS = {".mp4", ".mkv", ".avi", ".mov", ".webm",".mp3", ".m4a", ".m4b", ".opus"}

# ----------------------------
# OUTPUT ENFORCEMENT (clean + centralized)
# ----------------------------
REQUIRES_OUTPUT = {
    "chapters",
    "convert-av",
    "convert-audio",
    "convert-video",
    "remove-subs",
    "remove-audio",
    "extract-audio",
    "downsize",
    "filter-audio",
    "cut",
    "set-primary-audio",
    "metaset",
}
# ---------- 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_override_file(input_path, backup_files=False,silent=False):
    if input_path.is_file():
        input_path=input_path.parent
    for override_path in input_path.glob("*_override.*"):
        if not override_path.is_file():
            continue

        stem = override_path.stem
        suffix = override_path.suffix

        if not stem.endswith("_override"):
            continue

        base = stem.removesuffix("_override")

        original = input_path / f"{base}{suffix}"
        backup = input_path / f".backup_{base}{suffix}"
        new_name = input_path / f"{base}{suffix}"

        if not original.exists():
            print(f"Skip (missing original): {original}")
            continue

        if backup_files:
            # move original -> backup
            safe_rename(original, backup,silent=silent)

        # move override -> original name
        safe_rename(override_path, new_name,silent=silent)


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

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

    if loud:
        kwargs.update(
            capture_stdout=False,
            capture_stderr=True,
            quiet=True,
        )

    try:
        out, err = ffmpeg.run(stream, **kwargs)

        if loud and err:
            with open(LOG_FILE, "ab") as log:
                log.write(err)

    except ffmpeg.Error as e:
        if loud:
            with open(LOG_FILE, "ab") as log:
                log.write(e.stderr or b"")
            print(f"FFmpeg failed. See log: {LOG_FILE}")
        else:
            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"

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)

    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(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).
    """
    inp = ffmpeg.input(args.input)
    probe = ffmpeg.probe(args.input)

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

    streams_to_keep = []

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

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

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(args.input)

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

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

    extracted_audio = None

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

    out = ffmpeg.output(
        extracted_audio,
        output_audio,
        acodec='libmp3lame',
        audio_bitrate='192k'
    )
    return out

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 = []
    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,
        args.output,
        vcodec="copy",
        acodec="libvorbis",
        scodec=get_scodec(args.output)
    ).global_args(*metadata_args)


    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_show_metadata(args):
    """
    Prints metadata directly from ffprobe output.
    """

    probe = ffmpeg.probe(args.input)

    print("=== FORMAT METADATA ===")
    format_tags = probe.get("format", {}).get("tags", {})
    for k, v in format_tags.items():
        print(f"{k}: {v}")

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

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


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 = os.path.splitext(args.input)[1].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
        ]

        subprocess.run(cmd, check=True)
        return None  # Kid3 writes in-place

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

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

def cmd_apply_chapters(args):
    """
    Apply chapters 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)
        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
        ]

        subprocess.run(cmd, check=True)
    elif Path(args.input).suffix == ".mp3":
            filename = Path(args.input)
            print(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)

            print(f"{filename.name} -> {new_name}")
            audio = EasyID3(filename)
            audio["title"] = title
            audio["tracknumber"] = str(args.file_counter+1)
            audio.save()
            safe_rename(filename, new_path)

    return None

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)

        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)}]
        args.output = os.path.join(test_dir, "out_remove_audio" + os.path.splitext(test_file)[1])
        ret += run_and_probe("remove-audio-track", cmd_remove_audio_track, args, test_dict)

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

    for test_file in test_files:
        args.input = test_file
        test_dict = [{"video_count": 0, "audio_count": 1, "subtitle_count": 0}]
        args.output = os.path.join(test_dir, "extracted_audio_" + os.path.splitext(test_file)[1][1:] + ".mp3")
        ret += run_and_probe("remove-audio-track", cmd_extract_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)+1,
                      "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")
    base_p.add_argument("--override", action="store_true", help="Override input")
    base_p.add_argument("--backup", action="store_true", help="Backup original file when overriding")
    base_p.add_argument("--loud", action="store_false", help="Run ffmpeg loud")
    base_p.add_argument("--kidflag", action="store_true", help="Use kid3 instead of ffmpeg")
    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("-a", "--audio-index", type=int)
    p_ra.set_defaults(func=cmd_remove_audio_track, command="remove-audio")

    # ----------------------------
    # REMOVE AUDIO TRACK
    # ----------------------------
    p_ea = sub.add_parser("extract-audio", aliases=["ea"], parents=[base_p], help="Extract one audio track")
    p_ea.add_argument("-a", "--audio-index", type=int)
    p_ea.set_defaults(func=cmd_extract_audio_track, command="extract-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")

    # ----------------------------
    # CHAPTERS
    # ----------------------------
    p_ch = sub.add_parser("chapters",aliases=["ch"],parents=[base_p],help="Apply chapter metadata from ffmetadata file")
    p_ch.add_argument("-cf","--chapter-file",help="Path to ffmetadata chapters file (e.g. chapters.txt)")
    p_ch.set_defaults(func=cmd_apply_chapters, command="chapters")

    # ----------------------------
    # METADATA
    # ----------------------------

    # PRINT

    p_ch = sub.add_parser("metaprint", aliases=["mp"], parents=[base_p],
                          help="Show metadata")
    p_ch.set_defaults(func=cmd_show_metadata, command="metaprint")

    # SET

    p_ch = sub.add_parser("metaset",aliases=["ms"],parents=[base_p],help="Apply metadata key and value to input file")
    p_ch.add_argument("-k","--key",help="Metadata key name)")
    p_ch.add_argument("-v","--value",help="Metadata key value")
    p_ch.set_defaults(func=cmd_set_metadata, command="metaset")

    # =========================================================
    # 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()})")
        args.input = Path(args.input).resolve()
        if not Path(args.input).exists():
            parser.error("Input file does not exist!")
    if args.command in REQUIRES_OUTPUT and not getattr(args, "output", None):
        input_path = Path(args.input)
        if args.override:
            args.output = str(input_path.with_name(f"{input_path.stem}_override{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 = str(input_path.with_name(f"{input_path.stem}_{args.command}{input_path.suffix}"))
    print(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"]:
        args.audio_index = ensure(getattr(args, "audio_index", None), "Audio index")
        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 == "metaset":
        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)
        # ---- DIRECTORY ----
        if input_path.is_dir():
            counter=0
            sorted_files = sorted(input_path.iterdir(), key=natural_sort_key)
            for file in sorted_files:
                if file.suffix.lower() not in VIDEO_EXTS:
                    continue
                print(f"Processing File: {file}")
                args.input = str(file)
                if args.override:
                    out_path = file.with_name(f"{file.stem}_override{file.suffix}")
                else:
                    out_path = file.with_name(f"{file.stem}_{args.command}{file.suffix}")
                args.output = str(out_path)
                args.file_counter = counter
                outputs.append(args.func(args))
                counter+=1
        # ---- SINGLE FILE ----
        elif input_path.is_file():
            outputs.append(args.func(args))
        else:
            raise ValueError(f"Invalid input path: {args.input}")

        if any(item is None for item in outputs) or len(outputs) == 0:
            print("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.override:
            move_override_file(input_path, args.backup,args.loud)
    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()