path: root/converter/audio.py

"""Audio assembly: combining chunks, speed adjustment, cleanup."""

import logging
import re
import shutil
import subprocess
import traceback
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from . import config

logger = logging.getLogger(__name__)


def atempo_filters(speed: float) -> str:
    """Return a comma-joined ffmpeg ``atempo`` filter chain for ``speed``.

    ``atempo`` accepts 0.5..2.0 per filter; values outside that range are
    handled by chaining multiple filters. Returns "" when ``speed`` is 1.0.
    """
    if speed <= 0:
        raise ValueError(f"Speed must be a positive number, got {speed}")
    if abs(speed - 1.0) < 1e-6:
        return ""
    remaining = float(speed)
    chain = []
    while remaining > 2.0:
        chain.append("atempo=2.0")
        remaining /= 2.0
    while remaining < 0.5:
        chain.append("atempo=0.5")
        remaining /= 0.5
    chain.append(f"atempo={remaining:g}")
    return ",".join(chain)


def speed_export_params(speed: float) -> List[str]:
    """Return ffmpeg filter args for pitch-preserving speed adjustment."""
    filters = atempo_filters(speed)
    if not filters:
        return []
    return ["-filter:a", filters]


def _concat_escape(path: str) -> str:
    """Escape a path for use inside single quotes in an ffmpeg concat list."""
    return path.replace("'", "'\\''")


def _encode_args(output_format: str) -> List[str]:
    """Return ffmpeg output codec/bitrate args for the requested container."""
    if output_format == "m4b":
        return ["-c:a", "aac", "-b:a", config.AUDIO_BITRATE]
    if output_format == "ogg":
        return ["-c:a", "libvorbis", "-b:a", config.AUDIO_BITRATE]
    if output_format == "flac":
        return ["-c:a", "flac"]
    if output_format == "wav":
        # Lossless intermediate for per-chapter scratch audio; avoids
        # generational loss when the final m4b re-encodes to AAC.
        return ["-c:a", "pcm_s16le"]
    return ["-b:a", config.AUDIO_BITRATE]


_brand_supported: Optional[bool] = None


def _detect_brand_support() -> bool:
    """Check whether the local ffmpeg muxer accepts the ``-brand`` option."""
    if shutil.which("ffmpeg") is None:
        return False
    try:
        proc = subprocess.run(
            ["ffmpeg", "-hide_banner", "-h", "muxer=ipod"],
            capture_output=True, text=True, timeout=15,
        )
    except (OSError, subprocess.TimeoutExpired):
        return False
    return "-brand" in proc.stdout


def _m4b_container_args() -> List[str]:
    """Return per-output container flags for m4b files.

    ``+faststart`` moves the ``moov`` index to the front of the file so
    streaming players (and naive linear readers like web players) can index
    it; without it they may misreport the duration or refuse the file.
    The ``M4B `` major brand identifies the file as an audiobook to
    players that sniff brands instead of trusting the extension (the ffmpeg
    default brand for .m4b is ``M4A ``).
    """
    global _brand_supported
    if _brand_supported is None:
        _brand_supported = _detect_brand_support()
    args = ["-movflags", "+faststart"]
    if _brand_supported:
        args += ["-brand", "M4B "]
    return args


def build_concat_command(concat_list: Path, output_path: Path, output_format: str,
                         speed: float = 1.0, speed_path: Optional[Path] = None) -> List[str]:
    """Build the ffmpeg command that concatenates chunk audio into a book file."""
    encode = _encode_args(output_format)
    container = _m4b_container_args() if output_format == "m4b" else []
    filters = atempo_filters(speed)
    if filters:
        if speed_path is None:
            raise ValueError("speed_path is required when speed is not 1.0")
        return [
            "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
            "-filter_complex",
            f"[0:a]split=2[base][spd];[spd]{filters}[spdout]",
            "-map", "[base]", *encode, *container, str(output_path),
            "-map", "[spdout]", *encode, *container, str(speed_path),
        ]
    return [
        "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
        *encode, *container, str(output_path),
    ]


