UmbraAudio/tests/test_ar_export.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import json
import re
import shutil
import sqlite3
import subprocess
from typing import TYPE_CHECKING

import numpy as np
import pytest
import soundfile as sf

if TYPE_CHECKING:
    from pathlib import Path

from umbraaudio.assembly.ar_export import (
    ARChapter,
    ARPackageExporter,
    discover_chapters,
    export_chapter_mp3,
)

DUR = 24000
META = {
    "title": "Deal Hunter",
    "author": "B.D. Murphy",
    "narrator": "Umbra Audio",
    "publisher": "WorstAuthor LLC",
    "output_name": "Deal_Hunter.m4b",
}


def _fake_chapter_wav(path: Path, seconds: float = 1.0, amp: float = 0.1) -> Path:
    sf.write(str(path), (amp * np.random.randn(int(DUR * seconds))).astype(np.float32), DUR)
    return path


def _sample_book(data_dir: Path) -> list[ARChapter]:
    audio = data_dir / "chapter_audio"
    audio.mkdir(parents=True, exist_ok=True)
    return [
        ARChapter(1, 0, "Sample Track", "Welcome to the audio book Deal Hunter.",
                  _fake_chapter_wav(audio / "ch1_full.wav"), role="sample"),
        ARChapter(4, 3, "Copyright", "Copyright c 2023 by B.D. Murphy",
                  _fake_chapter_wav(audio / "ch4_full.wav"), role="chapter"),
        ARChapter(5, 4, "Author's Note", "Author's Note",
                  _fake_chapter_wav(audio / "ch5_full.wav"), role="chapter"),
        ARChapter(6, 5, "Salvage", "Deal Hunter. Chapter one, Salvage.",
                  _fake_chapter_wav(audio / "ch6_full.wav"), role="chapter"),
        ARChapter(7, 6, "Family Reunion", "The princess is back home in the palace on Albert-4.",
                  _fake_chapter_wav(audio / "ch7_full.wav"), role="chapter"),
        ARChapter(35, 34, "Closing Track", "Thank You for listening to Deal Hunter.",
                  _fake_chapter_wav(audio / "ch35_full.wav"), role="closing"),
    ]


def _build_book(data_dir: Path, book: list[ARChapter] | None = None, **kwargs) -> dict:
    exporter = ARPackageExporter(data_dir)
    return exporter.build(book or _sample_book(data_dir), dict(META), **kwargs)


def test_full_build_manifest(temp_dir: Path) -> None:
    manifest = _build_book(temp_dir / ".umbra")
    assert manifest["track_count"] == 6
    assert manifest["format"] == {
        "codec": "mp3", "sample_rate": 44100, "channels": 1,
        "bit_rate_kbps": 192, "cbr": True, "silence_s": 2.0,
    }
    roles = [t["role"] for t in manifest["tracks"]]
    # The closing credit is LAST, after the story it concludes.
    assert roles == ["sample", "chapter", "chapter", "chapter", "chapter", "closing"]
    assert [t["number"] for t in manifest["tracks"]] == ["", "", "", "1", "2", "ii"]
    assert [t["name"] for t in manifest["tracks"]] == [
        "Sample Track", "Copyright", "Author's Note",
        "Salvage", "Family Reunion", "Closing Track",
    ]
    assert [t["title"] for t in manifest["tracks"]] == [
        "Sample Track", "Copyright", "Author's Note",
        "1 Salvage", "2 Family Reunion", "ii Closing Track",
    ]

    files = sorted((temp_dir / ".umbra" / "export" / "ar_package").glob("*.mp3"))
    names = [f.name for f in files]
    assert len(names) == 6
    assert "Deal_Hunter_Sample_Track.mp3" in names
    assert "Deal_Hunter_Copyright.mp3" in names
    assert "Deal_Hunter_Authors_Note.mp3" in names
    assert "Deal_Hunter_01_Salvage.mp3" in names
    assert "Deal_Hunter_02_Family_Reunion.mp3" in names
    assert "Deal_Hunter_ii_Closing_Track.mp3" in names


def test_tracks_are_straight_chapter_audio_no_synthesis(temp_dir: Path) -> None:
    manifest = _build_book(temp_dir / ".umbra")
    for t in manifest["tracks"]:
        assert t["source"].endswith("_full.wav")
        assert t["source"] != ""
    assert all("announcement" not in t for t in manifest["tracks"])


