path: root/converter/audio.py

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

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

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]
    return ["-b:a", config.AUDIO_BITRATE]


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

    ``results`` maps chunk numbers to success flags; failed 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.
    """
    if shutil.which("ffmpeg") is None:
        logger.error("ffmpeg is required to combine audio chunks (install ffmpeg)")
        return False

    chunk_files = []
    missing_chunks = []
    for i in range(1, total_chunks + 1):
        # Skip chunks that failed if we have results tracking
        if results is not None and not results.get(i, False):
            missing_chunks.append(i)
            continue

        matches = sorted(config.CHUNKS_FOLDER.glob(f"chunk_{i:04d}.*"))
        if matches:
            chunk_files.append(matches[0])
        else:
            missing_chunks.append(i)

    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")

        filters = atempo_filters(speed)
        encode = _encode_args(output_format)
        if filters:
            speed_path = output_path.with_name(f"{output_path.stem}_{speed:g}{output_path.suffix}")
            cmd = [
                "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, str(output_path),
                "-map", "[spdout]", *encode, str(speed_path),
            ]
        else:
            speed_path = None
            cmd = [
                "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
                *encode, str(output_path),
            ]

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

        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 filters:
            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."""
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", str(path)],
        capture_output=True, text=True,
    )
    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 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={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)

        filters = atempo_filters(speed)
        encode = _encode_args("m4b")
        if filters:
            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)
            # The speed-adjusted stream needs rescaled chapter markers, so the
            # rescaled metadata is passed as a third input.
            cmd = [
                "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]{filters}[spdout]",
                "-map", "[base]", "-map_metadata", "1", *encode, str(output_path),
                "-map", "[spdout]", "-map_metadata", "2", *encode, str(speed_path),
            ]
        else:
            speed_path = None
            cmd = [
                "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
                "-i", str(metadata_file),
                "-map", "0:a", "-map_metadata", "1", *encode, str(output_path),
            ]

        proc = subprocess.run(cmd, capture_output=True, text=True)
        if proc.returncode != 0:
            logger.error("ffmpeg failed: %s", proc.stderr[-2000:])
            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