def build_m4b_chapters_command(concat_list: Path, metadata_file: Path, output_path: Path,
                               speed: float = 1.0, speed_path: Optional[Path] = None,
                               speed_metadata_file: Optional[Path] = None) -> List[str]:
    """Build the ffmpeg command that assembles chapter audio into one m4b.

    Chapter metadata inputs are bound to their outputs with explicit
    ``-map_chapters`` so the normal-speed and speed-adjusted copies each get
    their own (correctly scaled) chapter markers.
    """
    encode = _encode_args("m4b")
    container = _m4b_container_args()
    if speed_path is not None:
        return [
            "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
            "-i", str(metadata_file), "-i", str(speed_metadata_file),
            "-filter_complex",
            f"[0:a]split=2[base][spd];[spd]{atempo_filters(speed)}[spdout]",
            "-map", "[base]", "-map_metadata", "1", "-map_chapters", "1",
            *encode, *container, str(output_path),
            "-map", "[spdout]", "-map_metadata", "2", "-map_chapters", "2",
            *encode, *container, str(speed_path),
        ]
    return [
        "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
        "-i", str(metadata_file),
        "-map", "0:a", "-map_metadata", "1", "-map_chapters", "1",
        *encode, *container, str(output_path),
    ]


_DURATION_WARN_TOLERANCE = 0.05  # warn when output duration drifts >5%
_DURATION_FAIL_TOLERANCE = 0.25  # fail when output duration drifts >25%


def verify_output_duration(path: Path, expected_ms: int) -> bool:
    """Sanity-check an assembled file's duration against the expected total.

    Catches corrupt assembly (truncated concat, bogus container metadata)
    before the file reaches audiobook players. Duration drift beyond the
    warn tolerance is logged; drift beyond the fail tolerance is an error
    and the output is treated as broken. Returns True when unverifiable.
    """
    if expected_ms <= 0:
        return True
    actual_ms = probe_duration_ms(path)
    if actual_ms <= 0:
        logger.warning("Could not verify duration of %s (ffprobe failed)", path)
        return True
    drift = abs(actual_ms - expected_ms) / expected_ms
    if drift > _DURATION_FAIL_TOLERANCE:
        logger.error(
            "Duration mismatch for %s: expected ~%.1fs, got %.1fs (%.0f%% off); output is likely corrupt",
            path.name, expected_ms / 1000.0, actual_ms / 1000.0, drift * 100.0,
        )
        return False
    if drift > _DURATION_WARN_TOLERANCE:
        logger.warning(
            "Duration drift for %s: expected ~%.1fs, got %.1fs (%.0f%% off)",
            path.name, expected_ms / 1000.0, actual_ms / 1000.0, drift * 100.0,
        )
    return True


def _collect_chunk_files(total_chunks: int,
                         chunk_results: Optional[Dict[int, Optional[Path]]] = None
                         ) -> Tuple[List[Path], List[int]]:
    """Resolve chunk audio files in book order.

    When ``chunk_results`` is provided (chunk number -> path written, or None
    for a failed chunk), the recorded paths are used exactly as-is so stale
    files from a previous chapter can never leak in. Without it, the chunks
    folder is globbed per index (legacy discovery).
    """
    chunk_files: List[Path] = []
    missing: List[int] = []
    for i in range(1, total_chunks + 1):
        if chunk_results is not None:
            recorded = chunk_results.get(i)
            if recorded is not None and Path(recorded).exists():
                chunk_files.append(Path(recorded))
            else:
                missing.append(i)
            continue
        matches = sorted(config.CHUNKS_FOLDER.glob(f"chunk_{i:04d}.*"))
        if matches:
            chunk_files.append(matches[0])
        else:
            missing.append(i)
    return chunk_files, missing


