#!/usr/bin/env python3
"""Resumable overnight export of the three DSA audiobooks, one MP3 per chapter."""
from __future__ import annotations

import io
import hashlib
import json
import os
import re
import subprocess
import sys
import time
import wave
from pathlib import Path
from urllib.parse import quote

import requests

BASE = "http://127.0.0.1:7890"
ROOT = Path(__file__).resolve().parents[1]
CACHE = ROOT / "config" / "audiobook_batch_cache"
PAUSE_SENTENCE_MS = 120
PAUSE_PARAGRAPH_MS = 300
CHUNK_LIMIT = 900
BOOK_ORDER = [
    "Hohlbein & Hennen - DSA - Das Jahr des Greifen 01 - Der Sturm",
    "Hohlbein & Hennen - DSA - Das Jahr des Greifen 02 - Die Entdeckung",
    "Hohlbein & Hennen - DSA - Das Jahr des Greifen 03 - Die Amazone",
]

CHAPTER = re.compile(
    r"(?im)(?:^|\n\s*\n|\n\s*)(?P<header>(?:(?:\d{1,3}|[ivxlcdm]+)[.)]?\s*)?"
    r"(?:kapitel|chapter|prolog|epilog)\b[^\n]{0,80})"
)


def key(value: str) -> str:
    return str(value or "").strip().casefold()


def safe(value: str, limit: int = 80) -> str:
    value = re.sub(r"[^A-Za-z0-9ÄÖÜäöüß _.-]+", "", value).strip()
    return re.sub(r"\s+", "_", value)[:limit] or "chapter"


def split_text(text: str) -> list[tuple[str, int]]:
    """Return (text, pause-after-ms), preferring paragraph/sentence edges.

    The old exporter flattened whitespace and cut at commas/spaces, which made
    the TTS engine restart in the middle of a sentence. A paragraph is kept as
    one request whenever possible. Only an oversized paragraph is divided,
    first at sentence endings and only as a last resort at a clause/word edge.
    """
    raw = str(text or "").replace("\r\n", "\n").strip()
    paragraphs = [re.sub(r"\s+", " ", p).strip() for p in re.split(r"\n\s*\n+", raw) if p.strip()]
    result: list[tuple[str, int]] = []
    sentence_re = re.compile(r".+?(?:[.!?…]+[\"»”’']*|$)", re.S)
    for paragraph in paragraphs:
        if len(paragraph) <= CHUNK_LIMIT:
            result.append((paragraph, PAUSE_PARAGRAPH_MS))
            continue
        sentences = [s.strip() for s in sentence_re.findall(paragraph) if s.strip()]
        pieces: list[str] = []
        current = ""
        for sentence in sentences or [paragraph]:
            if len(sentence) > CHUNK_LIMIT:
                clauses = re.split(r"(?<=[,;:])\s+", sentence)
            else:
                clauses = [sentence]
            for clause in clauses:
                clause = clause.strip()
                if not clause:
                    continue
                joined = f"{current} {clause}".strip() if current else clause
                if len(joined) <= CHUNK_LIMIT:
                    current = joined
                else:
                    if current:
                        pieces.append(current)
                    current = clause
                    while len(current) > CHUNK_LIMIT:
                        cut = current.rfind(" ", 80, CHUNK_LIMIT + 1)
                        if cut <= 80:
                            cut = CHUNK_LIMIT
                        pieces.append(current[:cut].strip())
                        current = current[cut:].strip()
        if current:
            pieces.append(current)
        for index, piece in enumerate(pieces):
            result.append((piece, PAUSE_PARAGRAPH_MS if index == len(pieces) - 1 else PAUSE_SENTENCE_MS))
    return result


def wav_pcm(body: bytes) -> tuple[bytes, wave._wave_params]:
    with wave.open(io.BytesIO(body), "rb") as stream:
        params = stream.getparams()
        if params.sampwidth != 2:
            raise RuntimeError(f"TTS returned non-16-bit WAV: {params}")
        return stream.readframes(stream.getnframes()), params


