Server-scripts/audiobook_metadata/combine_audiobookmedia_to_m4b.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
import subprocess
import tempfile
from pathlib import Path
import time
from typing import Any

RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
RESET = "\033[0m"
logfile=Path(f"media_{time.strftime("%Y%m%d_%H%M%S")}.log")

def log_message(message):
    with open(logfile, "a") as log:
        log.write(f"{message}\n")
    print(message)

def get_duration_ms(media: Path) -> int:
    duration = subprocess.check_output(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            str(media),
        ],
        text=True,
    ).strip()

    return int(float(duration) * 1000)

def ffmpeg_escape(path: Path) -> str:
    return str(path.resolve()).replace("'", r"'\''")

def create_m4b(root_path: Path, book_dir: Path) -> None:
    media_files = sorted(book_dir.glob("*.m[p,4][3,b,a]"))
    if not media_files:
        return

    output_file_final = book_dir / f"{book_dir.name}.m4b"

    if output_file_final.exists():
        log_message(f"{RED}[SKIP]{RESET} {output_file_final}")
        log_message("Book already present")
        return

    log_message(f"[CREATE] {output_file_final}")

    with (tempfile.TemporaryDirectory(dir=root_path.parent) as tmp):
        tmpdir = Path(tmp)

        output_file_tmp = tmpdir / f"{book_dir.name}.m4b"
        concat_file = tmpdir / "concat.txt"
        metadata_file = tmpdir / "chapters.txt"

        current_ms = 0

        with concat_file.open("w", encoding="utf-8") as cf, \
            metadata_file.open("w", encoding="utf-8") as mf:
            mf.write(";FFMETADATA1\n")

            for media in media_files:
                cf.write(f"file '{ffmpeg_escape(media)}'\n")

                duration_ms = get_duration_ms(media)

                mf.write("[CHAPTER]\n")
                mf.write("TIMEBASE=1/1000\n")
                mf.write(f"START={current_ms}\n")
                mf.write(f"END={current_ms + duration_ms}\n")
                mf.write(f"title={media.stem}\n")

                current_ms += duration_ms
        cmd_backup = [ "ffmpeg","-y","-f","concat","-safe","0",
                "-i",str(concat_file),"-i",str(metadata_file),
                # cover from first media
                "-i",str(media_files[0]),
                # audio from concat
                "-map","0:a",
                # cover image from first media
                "-map","2:v?",
               # copy artist/album/composer/etc from first media file
               "-map_metadata", "2",
               # copy chapters from ffmetadata
               "-map_chapters", "1",

               "-c:a","copy",
                "-c:a", "aac",
                "-b:a", "128k",
                # copy cover image without re-encoding
                "-c:v","copy","-disposition:v:0","attached_pic",
                str(output_file_tmp),
            ]
        cmd_default = ["ffmpeg", "-y", "-f", "concat", "-safe", "0",
               "-i", str(concat_file), "-i", str(metadata_file),
               # cover from first media
               "-i", str(media_files[0]),
               # audio from concat
               "-map", "0:a",
               # cover image from first media
               "-map", "2:v?",

               # copy artist/album/composer/etc from first media file
               "-map_metadata", "2",
               # copy chapters from ffmetadata
                "-map_chapters", "1",
               "-c:a", "copy",
               # copy cover image without re-encoding
               "-c:v", "copy", "-disposition:v:0", "attached_pic",
               str(output_file_tmp),
               ]
        with logfile.open("a", encoding="utf-8") as log:
            result = subprocess.run(
                cmd_default,
                stdout=log,
                stderr=subprocess.STDOUT,
                text=True,
            )
            if result.returncode == 234:
                log_message(f"{YELLOW}[WARNING]{RESET} {result.returncode} retrying with aac encoder")
                result = subprocess.run(
                    cmd_backup,
                    stdout=log,
                    stderr=subprocess.STDOUT,
                    text=True,
                )
        if result.returncode == 0:
            output_file_tmp.rename(output_file_final)
            for media in media_files:
                media.unlink()
            log_message(f"{GREEN}[SUCCESS]{RESET}, delete media files")
        else:
            log_message(f"{RED}[ERROR]{RESET} Code: {result.returncode}, Stderr: {result.stderr}")

def find_book_dirs(root: Path) -> tuple[list[Any], int]:
    """
    A book directory is any directory containing more than one
    audiobook file (.mp3, .m4a, .m4b)
    """

    media_files = (
        list(root.rglob("*.mp3"))
        + list(root.rglob("*.m4a"))
        + list(root.rglob("*.m4b"))
    )

    counts: dict[Path, int] = {}

    for file in media_files:
        if counts.get(file.parent, 0) > 1:
            continue
        counts[file.parent] = counts.get(file.parent, 0) + 1

    directories = []
    size = 0
    for directory, count in counts.items():
        if count > 1:
            directories.append(directory)
            size += sum(file.stat().st_size for pattern in ("*.mp3", "*.m4a", "*.m4b") for file in directory.glob(pattern))


    return sorted(directories), size

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "path",
        type=Path,
        help="Book folder or root folder",
    )

    args = parser.parse_args()

    target = args.path.resolve()

    if not target.exists():
        raise SystemExit(f"Path does not exist: {target}")

    # Direct book folder
    if any(target.glob("*.m[p,4][3,b]")):
        create_m4b(target.parent,target)
        return

    book_dir_set, library_size = find_book_dirs(target)

    for book_dir in book_dir_set:
        log_message(f"Found Media: {book_dir}")

    print("Proceed? y/n")
    name = input()
    if name != "y":
        log_message(f"Aborted")
        exit(0)
    size = 0
    # Root folder containing many books
    for book_dir in book_dir_set:
        size += sum(file.stat().st_size for pattern in ("*.mp3", "*.m4a", "*.m4b") for file in book_dir.glob(pattern))
        print(f"{size / 1024**3:07.2f}/{library_size/ 1024**3:07.2f} - {size/library_size*100:.2f}% ... ",end="")
        try:
            create_m4b(target,book_dir)
        except Exception as exc:
            log_message(f"{RED}[ERROR]{RESET} {book_dir}: {exc}")


if __name__ == "__main__":
    main()