def combine_chunks(total_chunks: int, output_path: Path,
                   chunk_results: Optional[Dict[int, Optional[Path]]] = None,
                   speed: float = 1.0, output_format: str = "mp3",
                   intermediate: bool = False) -> bool:
    """Combine audio chunks into the final audiobook using ffmpeg's concat demuxer.

    ``chunk_results`` maps chunk numbers to the audio file each chunk produced
    (None for failed chunks); failed and missing chunks are skipped. When
    ``speed`` differs from 1.0, an additional speed-adjusted copy is written
    next to the normal-speed file. Chunks are streamed by ffmpeg, so the whole
    book is never held in memory. Set ``intermediate`` for scratch chapter
    audio on the way to a larger output (e.g. a chaptered m4b) so save
    messages don't present it as the final audiobook.
    """
    if shutil.which("ffmpeg") is None or shutil.which("ffprobe") is None:
        logger.error("ffmpeg and ffprobe are required to combine audio chunks (install ffmpeg)")
        return False

    chunk_files, missing_chunks = _collect_chunk_files(total_chunks, chunk_results)

    if not chunk_files:
        logger.error("No valid chunks found")
        return False

    if missing_chunks:
        logger.warning("Missing chunks: %s", missing_chunks)

    concat_list = config.CHUNKS_FOLDER / "_concat_list.txt"
    try:
        with open(concat_list, "w", encoding="utf-8") as list_file:
            for chunk_file in chunk_files:
                list_file.write(f"file '{_concat_escape(str(chunk_file))}'\n")

        speed_path = None
        if atempo_filters(speed):
            speed_path = output_path.with_name(f"{output_path.stem}_{speed:g}{output_path.suffix}")
        cmd = build_concat_command(concat_list, output_path, output_format,
                                   speed=speed, speed_path=speed_path)

        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode != 0:
            logger.error("ffmpeg failed: %s", proc.stderr[-2000:])
            return False

        # Verify the assembled duration against the sum of chunk durations
        # so corrupt output is caught before it reaches audiobook players.
        # A failed probe returns 0, which would deflate the expected total
        # and falsely fail the check, so unverifiable sums skip it.
        durations = [probe_duration_ms(chunk_file) for chunk_file in chunk_files]
        if any(duration <= 0 for duration in durations):
            logger.warning("Could not probe every chunk duration; skipping duration verification")
            duration_ok = True
        else:
            expected_ms = sum(durations)
            duration_ok = verify_output_duration(output_path, expected_ms)
            if speed_path is not None:
                duration_ok = verify_output_duration(speed_path, int(expected_ms / speed)) and duration_ok
        if not duration_ok:
            return False

        if intermediate:
            logger.info("Chapter audio saved (intermediate): %s (%d/%d chunks)",
                        output_path, len(chunk_files), total_chunks)
            print(f"[INFO] Saved chapter audio (intermediate): {output_path.name} "
                  f"({len(chunk_files)}/{total_chunks} chunks)")
        else:
            logger.info("Audiobook saved: %s (%d/%d chunks)", output_path, len(chunk_files), total_chunks)
            print(f"[INFO] Saved audiobook: {output_path.name} ({len(chunk_files)}/{total_chunks} chunks)")

        if speed_path is not None:
            logger.info("Saved speed-adjusted audiobook (%gx): %s", speed, speed_path)
            print(f"[INFO] Saved speed-adjusted audiobook: {speed_path.name} ({speed:g}x)")

        if missing_chunks:
            print(f"[WARNING] Missing chunks: {missing_chunks}")

        return True

    except FileNotFoundError:
        logger.error("ffmpeg not found on PATH (install ffmpeg and try again)")
        return False
    except Exception as exc:
        logger.error("Failed to combine chunks: %s", exc)
        logger.error(traceback.format_exc())
        return False
    finally:
        try:
            concat_list.unlink(missing_ok=True)
        except OSError:
            pass


def cleanup_chunks() -> None:
    """Remove temporary chunk and chapter files from the scratch folder."""
    try:
        chunk_count = 0
        for pattern in ("chunk_*", "chapter_*"):
            for chunk_file in config.CHUNKS_FOLDER.glob(pattern):
                try:
                    if chunk_file.is_file():
                        chunk_file.unlink()
                        chunk_count += 1
                except Exception as exc:
                    logger.warning("Failed to delete %s: %s", chunk_file, exc)

        if chunk_count > 0:
            logger.info("Cleaned up %d chunk files", chunk_count)
            print(f"[INFO] Cleaned up {chunk_count} chunk files")
    except Exception as exc:
        logger.warning("Cleanup failed: %s", exc)


def probe_duration_ms(path: Path) -> int:
    """Return audio duration in milliseconds using ffprobe."""
    try:
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=noprint_wrappers=1:nokey=1", str(path)],
            capture_output=True, text=True, timeout=30,
        )
    except subprocess.TimeoutExpired:
        logger.warning("ffprobe timed out for %s", path)
        return 0
    if result.returncode != 0:
        logger.warning("ffprobe failed for %s: %s", path, result.stderr[-200:])
        return 0
    try:
        return max(0, int(round(float(result.stdout.strip()) * 1000.0)))
    except ValueError:
        logger.warning("Could not parse ffprobe duration for %s", path)
        return 0


def _escape_ffmetadata_value(value: str) -> str:
    """Escape a metadata value for ffmpeg's FFMETADATA format.

    Backslash and the structural characters ``=``, ``;`` and ``#`` must be
    backslash-escaped; line breaks would corrupt the file and are collapsed
    to spaces.
    """
    value = value.replace("\\", "\\\\")
    value = re.sub(r"[\r\n]+", " ", value)
    return re.sub(r"[=;#]", r"\\\g<0>", value)


