"""Map chapter text positions onto audio time and tag the MP3 (ID3v2.4 CHAP).

Mapping: probe each chunk MP3's duration once (ffprobe) to get chunk start
times in ms; use chunk_offsets.json (start char-offset of each chunk in the
cleaned text) to binary-search each chapter's position. Each chapter spans
from its start to the next chapter's start (last: end of file).

Uses mutagen's CHAP class (1.x kwargs: start_time/end_time in ms,
element_id, sub_frames). Hand-rolling the CHAP bytes is error-prone.

Pitfalls baked in (all hit in the reference 2026-08 session):
  - work on a COPY of the audio (ID3() + save() to a missing path writes a
    tag-only file that ffprobe cannot read)
  - TXXX uses text= (value= silently stores empty)
  - many players display element_id, not the TXXX subframe, so the title
    goes into element_id too (Latin1, CHnnn fallback for non-Latin1)
"""
import json
import os
import shutil
import subprocess
import sys
from bisect import bisect_right


def probe_chunk_times(audio_dir: str) -> tuple:
    """Return (offsets_ms, total_ms): start time of each chunk in ms."""
    files = sorted(f for f in os.listdir(audio_dir) if f.endswith(".mp3"))
    offsets = []
    t = 0
    for fn in files:
        offsets.append(t)
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=nw=1:nk=1", os.path.join(audio_dir, fn)],
            capture_output=True, text=True, check=True,
        )
        t += int(float(r.stdout.strip()) * 1000)
    return offsets, t


def main(clean_path: str, audio_dir: str, in_mp3: str, out_mp3: str):
    # clean.py writes <clean>.chapters.json / <clean>.chunk_offsets.json
    stem = os.path.splitext(clean_path)[0]
    chapters = json.load(open(stem + ".chapters.json"))
    offsets_path = stem + ".chunk_offsets.json"
    if os.path.exists(offsets_path):
        char_offsets = json.load(open(offsets_path))
    else:
        # fallback: even distribution (crude, but keeps the tool working)
        total_chars = len(open(clean_path, encoding="utf-8").read())
        n = len([f for f in os.listdir(audio_dir) if f.endswith(".mp3")])
        char_offsets = [int(i * total_chars / n) for i in range(n)]
        print("WARNING: chunk_offsets.json missing; using even char distribution")

    offsets_ms, total_ms = probe_chunk_times(audio_dir)
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=nw=1:nk=1", in_mp3],
        capture_output=True, text=True, check=True,
    )
    file_end_ms = int(float(r.stdout.strip()) * 1000)

    entries = []
    for pos, title in chapters:
        i = bisect_right(char_offsets, pos) - 1
        i = max(0, min(i, len(offsets_ms) - 1))
        entries.append([offsets_ms[i], title])
    entries.sort(key=lambda e: e[0])
    # drop two chapters that land in the same chunk (zero-length span)
    dedup = []
    for st, title in entries:
        if dedup and dedup[-1][0] == st:
            continue
        dedup.append([st, title])
    entries = dedup
    if not entries:
        raise SystemExit("no usable chapters; check heading_pattern in the config")

    from mutagen.id3 import CHAP, ID3, ID3Tags, TXXX

    shutil.copy(in_mp3, out_mp3)
    tags = ID3(out_mp3)
    for n, (st, title) in enumerate(entries, 1):
        en = entries[n][0] if n < len(entries) else file_end_ms
        elem_id = title if title.encode("latin-1", "replace") == title.encode("latin-1") \
            else f"CH{n:03d}"
        sub = ID3Tags()
        sub.add(TXXX(desc="CHAPTER", text=title))
        tags.add(CHAP(element_id=elem_id, start_time=int(st),
                      end_time=int(en), sub_frames=sub))
    tags.save(out_mp3, v2_version=4)

    print(json.dumps({
        "out": out_mp3,
        "chapters": len(entries),
        "total_ms": total_ms,
        "file_end_ms": file_end_ms,
        "first": entries[0],
        "last": entries[-1],
    }, indent=1))
    print("verify: ffprobe -v error -show_chapters " + out_mp3)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])