#!/usr/bin/env python3
"""Read a plain text file aloud with the lektor_ai voice — no video, no SRT.

The dubbing pipeline is built around subtitles: every clip is placed at a
timestamp against a film. Prose has no timings, so this script does the one
thing that is genuinely missing — it splits the text into sentence-sized
chunks itself — and then reuses lektor_ai_comfy for everything else: the same
dictionary, the same Polish number normalizer, the same ComfyUI TTS call with
its truncation retries and crash recovery, the same per-clip trim and speed-up.
The clips are laid end to end with a short pause instead of at an SRT time.

    python3 tts_txt.py input/oko_proroka.txt
    python3 tts_txt.py input/oko_proroka.txt -o output/oko.mp3 --speed 1.1
    python3 tts_txt.py input/oko_proroka.txt --dry-run     # just show the chunks

Clips live in output/clips/<stem>/ exactly as for a film, so an interrupted
run resumes with --resume-from N and a re-assembly costs nothing.
"""
import argparse
import random
import re
import sys
from pathlib import Path

import lektor_ai_comfy as lek

# Chunk size is a compromise, not a limit of the model: longer chunks read with
# better prosody but are likelier to come back cut off mid-word, and a retry
# then costs the whole chunk. Subtitle blocks — what the voice was tuned on —
# sit well under this.
MAX_CHARS = 240

SENTENCE_GAP  = 0.35    # s of silence between chunks
PARAGRAPH_GAP = 0.90    # s at a blank line

# A dot after one of these is not the end of a sentence. Single letters are
# handled separately ("p. kasztelana", initials).
ABBREV = {
    "np.", "itd.", "itp.", "tzn.", "tj.", "ok.", "dr.", "prof.", "św.",
    "ul.", "nr.", "str.", "m.in.", "tzw.", "ds.", "godz.", "mln.", "mld.",
}

_SENT_SPLIT = re.compile(r'(?<=[.!?…])["»”\'’)\]]*\s+')
_CLOSERS = '"»”\'’)]'
_CLAUSE_SPLIT = re.compile(r'(?<=[,;:])\s+|\s+(?=—\s)')


def _terminal_punct(text: str) -> str:
    """Give the chunk an ending the model can land on.

    lek._ensure_terminal_punct is written for subtitles and does not know the
    single-character ellipsis or a closing quote, so it would turn `uciekło…`
    into `uciekło……` and `jeść!”` into `jeść!”...`.
    """
    if not text:
        return text
    bare = text.rstrip(_CLOSERS)
    if bare and bare[-1] in '.,!?;:…':
        return text
    return text + '…'


def _prepare(text: str, dictionary) -> str:
    text = lek._strip_sub_tags(text)
    # A chunk cut out of the middle of a long sentence starts lowercase, which
    # the TTS reads with the wrong intonation — the pipeline capitalises these
    # for the same reason.
    if text and text[0].islower():
        text = text[0].upper() + text[1:]
    if dictionary is not None:
        text = dictionary.apply(text)
    return _terminal_punct(text)


def _is_real_boundary(before: str) -> bool:
    """False when the dot that ended `before` belongs to an abbreviation."""
    tail = before.rstrip('"»”\'’)]').split()[-1] if before.split() else ""
    if not tail.endswith("."):
        return True                       # ended on ! ? … — always a boundary
    if tail.lower() in ABBREV:
        return False
    stem = tail[:-1]
    return not (len(stem) == 1 and stem.isalpha())     # "p.", initials


def split_sentences(paragraph: str) -> list:
    parts, pos = [], 0
    for m in _SENT_SPLIT.finditer(paragraph):
        candidate = paragraph[pos:m.start()]
        if not _is_real_boundary(candidate):
            continue
        parts.append(paragraph[pos:m.end()].strip())
        pos = m.end()
    tail = paragraph[pos:].strip()
    if tail:
        parts.append(tail)
    return parts