def build_ffmetadata(chapters: List[tuple], path: Path) -> None:
    """Write an ffmpeg FFMETADATA file with ``[CHAPTER]`` entries.

    ``chapters`` is a list of ``(start_ms, end_ms, title)`` tuples.
    """
    with open(path, "w", encoding="utf-8") as meta_file:
        meta_file.write(";FFMETADATA1\n")
        for start_ms, end_ms, title in chapters:
            meta_file.write("[CHAPTER]\n")
            meta_file.write("TIMEBASE=1/1000\n")
            meta_file.write(f"START={int(start_ms)}\n")
            meta_file.write(f"END={int(end_ms)}\n")
            meta_file.write(f"title={_escape_ffmetadata_value(title)}\n")


def combine_chapters_to_m4b(chapter_files: List[Path], titles: List[str],
                            output_path: Path, speed: float = 1.0) -> bool:
    """Concatenate per-chapter audio into a single m4b with embedded chapter markers.

    Chapter start/end times are derived from each chapter file's duration and
    written as ffmpeg chapter metadata. When ``speed`` differs from 1.0, a
    speed-adjusted copy (with rescaled chapter markers) is written alongside
    the normal-speed file.
    """
    if shutil.which("ffmpeg") is None or shutil.which("ffprobe") is None:
        logger.error("ffmpeg and ffprobe are required to build an m4b with chapters")
        return False

    if not chapter_files:
        logger.error("No chapter files provided")
        return False

    concat_list = config.CHUNKS_FOLDER / "_concat_list.txt"
    metadata_file = config.CHUNKS_FOLDER / "_chapters.txt"
    speed_metadata_file = config.CHUNKS_FOLDER / "_chapters_speed.txt"
    try:
        chapters = []
        start_ms = 0
        with open(concat_list, "w", encoding="utf-8") as list_file:
            for chapter_file, title in zip(chapter_files, titles):
                list_file.write(f"file '{_concat_escape(str(chapter_file))}'\n")
                duration_ms = probe_duration_ms(chapter_file)
                end_ms = start_ms + duration_ms
                chapters.append((start_ms, end_ms, title or "Chapter"))
                start_ms = end_ms

        build_ffmetadata(chapters, metadata_file)

        speed_path = None
        if atempo_filters(speed):
            speed_path = output_path.with_name(f"{output_path.stem}_{speed:g}{output_path.suffix}")
            scaled = [(int(s / speed), int(e / speed), t) for s, e, t in chapters]
            build_ffmetadata(scaled, speed_metadata_file)
        cmd = build_m4b_chapters_command(concat_list, metadata_file, output_path,
                                         speed=speed, speed_path=speed_path,
                                         speed_metadata_file=speed_metadata_file
                                         if speed_path is not None else None)

        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode != 0:
            logger.error("ffmpeg failed: %s", proc.stderr[-2000:])
            return False

        # The last chapter's end time is the expected total duration; skip
        # verification when any chapter duration probe failed (returned 0)
        # so an unprobed chapter can't falsely fail the whole output.
        expected_ms = chapters[-1][1]
        if any(end_ms - start_ms <= 0 for start_ms, end_ms, _ in chapters):
            logger.warning("Could not probe every chapter duration; skipping duration verification")
            duration_ok = True
        else:
            duration_ok = verify_output_duration(output_path, expected_ms)
            if speed_path is not None:
                duration_ok = verify_output_duration(speed_path, int(expected_ms / speed)) and duration_ok
        if not duration_ok:
            return False

        logger.info("Audiobook saved: %s (%d chapters)", output_path, len(chapter_files))
        print(f"[INFO] Saved audiobook: {output_path.name} ({len(chapter_files)} chapters)")

        if speed_path is not None:
            logger.info("Saved speed-adjusted audiobook (%gx): %s", speed, speed_path)
            print(f"[INFO] Saved speed-adjusted audiobook: {speed_path.name} ({speed:g}x)")

        return True

    except FileNotFoundError:
        logger.error("ffmpeg/ffprobe not found on PATH (install ffmpeg and try again)")
        return False
    except Exception as exc:
        logger.error("Failed to combine chapters: %s", exc)
        logger.error(traceback.format_exc())
        return False
    finally:
        for scratch in (concat_list, metadata_file, speed_metadata_file):
            try:
                scratch.unlink(missing_ok=True)
            except OSError:
                pass