def make_wav(raw: bytes, params: wave._wave_params) -> bytes:
    output = io.BytesIO()
    with wave.open(output, "wb") as stream:
        stream.setnchannels(params.nchannels)
        stream.setsampwidth(params.sampwidth)
        stream.setframerate(params.framerate)
        stream.writeframes(raw)
    return output.getvalue()


def request_json(method: str, url: str, **kwargs):
    kwargs.setdefault("timeout", 90)
    for attempt in range(1, 4):
        try:
            response = requests.request(method, url, **kwargs)
            if response.ok:
                return response
            if response.status_code < 500:
                raise RuntimeError(f"{method} {url}: {response.status_code} {response.text[:300]}")
            raise RuntimeError(f"server {response.status_code}: {response.text[:300]}")
        except Exception:
            if attempt == 3:
                raise
            time.sleep(attempt * 3)


def tts(voice: str, text: str) -> tuple[bytes, wave._wave_params]:
    seed = int.from_bytes(hashlib.sha256(voice.encode("utf-8")).digest()[:8], "big") % 2147483647
    response = request_json(
        "POST", f"{BASE}/api/tts-preview",
        json={"voice": voice, "text": text, "response_format": "wav", "backend": "voice_clone", "seed": seed},
        timeout=180,
    )
    return wav_pcm(response.content)


def chapter_parts(segments: list[dict]) -> list[dict]:
    """Split segments at headings even when OCR put a heading inside a segment."""
    buckets: list[dict] = [{"title": "", "segments": []}]
    for segment in segments:
        text = str(segment.get("text") or "").strip()
        if not text:
            continue
        matches = list(CHAPTER.finditer(text))
        if not matches:
            buckets[-1]["segments"].append({**segment, "text": text})
            continue
        prefix = text[:matches[0].start()].strip()
        if prefix:
            buckets[-1]["segments"].append({**segment, "text": prefix})
        for index, match in enumerate(matches):
            if buckets[-1]["segments"] or buckets[-1]["title"]:
                buckets.append({"title": match.group("header").strip(), "segments": []})
            else:
                buckets[-1]["title"] = match.group("header").strip()
            end = matches[index + 1].start() if index + 1 < len(matches) else len(text)
            body = text[match.end():end].strip()
            if body:
                buckets[-1]["segments"].append({**segment, "text": body, "type": "narration", "speaker": "Narrator"})
    return [bucket for bucket in buckets if bucket["segments"]]


def voice_map(book: str) -> tuple[dict[str, str], str]:
    characters = request_json("GET", f"{BASE}/api/characters").json().get("characters", [])
    mapping: dict[str, str] = {}
    for character in characters:
        if character.get("book") != book or not character.get("voice"):
            continue
        voice = str(character["voice"])
        mapping[key(character.get("name"))] = voice
        for alias in re.split(r"[,;|]", str((character.get("sheet") or {}).get("aliases") or "")):
            if alias.strip():
                mapping[key(alias)] = voice
    narrator = mapping.get("narrator")
    if not narrator:
        raise RuntimeError(f"{book}: no Narrator voice assignment found")
    return mapping, narrator