def _split_long(sentence: str, max_chars: int) -> list:
    """A sentence longer than one chunk, broken at commas and dashes.

    Nineteenth-century prose runs to sentences of several hundred characters;
    handing one to the model whole is exactly the case that comes back
    truncated.
    """
    pieces, current = [], ""
    for part in _CLAUSE_SPLIT.split(sentence):
        if current and len(current) + 1 + len(part) > max_chars:
            pieces.append(current)
            current = part
        else:
            current = f"{current} {part}".strip()
    if current:
        pieces.append(current)
    return pieces


def chunk_text(raw: str, max_chars: int = MAX_CHARS) -> list:
    """Text -> [(chunk_text, ends_paragraph)].

    Hard line wraps inside a paragraph are joined: they are an artefact of how
    the file was typeset, not pauses the reader should hear. A blank line is a
    real break and becomes the longer pause.
    """
    chunks = []
    paragraphs = [p for p in re.split(r'\n\s*\n', raw) if p.strip()]
    for para in paragraphs:
        joined = re.sub(r'\s*\n\s*', ' ', para.strip())
        joined = re.sub(r'([?!.…])(—)', r'\1 \2', joined)   # "jak?— Ja" -> "jak? — Ja"
        joined = re.sub(r'\s{2,}', ' ', joined)
        buf = []
        for sentence in split_sentences(joined):
            for piece in ([sentence] if len(sentence) <= max_chars
                          else _split_long(sentence, max_chars)):
                if buf and len(" ".join(buf)) + 1 + len(piece) > max_chars:
                    chunks.append([" ".join(buf), False])
                    buf = []
                buf.append(piece)
        if buf:
            chunks.append([" ".join(buf), False])
        if chunks:
            chunks[-1][1] = True          # last chunk of the paragraph
    return [(t, end) for t, end in chunks]


