#!/usr/bin/env python3
"""
Prepend N seconds of silence to audio files using ffmpeg as a subprocess.

- Primary path: concat demuxer + -c copy (avoids re-encoding main audio)
- Fallback: filter_complex concat (re-encodes output) when stream-copy can't be used
- Writes outputs to <folder>/output/
- Optional recursion, skip-existing, and extension filtering

Usage:
    python extend_audio.py /path/to/folder [duration_seconds]
Examples:
    python extend_audio.py "/folder_that_contain_many_audio" 1
    python extend_audio.py "/folder_that_contain_many_audio" 1 --recursive --skip-existing
"""

import os
import sys
import json
import shlex
import subprocess
from pathlib import Path

# ------------------- helpers -------------------
def check_binary(name: str):
    if shutil.which(name) is None:
        print(f"Error: '{name}' not found on PATH. Install it and try again.")
        sys.exit(1)

def run_probe(args):
    return json.loads(subprocess.run(args, capture_output=True, text=True, check=True).stdout)

def get_audio_info(file_path):
    """
    Probe first audio stream; return codec_name, sample_rate, channels, channel_layout.
    Handle files without audio gracefully.
    """
    try:
        data = run_probe([
            "ffprobe", "-v", "error",
            "-select_streams", "a:0",
            "-show_entries", "stream=codec_name,channels,sample_rate,channel_layout,bit_rate",
            "-of", "json", file_path
        ])
        s = data["streams"][0]
        codec = s.get("codec_name")
        channels = int(s.get("channels", 2))
        sample_rate = int(s.get("sample_rate", 48000))
        layout = s.get("channel_layout") or ("mono" if channels == 1 else "stereo")
        bitrate = s.get("bit_rate")
        return {
            "codec": codec,
            "channels": channels,
            "sample_rate": sample_rate,
            "layout": layout,
            "bitrate": bitrate,
            "has_audio": True
        }
    except Exception:
        # If the container has no audio stream, still allow generating silence only.
        return {
            "codec": "aac",
            "channels": 2,
            "sample_rate": 48000,
            "layout": "stereo",
            "bitrate": None,
            "has_audio": False
        }

# ------------------- encoders -------------------
AUDIO_ENCODER_MAP = {
    "aac": ("aac", ["-b:a", "192k"]),
    "mp3": ("libmp3lame", ["-b:a", "192k"]),
    "libmp3lame": ("libmp3lame", ["-b:a", "192k"]),
    "opus": ("libopus", ["-b:a", "160k"]),
    "vorbis": ("libvorbis", ["-q:a", "4"]),
    "ac3": ("ac3", ["-b:a", "384k"]),
    "eac3": ("eac3", ["-b:a", "384k"]),
    "flac": ("flac", []),
    "alac": ("alac", []),
    "pcm_s16le": ("pcm_s16le", []),
    "wav": ("pcm_s16le", []),  # convenience
}

def pick_audio_encoder(codec_name):
    # Normalize codec name to the encoder we want to use
    if not codec_name:
        return "aac", ["-b:a", "192k"]
    enc, extra = AUDIO_ENCODER_MAP.get(codec_name, ("aac", ["-b:a", "192k"]))
    return enc, extra

# ------------------- build intro (silence) -------------------
def build_silence_cmd(info, seconds, intro_path):
    """
    Create a short silent file encoded with the SAME codec + params as the input,
    so we can try concat demuxer + -c copy.
    """
    enc, extra = pick_audio_encoder(info["codec"])
    ch = info["channels"]
    sr = info["sample_rate"]
    layout = info["layout"]

    # Important: tell ffmpeg that anullsrc is a filter input (-f lavfi). [3](https://stackoverflow.com/questions/42147512/ffmpeg-adding-silence-struggling-to-use-i-anullsrc-option)
    cmd = [
        "ffmpeg",
        "-hide_banner", "-loglevel", "error", "-y",
        "-f", "lavfi", "-t", str(seconds),
        "-i", f"anullsrc=channel_layout={layout}:sample_rate={sr}:d={seconds}",
        "-c:a", enc, *extra,
        "-ac", str(ch), "-ar", str(sr),
        intro_path
    ]
    return cmd

# ------------------- concat methods -------------------
def concat_demuxer_copy(intro_path, input_path, output_path):
    """
    Concatenate via the concat demuxer and stream copy. Requires same codec & params. [1](https://trac.ffmpeg.org/wiki/Concatenate)
    """
    # Build a temporary list file
    list_file = Path(output_path).with_suffix(".concat.txt")
    list_file.write_text(f"file '{intro_path}'\nfile '{input_path}'\n", encoding="utf-8")

    cmd = [
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
        "-f", "concat", "-safe", "0",
        "-i", str(list_file),
        "-c", "copy",
        output_path
    ]
    print("Concat demuxer command:\n " + " ".join(shlex.quote(x) for x in cmd))
    proc = subprocess.run(cmd)
    try:
        list_file.unlink()
    except Exception:
        pass
    return proc.returncode

