backend/export.py

import json
import shutil
import subprocess
import wave
import zipfile
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

from mutagen.easyid3 import EasyID3
from mutagen.id3 import ID3, APIC, TIT2, TPE1


def export_audiobook(
    *,
    session_root: Path,
    chapters: List[Dict[str, object]],
    output_format: str,
    metadata: Dict[str, str],
    embed_markers: bool,
    cover_path: Path | None = None,
) -> Dict[str, object]:
    exports_dir = session_root / "exports"
    exports_dir.mkdir(parents=True, exist_ok=True)
    rendered = _rendered_files(session_root)
    if not rendered:
        raise ValueError("No rendered chapters available for export")

    if output_format == "zip":
        archive = exports_dir / "reading-room-chapters.zip"
        with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for source in rendered:
                zf.write(source, arcname=source.name)
        return {"file": str(archive), "format": "zip", "size_bytes": archive.stat().st_size}

    base_wav = exports_dir / "reading-room.wav"
    total_seconds = _concat_wav(rendered, base_wav)
    chapter_offsets = _chapter_offsets(chapters)

    if output_format == "m4a":
        output = exports_dir / "reading-room.m4a"
        _ffmpeg_transcode(
            source=base_wav,
            target=output,
            codec_args=["-c:a", "aac", "-b:a", "96k"],
            metadata=metadata,
            chapter_offsets=chapter_offsets if embed_markers else None,
        )
    elif output_format == "mp3":
        output = exports_dir / "reading-room.mp3"
        _ffmpeg_transcode(
            source=base_wav,
            target=output,
            codec_args=["-c:a", "libmp3lame", "-b:a", "128k"],
            metadata=metadata,
            chapter_offsets=None,
        )
        _tag_mp3(output, metadata, cover_path)
    else:
        raise ValueError("Unsupported export format")

    return {
        "file": str(output),
        "format": output_format,
        "runtime_seconds": total_seconds,
        "size_bytes": output.stat().st_size,
        "chapter_count": len(chapters),
    }


def _rendered_files(session_root: Path) -> List[Path]:
    render_dir = session_root / "renders"
    return sorted(render_dir.glob("*.wav"))


def _concat_wav(inputs: Iterable[Path], output: Path) -> int:
    inputs = list(inputs)
    params = None
    total_frames = 0
    frames: List[bytes] = []
    for wav_path in inputs:
        with wave.open(str(wav_path), "rb") as handle:
            if params is None:
                params = handle.getparams()
            elif handle.getframerate() != params.framerate:
                raise ValueError("Mismatched sample rates across rendered chapters")
            chunk = handle.readframes(handle.getnframes())
            frames.append(chunk)
            total_frames += handle.getnframes()

    if params is None:
        raise ValueError("No WAV inputs to concatenate")

    with wave.open(str(output), "wb") as handle:
        handle.setparams(params)
        for chunk in frames:
            handle.writeframes(chunk)

    return int(total_frames / params.framerate)


def _ffmpeg_transcode(
    *,
    source: Path,
    target: Path,
    codec_args: List[str],
    metadata: Dict[str, str],
    chapter_offsets: List[Tuple[str, int, int]] | None,
) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    ffmetadata_path = None
    cmd = ["ffmpeg", "-y", "-i", str(source)]
    if chapter_offsets:
        ffmetadata_path = target.with_suffix(".ffmeta")
        ffmetadata_path.write_text(_ffmetadata(metadata, chapter_offsets), encoding="utf-8")
        cmd.extend(["-i", str(ffmetadata_path), "-map_metadata", "1"])
    else:
        for key, value in metadata.items():
            cmd.extend(["-metadata", f"{key}={value}"])
    cmd.extend(codec_args)
    cmd.append(str(target))
    subprocess.run(cmd, check=True, capture_output=True)


def _ffmetadata(metadata: Dict[str, str], chapter_offsets: List[Tuple[str, int, int]]) -> str:
    lines = [";FFMETADATA1"]
    for key, value in metadata.items():
        lines.append(f"{key}={value}")
    for title, start, end in chapter_offsets:
        lines.extend(
            [
                "[CHAPTER]",
                "TIMEBASE=1/1000",
                f"START={start}",
                f"END={end}",
                f"title={title}",
            ]
        )
    return "\n".join(lines)


def _chapter_offsets(chapters: List[Dict[str, object]]) -> List[Tuple[str, int, int]]:
    offsets = []
    start_ms = 0
    for chapter in chapters:
        duration_ms = int(chapter.get("duration_seconds", 0) * 1000)
        offsets.append((str(chapter["title"]), start_ms, start_ms + duration_ms))
        start_ms += duration_ms
    return offsets


def _tag_mp3(path: Path, metadata: Dict[str, str], cover_path: Path | None) -> None:
    try:
        tags = EasyID3(str(path))
    except Exception:
        tags = EasyID3()
        tags.save(str(path))
        tags = EasyID3(str(path))
    if metadata.get("title"):
        tags["title"] = [metadata["title"]]
    if metadata.get("artist"):
        tags["artist"] = [metadata["artist"]]
    tags.save()

    if cover_path and cover_path.exists():
        audio = ID3(str(path))
        audio.add(
            APIC(
                encoding=3,
                mime="image/jpeg",
                type=3,
                desc="Cover",
                data=cover_path.read_bytes(),
            )
        )
        if metadata.get("title"):
            audio.add(TIT2(encoding=3, text=metadata["title"]))
        if metadata.get("artist"):
            audio.add(TPE1(encoding=3, text=metadata["artist"]))
        audio.save(str(path))