def render_book(doc: dict) -> None:
    book = doc["title"]
    doc_id = doc["id"]
    cast = request_json("GET", f"{BASE}/api/reader/docs/{doc_id}/scripts/cast").json()
    mapping, narrator = voice_map(book)
    buckets = chapter_parts(cast.get("segments", []))
    book_dir = CACHE / safe(book)
    book_dir.mkdir(parents=True, exist_ok=True)
    print(f"[{book}] {len(buckets)} chapter buckets", flush=True)
    for chapter_index, bucket in enumerate(buckets, 1):
        title = bucket["title"] or f"Chapter {chapter_index}"
        chapter_dir = book_dir / f"{chapter_index:03d}_{safe(title, 48)}"
        # New cache namespace: chunk boundaries changed from unsafe fixed
        # character cuts to paragraph/sentence boundaries.
        chunks_dir = chapter_dir / "chunks_v2_sentence_aware"
        chunks_dir.mkdir(parents=True, exist_ok=True)
        chunks: list[tuple[str, str, int]] = []
        fallback = 0
        for segment in bucket["segments"]:
            speaker = key(segment.get("speaker") or "narrator")
            voice = mapping.get(speaker, narrator)
            if speaker not in mapping and speaker != "narrator":
                fallback += 1
            for text, pause_ms in split_text(segment.get("text")):
                chunks.append((voice, text, pause_ms))
        pcm_parts: list[bytes] = []
        params = None
        for index, (voice, text, pause_ms) in enumerate(chunks):
            chunk_path = chunks_dir / f"{index:05d}.pcm"
            meta_path = chunks_dir / f"{index:05d}.json"
            if chunk_path.is_file() and meta_path.is_file():
                pcm_parts.append(chunk_path.read_bytes())
                if params is None:
                    params = wave._wave_params(**json.loads(meta_path.read_text()))
                continue
            raw, current = tts(voice, text)
            if params is None:
                params = current
            elif (current.nchannels, current.sampwidth, current.framerate, current.comptype, current.compname) != (params.nchannels, params.sampwidth, params.framerate, params.comptype, params.compname):
                raise RuntimeError(f"{book} chapter {chapter_index}: WAV format changed mid-chapter")
            chunk_path.write_bytes(raw)
            meta_path.write_text(json.dumps({"nchannels": current.nchannels, "sampwidth": current.sampwidth, "framerate": current.framerate, "nframes": len(raw) // current.sampwidth // current.nchannels, "comptype": current.comptype, "compname": current.compname}))
            pcm_parts.append(raw)
            if (index + 1) % 25 == 0:
                print(f"[{book}] chapter {chapter_index} {index + 1}/{len(chunks)} chunks", flush=True)
        if not pcm_parts or params is None:
            continue
        merged_parts = []
        for i, part in enumerate(pcm_parts):
            if i:
                pause_ms = chunks[i - 1][2]
                silence = b"\0" * int(params.framerate * pause_ms / 1000) * params.sampwidth * params.nchannels
                merged_parts.append(silence)
            merged_parts.append(part)
        merged = b"".join(merged_parts)
        wav = make_wav(merged, params)
        encoded = request_json("POST", f"{BASE}/api/audio/encode-mp3", data=wav, timeout=180).content
        filename = f"{safe(book, 45)} - {chapter_index:02d} - {safe(title, 55)}.mp3"
        upload_url = f"{BASE}/api/audiobook-export/{quote(book, safe='')}/{quote(filename, safe='')}"
        request_json("POST", upload_url, data=encoded, timeout=180)
        (chapter_dir / "complete.json").write_text(json.dumps({"filename": filename, "title": title, "chunks": len(chunks), "narrator_fallbacks": fallback, "bytes": len(encoded)}))
        print(f"[{book}] COMPLETE chapter {chapter_index}/{len(buckets)}: {filename} ({len(encoded) / 1048576:.1f} MiB)", flush=True)


def main() -> None:
    docs = request_json("GET", f"{BASE}/api/reader/docs").json().get("docs", [])
    by_title = {doc.get("title"): doc for doc in docs}
    for title in BOOK_ORDER:
        if title not in by_title:
            raise RuntimeError(f"Book document not found: {title}")
        render_book(by_title[title])
    (CACHE / "ALL_THREE_BOOKS_COMPLETE").write_text(time.strftime("%Y-%m-%dT%H:%M:%S%z"))
    print("ALL THREE AUDIOBOOKS COMPLETE", flush=True)
    if os.environ.get("RUN_IMAGE_SAMPLES") == "1":
        print("Restarting Voice Creator before the ten image samples…", flush=True)
        subprocess.run(["docker", "compose", "restart", "tts-voice-creator-clone-and-design-2"], cwd=ROOT, check=True)
        time.sleep(3)
        subprocess.run([sys.executable, str(ROOT / "scripts" / "generate_character_image_samples.py")], cwd=ROOT, check=True)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("Interrupted; completed chunks remain cached for resume.", file=sys.stderr)
        raise
    except Exception as exc:
        print(f"BATCH FAILED: {exc}", file=sys.stderr)
        raise