def concat_filter_reencode(input_path, seconds, output_path, info):
    """
    Fallback: use filter_complex concat to prepend the silent intro and re-encode output. [1](https://trac.ffmpeg.org/wiki/Concatenate)
    """
    enc, extra = pick_audio_encoder(info["codec"])
    ch = info["channels"]
    sr = info["sample_rate"]
    layout = info["layout"]

    # Build filter path: [silence][input] concat to 1 audio stream. [4](https://superuser.com/questions/579008/add-1-second-of-silence-to-audio-through-ffmpeg)
    cmd = [
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
        "-f", "lavfi", "-t", str(seconds),
        "-i", f"anullsrc=channel_layout={layout}:sample_rate={sr}:d={seconds}",
        "-i", input_path,
        "-filter_complex", "[0:a][1:a]concat=n=2:v=0:a=1[a]",  # concat filter for audio
        "-map", "[a]",
        "-c:a", enc, *extra,
        "-ac", str(ch), "-ar", str(sr),
        output_path
    ]
    print("Concat filter (fallback) command:\n " + " ".join(shlex.quote(x) for x in cmd))
    return subprocess.run(cmd).returncode

# ------------------- batch logic -------------------
import shutil

SUPPORTED_EXTS = [".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg", ".opus"]

def find_audio_files(folder: Path, recursive: bool, exts):
    pattern = "**/*" if recursive else "*"
    exts_norm = {e.lower() for e in exts}
    return sorted([p for p in folder.glob(pattern) if p.is_file() and p.suffix.lower() in exts_norm])

def process_one(input_file: Path, out_dir: Path, seconds: float, skip_existing: bool):
    info = get_audio_info(str(input_file))
    output_file = out_dir / input_file.name
    intro_file = out_dir / (input_file.stem + "_intro" + input_file.suffix)

    if skip_existing and output_file.exists():
        print(f"Skip (exists): {output_file}")
        return

    # 1) Create silence intro encoded like the source
    intro_cmd = build_silence_cmd(info, seconds, str(intro_file))
    print("Intro command:\n " + " ".join(shlex.quote(x) for x in intro_cmd))
    rc = subprocess.run(intro_cmd).returncode
    if rc != 0:
        print(f"Failed to make intro for {input_file.name}")
        return

    # 2) Try demuxer (stream copy) first
    rc = concat_demuxer_copy(str(intro_file), str(input_file), str(output_file))
    if rc == 0:
        print(f"OK (copy): {input_file.name} -> {output_file.name}")
    else:
        print(f"Demuxer concat failed; trying filter fallback…")
        rc = concat_filter_reencode(str(input_file), seconds, str(output_file), info)
        if rc == 0:
            print(f"OK (re-encode): {input_file.name} -> {output_file.name}")
        else:
            print(f"FAILED: {input_file}")
    # Cleanup the intro file
    try:
        intro_file.unlink()
    except Exception:
        pass

def process_folder(root: Path, seconds: float, recursive: bool, skip_existing: bool, output_dir_name: str, exts):
    out_dir = root / output_dir_name
    out_dir.mkdir(parents=True, exist_ok=True)
    files = find_audio_files(root, recursive, exts)
    # Don't process anything in output/
    files = [f for f in files if out_dir not in f.parents]

    if not files:
        print("No matching audio files found.")
        return

    print(f"Found {len(files)} file(s). Output dir: {out_dir}")
    for f in files:
        try:
            process_one(f, out_dir, seconds, skip_existing)
        except KeyboardInterrupt:
            print("\nInterrupted.")
            sys.exit(130)
        except Exception as e:
            print(f"Error on '{f}': {e}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(
        description="Prepend N seconds of silence to audio files. Uses concat demuxer when possible; falls back to re-encode."
    )
    parser.add_argument("folder", help="Folder containing audio files")
    parser.add_argument("seconds", nargs="?", type=float, default=1.0, help="Silence duration in seconds (default: 1.0)")
    parser.add_argument("--recursive", action="store_true", help="Process subfolders")
    parser.add_argument("--skip-existing", action="store_true", help="Skip if output already exists")
    parser.add_argument("--output-dir-name", default="output", help="Subfolder name for outputs (default: output)")
    parser.add_argument("--exts", nargs="+", default=SUPPORTED_EXTS, help="Extensions to include")
    args = parser.parse_args()

    check_binary("ffmpeg")
    check_binary("ffprobe")

    root = Path(args.folder).expanduser().resolve()
    if not root.exists() or not root.is_dir():
        print(f"Error: '{root}' is not a folder.")
        sys.exit(1)

    print(f"Extending audio in: {root} | silence: {args.seconds}s\n")
    process_folder(root, args.seconds, args.recursive, args.skip_existing, args.output_dir_name, args.exts)