def main():
    pre = argparse.ArgumentParser(add_help=False)
    pre.add_argument("--settings", default=str(lek.SETTINGS_FILE), metavar="FILE")
    settings_path = Path(pre.parse_known_args()[0].settings).expanduser()
    cfg = lek.load_settings(settings_path)

    p = argparse.ArgumentParser(
        description="Turn a plain text file into a lektor-voiced audio file.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument("text", help="UTF-8 .txt file to read aloud")
    p.add_argument("-o", "--output", default=None, metavar="FILE",
                   help="output audio (extension picks the format); "
                        "default <output_dir>/<stem>.mp3")
    p.add_argument("--settings", default=str(lek.SETTINGS_FILE), metavar="FILE")
    p.add_argument("--voice", default=cfg["voice"],
                   help="voice name in voices/ or a path to a WAV")
    # 1.0 rather than the settings default: the pipeline's speed-up exists to
    # fit a clip into its subtitle slot, and prose has no slot to fit.
    p.add_argument("--speed", type=float, default=1.0, metavar="FACTOR")
    p.add_argument("--temperature", type=float, default=cfg["temperature"], metavar="T")
    p.add_argument("--max-tokens", type=int, default=cfg["max_tokens"], metavar="N")
    p.add_argument("--seed", type=int, default=0, help="0 = random")
    p.add_argument("--truncation-retries", type=int,
                   default=cfg["truncation_retries"], metavar="N")
    p.add_argument("--strip-leading-ms", type=float, default=cfg["strip_leading_ms"])
    p.add_argument("--trailing-fade-ms", type=float, default=cfg["trailing_fade_ms"])
    p.add_argument("--dictionary", default=cfg["dictionary"], metavar="FILE",
                   help='replacement list; "" disables')
    p.add_argument("--dictionary-extra", default=None, metavar="FILE",
                   help="second list merged on top, for rules specific to this text")
    numbers = p.add_mutually_exclusive_group()
    numbers.add_argument("--normalize-numbers", dest="normalize_numbers",
                         action="store_true", default=cfg["normalize_numbers"])
    numbers.add_argument("--no-normalize-numbers", dest="normalize_numbers",
                         action="store_false")
    p.add_argument("--max-chars", type=int, default=MAX_CHARS, metavar="N",
                   help="longest chunk handed to the model in one go")
    p.add_argument("--sentence-gap", type=float, default=SENTENCE_GAP, metavar="S")
    p.add_argument("--paragraph-gap", type=float, default=PARAGRAPH_GAP, metavar="S")
    p.add_argument("--limit", type=int, default=0, metavar="N",
                   help="generate only the first N chunks (for a quick listen)")
    p.add_argument("--resume-from", type=int, default=0, metavar="N",
                   help="reuse clips 1..N-1 already in the job folder")
    p.add_argument("--rebuild", action="store_true",
                   help="re-assemble from the clips on disk, no TTS")
    p.add_argument("--dry-run", action="store_true",
                   help="print the chunks and exit without touching ComfyUI")
    p.add_argument("--output-dir", default=cfg["output_dir"])
    p.add_argument("--comfy-url", default=cfg["comfy_url"])
    p.add_argument("--no-comfy-autostart", action="store_true")
    args = p.parse_args()

    lek.COMFY_URL = args.comfy_url
    lek.COMFY_AUTOSTART = cfg["comfy_autostart"] and not args.no_comfy_autostart
    start_cmd = Path(cfg["comfy_start_command"]).expanduser()
    lek.COMFY_START_SCRIPT = (start_cmd if start_cmd.is_absolute()
                              else Path(lek.__file__).resolve().parent / start_cmd)

    src = Path(args.text)
    if not src.is_file():
        sys.exit(f"No such text file: {src}")
    out_dir = Path(args.output_dir)
    clips_dir = lek.job_dir(out_dir, str(src))
    clips_dir.mkdir(parents=True, exist_ok=True)
    final = Path(args.output) if args.output else out_dir / f"{src.stem}.mp3"

    print(f"  Settings: {settings_path if settings_path.is_file() else 'built-in defaults'}")
    print(f"  Job folder: {clips_dir}")

    # ── Step 1: Text -> blocks ────────────────────────────────────────────────
    print(f"\n[1/3] Splitting {src.name}…")
    chunks = chunk_text(src.read_text(encoding="utf-8-sig"), args.max_chars)
    if not chunks:
        sys.exit(f"{src} has no text.")
    if args.limit:
        chunks = chunks[:args.limit]
        print(f"      --limit {args.limit}: only the first {len(chunks)} chunk(s).")

    wanted = []
    if args.dictionary.strip():
        wanted.append(lek.resolve_dictionary_path(args.dictionary))
    if args.dictionary_extra and args.dictionary_extra.strip():
        wanted.append(Path(args.dictionary_extra).expanduser())
    dictionary = lek.Dictionary.load_all(wanted) if wanted else None
    if dictionary is not None:
        for path, count in dictionary.sources:
            print(f"  Dictionary: {count} rule(s) from {path}")

    # The synthetic timings are placeholders — every clip asks to start at 0 and
    # generate_clips pushes it after the previous one. The real spacing is
    # applied after generation, when the durations are known.
    blocks, gaps = [], []
    for text, ends_para in chunks:
        clean = _prepare(text, dictionary)
        blocks.append((0.0, 0.0, clean))
        gaps.append(args.paragraph_gap if ends_para else args.sentence_gap)
    if dictionary is not None and dictionary.hits:
        print(f"      {dictionary.hits} dictionary replacement(s) applied.")


    if args.normalize_numbers:
        py = Path(cfg["normalizer_python"]).expanduser()
        if not py.is_absolute():
            py = Path(lek.__file__).resolve().parent / py
        llm = Path(cfg["normalizer_llm"]).expanduser() if cfg["normalizer_llm"] else None
        if llm and not llm.is_absolute():
            llm = Path(lek.__file__).resolve().parent / llm
        blocks = lek.normalize_block_numbers(blocks, py, cfg["normalizer_timeout"],
                                             free_vram=cfg["free_comfy_vram"],
                                             llm_dir=llm)

    chars = sum(len(t) for _s, _e, t in blocks)
    print(f"      {len(blocks)} chunk(s), {chars} characters "
          f"(longest {max(len(t) for _s, _e, t in blocks)}).")

    preview = clips_dir / "tts_input.txt"
    preview.write_text("\n\n".join(f"[{i + 1}] {t}" for i, (_s, _e, t) in enumerate(blocks)),
                       encoding="utf-8")
    print(f"      Saved to {preview} for inspection.")

    if args.dry_run:
        for i, (_s, _e, text) in enumerate(blocks, 1):
            print(f"  [{i:4d}] ({len(text):3d}) {text}")
        return

    # ── Step 2: TTS ───────────────────────────────────────────────────────────
    if args.rebuild:
        clips = [(0.0, path) for path in
                 sorted(clips_dir.glob("clip_[0-9][0-9][0-9][0-9].flac"))]
        if not clips:
            sys.exit(f"--rebuild: no clips in {clips_dir}.")
        print(f"\n[2/3] Rebuilding from {len(clips)} existing clip(s) — no TTS.")
    else:
        print(f"\n[2/3] Generating TTS ({len(blocks)} clips, speed {args.speed}×)…")
        wav_path = Path(args.voice)
        if not wav_path.is_file():
            wav_path = lek.VOICE_PROMPTS_DIR / f"{args.voice}.wav"
        if not wav_path.is_file():
            sys.exit(f"Voice WAV not found: {wav_path}")
        print(f"  Voice: {wav_path.name}")

        # A cold backend answers the first submission with a connection error,
        # which run_workflow would count against the clip's crash budget.
        if not lek.wait_for_backend(2.0):
            print("  ComfyUI is not up — starting it…")
            if not (lek.COMFY_AUTOSTART and lek.start_comfy_backend(lek.COMFY_START_SCRIPT)
                    and lek.wait_for_backend(lek.BACKEND_HEALTH_TIMEOUT)):
                sys.exit(f"ComfyUI did not come up at {lek.COMFY_URL}.")

        base_seed = args.seed if args.seed != 0 else random.randint(1, 2**31 - 1)
        print(f"  Base seed: {base_seed}")
        voice_filename, voice_ref_text = lek.upload_voice_wav(wav_path)

        clips = lek.generate_clips(blocks, voice_filename, voice_ref_text,
                                   base_seed, args.speed, clips_dir,
                                   strip_leading_ms=args.strip_leading_ms,
                                   resume_from=args.resume_from,
                                   temperature=args.temperature,
                                   trailing_fade_ms=args.trailing_fade_ms,
                                   max_new_tokens=args.max_tokens,
                                   retries=args.truncation_retries)
        if cfg["free_comfy_vram"]:
            lek.free_comfy_vram()

    # ── Step 3: Lay the clips end to end ─────────────────────────────────────
    print(f"\n[3/3] Assembling…")
    timed, cursor = [], 0.0
    for _start, path in clips:
        # generate_clips skips a clip that repeatedly killed the backend, so the
        # gap belongs to the block the clip came from, not to its position here.
        idx = int(re.search(r'(\d+)', path.stem).group(1)) - 1
        timed.append((cursor, path))
        cursor += lek.get_audio_duration(path)
        cursor += gaps[idx] if idx < len(gaps) else args.sentence_gap

    assembled = clips_dir / "tts_assembled.flac"
    lek.assemble_timed_audio(timed, str(assembled))

    final.parent.mkdir(parents=True, exist_ok=True)
    if final.suffix.lower() == ".flac":
        assembled.replace(final)
    else:
        lek._run(["ffmpeg", "-y", "-i", str(assembled), "-b:a", "192k", str(final)],
                 f"Encoding {final.name}")

    total = lek.get_audio_duration(final)
    print(f"\nDone! Output: {final}  ({int(total // 60)}:{int(total % 60):02d}, "
          f"{final.stat().st_size // 1024} KB)")


if __name__ == "__main__":
    main()