def test_chapters_without_audio_are_skipped(temp_dir: Path) -> None:
    book = _sample_book(temp_dir / ".umbra")
    book[3] = ARChapter(6, 5, "Salvage", "Deal Hunter. Chapter one, Salvage.",
                        None, role="chapter")
    manifest = _build_book(temp_dir / ".umbra", book)
    assert manifest["track_count"] == 5
    assert all(t["title"] != "Salvage" for t in manifest["tracks"])


def test_loudness_warning_for_quiet_chapter(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    book = _sample_book(data_dir)
    quiet = book[3].audio_path
    assert quiet is not None
    _fake_chapter_wav(quiet, amp=0.0001)
    book[3] = ARChapter(6, 5, "Salvage", "Deal Hunter. Chapter one, Salvage.",
                        quiet, role="chapter")
    manifest = _build_book(data_dir, book)
    assert any("RMS" in w for w in manifest["warnings"])


def test_loudness_normalization_brings_quiet_book_to_spec(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    audio = data_dir / "chapter_audio"
    audio.mkdir(parents=True, exist_ok=True)
    book = [
        ARChapter(1, 0, "Sample Track", "Welcome to the audio book Deal Hunter.",
                  _fake_chapter_wav(audio / "ch1_full.wav", amp=0.02),
                  role="sample"),
        ARChapter(4, 3, "Copyright", "Copyright c 2023 by B.D. Murphy",
                  _fake_chapter_wav(audio / "ch4_full.wav", amp=0.02),
                  role="chapter"),
        ARChapter(35, 34, "Closing Track", "Thank You for listening to Deal Hunter.",
                  _fake_chapter_wav(audio / "ch35_full.wav", amp=0.02),
                  role="closing"),
    ]
    manifest = _build_book(data_dir, book)
    assert manifest["warnings"] == [], manifest["warnings"]
    for t in manifest["tracks"]:
        assert -22.5 <= t["rms_db"] <= -18.5, t
        assert t["peak_db"] <= -3.0, t
        assert t["loudness_ok"] and t["peak_ok"], t


def test_peak_limiter_caps_clipping_source(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    audio = data_dir / "chapter_audio"
    audio.mkdir(parents=True, exist_ok=True)
    hot = audio / "ch1_full.wav"
    sf.write(str(hot), (4.0 * np.random.randn(DUR)).astype(np.float32), DUR)
    book = _sample_book(data_dir)
    book[0] = ARChapter(1, 0, "Sample Track", "Welcome to the audio book Deal Hunter.",
                        hot, role="sample")
    manifest = _build_book(data_dir, book)
    sample = next(t for t in manifest["tracks"] if t["role"] == "sample")
    assert sample["peak_db"] <= -3.0
    assert sample["peak_ok"]


def test_rebuild_is_idempotent(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    exporter = ARPackageExporter(data_dir)
    first = exporter.build(_sample_book(data_dir), dict(META))
    second = exporter.build(_sample_book(data_dir), dict(META))
    assert isinstance(second, dict)
    assert second["track_count"] == first["track_count"]
    assert [t["title"] for t in second["tracks"]] == [t["title"] for t in first["tracks"]]


def test_manifest_written_and_total_duration(temp_dir: Path) -> None:
    manifest_path = temp_dir / ".umbra" / "export" / "ar_package.json"
    manifest = _build_book(temp_dir / ".umbra")
    assert manifest_path.exists()
    on_disk = json.loads(manifest_path.read_text())
    assert on_disk["track_count"] == manifest["track_count"]
    expected_duration = sum(t["duration_s"] for t in manifest["tracks"])
    assert abs(manifest["total_duration_s"] - expected_duration) < 0.01


@pytest.mark.skipif(shutil.which("ffmpeg") is None or shutil.which("ffprobe") is None,
                    reason="ffmpeg/ffprobe required")
def test_encoded_files_are_44100_mono_192k_cbr(temp_dir: Path) -> None:
    _build_book(temp_dir / ".umbra")
    mp3 = next((temp_dir / ".umbra" / "export" / "ar_package").glob("*.mp3"))
    probe = subprocess.run(
        [
            "ffprobe", "-v", "error",
            "-show_entries", "stream=codec_name,sample_rate,channels,bit_rate",
            "-show_entries", "format=duration",
            "-of", "json",
            str(mp3),
        ],
        capture_output=True, text=True,
    )
    assert probe.returncode == 0, probe.stderr
    data = json.loads(probe.stdout)
    stream = data["streams"][0]
    assert stream["codec_name"] == "mp3"
    assert stream["sample_rate"] == "44100"
    assert int(stream["channels"]) == 1
    assert stream["bit_rate"] == "192000"


@pytest.mark.skipif(shutil.which("ffprobe") is None, reason="ffprobe required")
def test_chapter_mp3_export_is_ar_compliant(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    audio = data_dir / "chapter_audio"
    audio.mkdir(parents=True, exist_ok=True)
    wav = _fake_chapter_wav(audio / "ch7_full.wav", seconds=1.0, amp=0.1)
    out = audio / "ch7_full.mp3"

    stats = export_chapter_mp3(data_dir, wav, out)
    assert stats["duration_s"] == 5.0
    assert stats["size_bytes"] == out.stat().st_size
    assert stats["loudness_ok"] is True
    assert stats["peak_ok"] is True

    probe = subprocess.run(
        [
            "ffprobe", "-v", "error",
            "-show_entries", "stream=codec_name,sample_rate,channels,bit_rate",
            "-show_entries", "format=duration",
            "-of", "json",
            str(out),
        ],
        capture_output=True, text=True,
    )
    assert probe.returncode == 0, probe.stderr
    data = json.loads(probe.stdout)
    stream = data["streams"][0]
    assert stream["codec_name"] == "mp3"
    assert stream["sample_rate"] == "44100"
    assert int(stream["channels"]) == 1
    assert stream["bit_rate"] == "192000"
    assert abs(float(data["format"]["duration"]) - 5.0) < 0.3


@pytest.mark.skipif(shutil.which("ffprobe") is None, reason="ffprobe required")
def test_track_duration_matches_file_length(temp_dir: Path) -> None:
    manifest = _build_book(temp_dir / ".umbra")
    pkg = temp_dir / ".umbra" / "export" / "ar_package"
    for t in manifest["tracks"]:
        probe = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "json", str(pkg / t["file"])],
            capture_output=True, text=True,
        )
        assert probe.returncode == 0
        actual = float(json.loads(probe.stdout)["format"]["duration"])
        assert abs(actual - t["duration_s"]) < 0.3


def _make_project_db(data_dir: Path) -> None:
    data_dir.mkdir(parents=True, exist_ok=True)
    db = sqlite3.connect(data_dir / "project.db")
    db.executescript(
        """
        CREATE TABLE chapters (id INTEGER PRIMARY KEY, chapter_index INTEGER,
                               title TEXT, source_text TEXT);
        CREATE TABLE utterances (
            id INTEGER PRIMARY KEY,
            chapter_id INTEGER,
            utterance_index INTEGER,
            type TEXT,
            speaker_name TEXT,
            text TEXT);
        """
    )
    db.executemany(
        "INSERT INTO chapters (id, chapter_index, title, source_text) VALUES (?,?,?,?)",
        [
            (1, 0, "Sample Track", ""),
            (2, 1, "Title Page", ""),
            (3, 2, "Opening Track", ""),
            (4, 3, "Copyright", ""),
            (6, 5, "Salvage", ""),
            (35, 34, "Closing Track", ""),
        ],
    )
    db.executemany(
        "INSERT INTO utterances "
        "(chapter_id, utterance_index, type, speaker_name, text) VALUES (?,?,?,?,?)",
        [
            (1, 0, "narrator", "Narrator", "Welcome to the audio book Deal Hunter."),
            (4, 0, "narrator", "Narrator", "Copyright c 2023"),
            (6, 0, "narrator", "Narrator", "Deal Hunter. Chapter one, Salvage."),
            (35, 0, "narrator", "Narrator", "Thank You for listening to Deal Hunter."),
        ],
    )
    db.commit()
    db.close()
    audio = data_dir / "chapter_audio"
    audio.mkdir()
    _fake_chapter_wav(audio / "ch1_full.wav")
    _fake_chapter_wav(audio / "ch4_full.wav")
    _fake_chapter_wav(audio / "ch6_full.wav")
    _fake_chapter_wav(audio / "ch35_full.wav")


def test_discover_chapters_assigns_roles(temp_dir: Path) -> None:
    data_dir = temp_dir / ".umbra"
    _make_project_db(data_dir)
    chapters = discover_chapters(data_dir)
    roles = {c.role for c in chapters}
    assert roles == {"sample", "chapter", "closing"}
    sample = next(c for c in chapters if c.role == "sample")
    closing = next(c for c in chapters if c.role == "closing")
    assert sample.chapter_id == 1
    assert sample.title == "Sample Track"
    assert closing.chapter_id == 35
    assert closing.title == "Closing Track"
    assert all(c.title not in ("Title Page", "Opening Track") for c in chapters)
    assert [c.title for c in chapters] == [
        "Sample Track", "Copyright", "Salvage", "Closing Track",
    ]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])


# ---------------------------------------------------------------------------
# Peak gate: measure the delivered file, and enforce the ceiling on it.
#
# Regression cover for a real defect: _encode_track measured the pre-encode
# float buffer, and _normalize hard-clipped to exactly -3.0 dB, so all 32
# shipped tracks recorded peak_db = -3.0 and peak_ok = True. libmp3lame then
# overshot the ceiling, so the files listeners received actually peaked near
# -2.0 dB -- above the -3 dB limit that Author's Republic, INaudio and Spotify
# all require. The gate certified intent, not the artifact.
# ---------------------------------------------------------------------------


def _ffmpeg_peak_db(path: Path) -> float:
    """Measure a delivered file's peak independently of ar_export."""
    proc = subprocess.run(
        ["ffmpeg", "-hide_banner", "-nostats", "-i", str(path),
         "-af", "volumedetect", "-f", "null", "-"],
        capture_output=True, text=True, check=True,
    )
    match = re.search(r"max_volume:\s*([-\d.]+) dB", proc.stderr)
    assert match, proc.stderr[-500:]
    return float(match.group(1))


@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg required")
def test_manifest_peak_describes_the_encoded_file_not_the_source(temp_dir: Path) -> None:
    """The recorded peak must equal the delivered file's peak."""
    manifest = _build_book(temp_dir / ".umbra")
    package = temp_dir / ".umbra" / "export" / "ar_package"
    for track in manifest["tracks"]:
        recorded = track["peak_db"]
        actual = _ffmpeg_peak_db(package / track["file"])
        assert recorded == pytest.approx(actual, abs=0.1), (
            f"{track['title']}: manifest says {recorded} but the file measures {actual}"
        )


@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg required")
def test_every_delivered_file_meets_the_peak_ceiling(temp_dir: Path) -> None:
    """No shipped MP3 may peak above -3.0 dB, whatever the encoder did."""
    manifest = _build_book(temp_dir / ".umbra")
    package = temp_dir / ".umbra" / "export" / "ar_package"
    for track in manifest["tracks"]:
        actual = _ffmpeg_peak_db(package / track["file"])
        assert actual <= -3.0, f"{track['title']} peaks at {actual} dB, over the -3.0 limit"
        assert track["peak_ok"] is True, track


@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg required")
def test_hot_source_is_trimmed_until_the_file_complies(tmp_path: Path) -> None:
    """A deliberately clipping source must still deliver under the ceiling."""
    data_dir = tmp_path / ".umbra"
    audio = data_dir / "chapter_audio"
    audio.mkdir(parents=True, exist_ok=True)
    hot = audio / "ch1_full.wav"
    sf.write(str(hot), (4.0 * np.random.randn(DUR)).astype(np.float32), DUR)
    exporter = ARPackageExporter(data_dir)
    out = data_dir / "export" / "ar_package" / "hot.mp3"
    out.parent.mkdir(parents=True, exist_ok=True)
    stats = exporter.encode_chapter_mp3(hot, out)

    assert _ffmpeg_peak_db(out) <= -3.0
    assert stats["peak_ok"] is True
    # RMS is deliberately NOT asserted here: NORM_MAX_GAIN_DB caps gain at
    # +/-12 dB, so a source this hot is not brought to the RMS target.


def test_closing_credit_is_the_final_track(temp_dir: Path) -> None:
    """The end-of-book credit belongs last, not at position 3."""
    manifest = _build_book(temp_dir / ".umbra")
    titles = [t["title"] for t in manifest["tracks"]]
    roles = [t["role"] for t in manifest["tracks"]]
    assert roles[-1] == "closing", roles
    assert titles[-1].endswith("Closing Track"), titles
    assert titles[0] == "Sample Track", titles
    assert roles.count("closing") == 1


def test_track_sequence_numbers_stay_contiguous(temp_dir: Path) -> None:
    manifest = _build_book(temp_dir / ".umbra")
    seqs = [t["seq"] for t in manifest["tracks"]]
    assert seqs == list(range(1, len(seqs) + 1)), seqs