#!/usr/bin/env python3

import argparse
import re
import subprocess
import tempfile
from pathlib import Path

EXTENSIONS = {
    ".mp3", ".m4a", ".wav", ".flac",
    ".aac", ".ogg", ".aiff", ".aif"
}


def run(command):
    subprocess.run(command, check=True)


def create_snippets(
    source,
    output,
    snippet_seconds=8,
    silence_seconds=4
):
    source = Path(source).expanduser().resolve()
    output = Path(output).expanduser().resolve()

    if not source.is_dir():
        raise SystemExit(f"Folder not found: {source}")

    if snippet_seconds <= 0 or silence_seconds < 0:
        raise SystemExit("Invalid snippet or silence duration")

    tracks = sorted(
        (
            p for p in source.iterdir()
            if p.is_file()
            and p.suffix.lower() in EXTENSIONS
            and p != output
        ),
        key=lambda p: [
            int(part) if part.isdigit() else part.lower()
            for part in re.split(r"(\d+)", p.name)
        ]
    )

    if not tracks:
        raise SystemExit("No supported audio files found")

    print(f"Found {len(tracks)} tracks")

    output.parent.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as temp:
        temp = Path(temp)
        segments = []

        for index, track in enumerate(tracks, start=1):
            print(f"[{index}/{len(tracks)}] {track.name}")

            snippet = temp / f"snippet_{index:04d}.wav"

            run([
                "ffmpeg",
                "-hide_banner",
                "-loglevel", "error",
                "-y",
                "-i", str(track),
                "-t", str(snippet_seconds),
                "-vn",
                "-ac", "2",
                "-ar", "44100",
                "-c:a", "pcm_s16le",
                str(snippet)
            ])

            segments.append(snippet)

            if index < len(tracks) and silence_seconds > 0:
                silence = temp / f"silence_{index:04d}.wav"

                run([
                    "ffmpeg",
                    "-hide_banner",
                    "-loglevel", "error",
                    "-y",
                    "-f", "lavfi",
                    "-i",
                    "anullsrc=r=44100:cl=stereo",
                    "-t", str(silence_seconds),
                    "-c:a", "pcm_s16le",
                    str(silence)
                ])

                segments.append(silence)

        playlist = temp / "playlist.txt"

        playlist.write_text(
            "".join(
                f"file '{str(segment)}'\n"
                for segment in segments
            )
        )

        print("Combining snippets...")

        run([
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-y",
            "-f", "concat",
            "-safe", "0",
            "-i", str(playlist),
            "-c:a", "libmp3lame",
            "-b:a", "192k",
            str(output)
        ])

    print()
    print("Complete!")
    print(f"Output: {output}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("source")
    parser.add_argument("output")

    parser.add_argument(
        "--snippet",
        type=float,
        default=5
    )

    parser.add_argument(
        "--silence",
        type=float,
        default=5
    )

    args = parser.parse_args()

    create_snippets(
        args.source,
        args.output,
        args.snippet,
        args.silence
    )