#!/usr/bin/env python3
"""
lektor_ai_comfy — Convert SRT subtitles to AI voiceover via ComfyUI + Higgs Audio V3.

Usage:
  python3 lektor_ai_comfy.py input/film.srt input/film.mkv --start 00:00 --end 01:40

Requires ComfyUI with TTS-Audio-Suite running at http://127.0.0.1:8189
"""

import argparse
import configparser
import json
import os
import random
import re
import shlex
import shutil
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path
from urllib.request import urlopen, Request

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

COMFY_OUTPUT_DIR  = Path("/home/lukas/comfy-tts-app/ComfyUI/output")
COMFY_INPUT_DIR   = Path("/home/lukas/comfy-tts-app/ComfyUI/input")
VOICE_PROMPTS_DIR = Path(__file__).resolve().parent / "voices"

SETTINGS_FILE = Path(__file__).resolve().parent / "settings.ini"

def _as_bool(raw) -> bool:
    """yes/no setting. configparser's own getboolean is bypassed by the
    per-key cast in load_settings, so accept the same spellings here."""
    value = str(raw).strip().lower()
    if value in ("1", "true", "yes", "on"):
        return True
    if value in ("0", "false", "no", "off"):
        return False
    raise ValueError(f"{raw!r} is not a yes/no value")


# Fallback values, used when settings.ini is absent or omits a key. Keeping them
# here means the script still runs correctly with no settings file at all.
# Each entry is (ini_section, ini_key, type, value).
_DEFAULTS = {
    "voice":            ("tts",   "voice",            str,   "gudowski_short"),
    "speed":            ("tts",   "speed",            float, 1.2),
    "temperature":      ("tts",   "temperature",      float, 0.9),
    "max_tokens":       ("tts",   "max_tokens",       int,   4096),
    "truncation_retries": ("tts", "truncation_retries", int,  2),
    "strip_leading_ms": ("tts",   "strip_leading_ms", float, 0.0),
    "trailing_fade_ms": ("tts",   "trailing_fade_ms", float, 80.0),
    "orig_volume":      ("mix",   "orig_volume",      float, 0.60),
    "tts_volume":       ("mix",   "tts_volume",       float, 1.8),
    "output_dir":       ("paths", "output_dir",       str,   "output"),
    "comfy_url":        ("comfy", "comfy_url",        str,   "http://127.0.0.1:8189"),
    "comfy_autostart":  ("comfy", "autostart",     _as_bool, True),
    "comfy_start_command": ("comfy", "start_command", str,  "start_comfyui.sh"),
}

COMFY_URL = _DEFAULTS["comfy_url"][3]   # rebound from settings/CLI in main()

# Rebound in main() too; used when the backend stops answering.
COMFY_AUTOSTART    = True
COMFY_START_SCRIPT = Path(__file__).resolve().parent / "start_comfyui.sh"


def load_settings(path: Path) -> dict:
    """Read settings.ini, falling back to _DEFAULTS for anything missing.

    Values here become argparse's defaults, so precedence ends up as
    CLI flag > settings.ini > _DEFAULTS without any extra bookkeeping.
    A malformed or unreadable file is reported and then ignored rather than
    being fatal — a typo in a config file should not cost you a long run.
    """
    values = {key: spec[3] for key, spec in _DEFAULTS.items()}
    if not path.is_file():
        return values

    parser = configparser.ConfigParser(inline_comment_prefixes=("#", ";"))
    try:
        parser.read(path, encoding="utf-8")
    except (configparser.Error, OSError) as exc:
        print(f"  WARNING: could not read {path.name} ({exc}) — using built-in defaults",
              file=sys.stderr)
        return values

    for key, (section, option, cast, fallback) in _DEFAULTS.items():
        if not parser.has_option(section, option):
            continue
        raw = parser.get(section, option).strip()
        if not raw:
            continue
        try:
            values[key] = cast(raw)
        except ValueError:
            print(f"  WARNING: {path.name} [{section}] {option}={raw!r} is not a valid "
                  f"{cast.__name__} — using {fallback!r}", file=sys.stderr)
    return values

# Subtitle codecs that carry real text and can therefore be extracted and fed
# to the TTS engine. Used only by the --embedded-subs picker; image-based
# formats (hdmv_pgs_subtitle, dvd_subtitle) hold bitmaps and would need OCR.
# The final mux copies every subtitle track regardless of codec.
_TEXT_SUB_CODECS = {"subrip", "srt", "ass", "ssa", "webvtt", "mov_text"}

# ---------------------------------------------------------------------------
# SRT helpers
# ---------------------------------------------------------------------------

_SRT_TIME_RE  = re.compile(r'(\d+):(\d+):(\d+)[,.](\d+)')



def _parse_srt_time(s: str) -> float:
    m = _SRT_TIME_RE.match(s.strip())
    if not m:
        raise ValueError(f"Bad SRT timestamp: {s!r}")
    h, mn, sec, ms = int(m[1]), int(m[2]), int(m[3]), int(m[4])
    return h * 3600 + mn * 60 + sec + ms / 1000


def _fmt_srt_time(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int(round((seconds % 1) * 1000))
    if ms == 1000:
        ms = 999
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def _parse_mmss(value: str) -> float:
    parts = value.strip().split(":")
    if len(parts) == 2:
        return int(parts[0]) * 60 + float(parts[1])
    if len(parts) == 3:
        return int(parts[0]) * 3600 + int(parts[1]) * 60 + float(parts[2])
    raise ValueError(f"Cannot parse time: {value!r}")


def _ensure_terminal_punct(text: str, single_dot: bool = False) -> str:
    if not text or text[-1] in ('.', ',', '!', '?', ';', ':'):
        return text
    return text + ('.' if single_dot else '...')


def _strip_sub_tags(line: str) -> str:
    """Remove SRT/ASS formatting codes from a subtitle line."""
    line = re.sub(r'\{[^}]*\}', '', line)   # ASS overrides: {\an8}, {\pos(…)}, {\c&H…&}
    line = re.sub(r'<[^>]+>', '', line)      # HTML tags: <i>, </i>, <b>, <font …>
    return line


def _clean_text(lines: list) -> str:
    # Strip formatting tags first so they don't interfere with allcaps detection
    # or the leading-dash check.
    lines = [_strip_sub_tags(l) for l in lines]
    # Only split into multi-part (for dialogue splitting later) when lines
    # have explicit speaker dashes. Wrapped single-speaker lines join with space.
    has_dialogue = any(re.match(r"^-\s*\S", l.strip()) for l in lines if l.strip())
    out = []
    for line in lines:
        line = line.strip()
        line = re.sub(r"^-\s*", "", line)
        if line:
            out.append(line)
    result = "\n".join(out) if has_dialogue else " ".join(out)
    # If the whole block is ALL-CAPS, convert to sentence case and end with a single dot.
    letters = [c for c in result if c.isalpha()]
    was_allcaps = bool(letters) and all(c == c.upper() for c in letters)
    if was_allcaps:
        result = result.lower()
    # Sentence fragments that continue from a previous subtitle start with a
    # lowercase letter and confuse TTS models — capitalise the first character.
    if result and result[0].islower():
        result = result[0].upper() + result[1:]
    # Ensure each sentence ends with punctuation so TTS produces a natural tail.
    if has_dialogue:
        result = "\n".join(_ensure_terminal_punct(p, was_allcaps) for p in result.split("\n"))
    else:
        result = _ensure_terminal_punct(result, was_allcaps)
    return result


def load_tts_input_srt(path: Path) -> list:
    """Read a pre-processed tts_input.srt into (t_start, t_end, text) tuples.
    Unlike parse_subtitle_blocks this does NOT re-run _clean_text or dialogue
    splitting — the file is already in final form.
    """
    text = path.read_text(encoding="utf-8")
    blocks = re.split(r"\n{2,}", text.strip())
    result = []
    for block in blocks:
        lines = block.strip().splitlines()
        ti = next((i for i, l in enumerate(lines) if "-->" in l), None)
        if ti is None:
            continue
        parts = lines[ti].split("-->")
        if len(parts) != 2:
            continue
        t_start = _parse_srt_time(parts[0])
        t_end   = _parse_srt_time(parts[1])
        payload = " ".join(l.strip() for l in lines[ti + 1:] if l.strip())
        if payload:
            result.append((t_start, t_end, payload))
    return result


def parse_subtitle_blocks(srt_path: str, start_sec: float, end_sec: float) -> list:
    """Return list of (start_relative, end_relative, plain_text) blocks."""
    text = Path(srt_path).read_text(encoding="utf-8-sig")
    blocks = re.split(r"\n{2,}", text.strip())
    result = []
    for block in blocks:
        lines = block.strip().splitlines()
        ti = next((i for i, l in enumerate(lines) if "-->" in l), None)
        if ti is None:
            continue
        parts = lines[ti].split("-->")
        if len(parts) != 2:
            continue
        t_start = _parse_srt_time(parts[0])
        t_end   = _parse_srt_time(parts[1])
        if t_start < start_sec or t_start > end_sec:
            continue
        plain = _clean_text(lines[ti + 1:])
        if plain:
            result.append((t_start - start_sec, t_end - start_sec, plain))
    return result


def split_dialogue_blocks(blocks: list) -> list:
    """Split multi-line blocks into one entry per line with proportional timing."""
    result = []
    for t_start, t_end, text in blocks:
        lines = [l for l in text.split("\n") if l.strip()]
        if len(lines) <= 1:
            result.append((t_start, t_end, text))
            continue
        duration = t_end - t_start
        total_chars = sum(len(l) for l in lines)
        current = t_start
        for i, line in enumerate(lines):
            share = len(line) / total_chars if total_chars else 1 / len(lines)
            end = t_end if i == len(lines) - 1 else current + duration * share
            result.append((current, end, line))
            current = end
    return result


def blocks_to_srt(blocks: list) -> str:
    """Convert (start, end, text) tuples back to SRT string for inspection."""
    lines = []
    for i, (t_start, t_end, text) in enumerate(blocks, 1):
        lines.append(str(i))
        lines.append(f"{_fmt_srt_time(t_start)} --> {_fmt_srt_time(t_end)}")
        lines.append(text)
        lines.append("")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# FFmpeg helpers
# ---------------------------------------------------------------------------

def _ffmpeg_ts(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}"


def _run(cmd: list, label: str):
    print(f"  {label}...")
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"ffmpeg error ({label}):\n{result.stderr[-2000:]}", file=sys.stderr)
        sys.exit(1)


def _run_progress(cmd: list, label: str, total_seconds: float):
    """Run ffmpeg with a live progress bar that rewrites the current line."""
    # Inject -progress pipe:1 before the output path (last arg)
    prog_cmd = cmd[:-1] + ["-progress", "pipe:1", "-nostats", cmd[-1]]
    proc = subprocess.Popen(
        prog_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, bufsize=1,
    )
    stderr_chunks: list = []
    def _drain():
        for line in proc.stderr:
            stderr_chunks.append(line)
    threading.Thread(target=_drain, daemon=True).start()

    BAR = 28
    print(f"  {label}...", end="", flush=True)
    for line in proc.stdout:
        if not line.startswith("out_time_us="):
            continue
        try:
            us = int(line.split("=", 1)[1])
            if us < 0:
                continue
            sec = us / 1_000_000
        except (ValueError, IndexError):
            continue
        pct   = min(sec / total_seconds, 1.0) if total_seconds > 0 else 0.0
        filled = int(BAR * pct)
        bar    = "█" * filled + "░" * (BAR - filled)
        t_cur  = f"{int(sec // 60):02d}:{int(sec % 60):02d}"
        t_tot  = f"{int(total_seconds // 60):02d}:{int(total_seconds % 60):02d}"
        print(f"\r  {label}... [{bar}] {t_cur}/{t_tot}", end="", flush=True)

    proc.wait()
    if proc.returncode != 0:
        err = "".join(stderr_chunks)
        print(f"\nffmpeg error ({label}):\n{err[-2000:]}", file=sys.stderr)
        sys.exit(1)
    t_tot = f"{int(total_seconds // 60):02d}:{int(total_seconds % 60):02d}"
    print(f"\r  {label}... [{'█' * BAR}] {t_tot}/{t_tot}")


def get_audio_duration(path: Path) -> float:
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", str(path)],
        capture_output=True, text=True,
    )
    info = json.loads(result.stdout)
    for stream in info.get("streams", []):
        if "duration" in stream:
            return float(stream["duration"])
    raise RuntimeError(f"Cannot determine duration of {path}")


def speed_up_audio(input_path: Path, output_path: Path, speed: float,
                   fade_out_ms: float = 0):
    """Apply atempo (and optional fade-out) to a single clip."""
    filters = []
    if speed != 1.0:
        remaining = speed
        while remaining > 2.0:
            filters.append("atempo=2.0")
            remaining /= 2.0
        filters.append(f"atempo={remaining:.4f}")
    if fade_out_ms > 0:
        fade_sec = fade_out_ms / 1000
        # Compute fade start from the post-speed duration so we can use a plain
        # afade=t=out (areverse caused MP3 tail truncation on VBR files).
        raw_dur = get_audio_duration(input_path)
        post_dur = raw_dur / speed if speed != 1.0 else raw_dur
        fade_start = max(0.0, post_dur - fade_sec)
        filters.append(f"afade=t=out:st={fade_start:.3f}:d={fade_sec:.3f}")
    if not filters:
        shutil.copy2(input_path, output_path)
        return
    label = f"Speed {speed:.2f}×"
    if fade_out_ms:
        label += f" + fade-out {fade_out_ms:.0f}ms"
    _run([
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-filter:a", ",".join(filters),
        str(output_path),
    ], label)


def trim_audio(input_path: Path, output_path: Path, max_duration: float):
    """Hard-trim a clip to max_duration seconds."""
    _run([
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-t", f"{max_duration:.3f}",
        "-c", "copy",
        str(output_path),
    ], f"Trim to {max_duration:.2f}s")


def strip_leading_audio(input_path: Path, output_path: Path, strip_ms: float):
    """Remove the first strip_ms milliseconds from a clip."""
    _run([
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-ss", f"{strip_ms / 1000:.3f}",
        "-c", "copy",
        str(output_path),
    ], f"Strip {strip_ms:.0f}ms leading audio")


# A clip's trailing silence must not count towards its duration: placement uses
# a push-forward cursor, so padding at the end of one clip delays every clip
# after it. Usually that padding is a fraction of a second, but when the model
# fails to emit a stop token it generates until --max-tokens runs out — a short
# line then yields a clip of speech followed by two minutes of digital silence,
# which shunts the rest of the film out of sync.
SILENCE_DB       = -45.0    # anything quieter counts as silence
SILENCE_MIN      = 0.35     # ignore shorter gaps; commas and beats are not padding
SILENCE_KEEP_MS  = 120      # leave this much tail so speech does not end abruptly


def trailing_silence_start(path: Path) -> float:
    """Time at which the clip's final run of silence begins, or None."""
    result = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(path),
         "-af", f"silencedetect=noise={SILENCE_DB}dB:d={SILENCE_MIN}",
         "-f", "null", "-"],
        capture_output=True, text=True,
    )
    try:
        duration = get_audio_duration(path)
    except RuntimeError:
        return None

    open_at = None
    for kind, value in re.findall(r"silence_(start|end):\s*([0-9.]+)", result.stderr):
        t = float(value)
        if kind == "start":
            open_at = t
        elif t < duration - 0.05:
            open_at = None          # that silence ended mid-clip, so it is a pause
    return open_at


# Some runaway generations pad with a steady tone instead of digital silence.
# It sits above any sane silence threshold, so silencedetect cannot see it —
# but real speech swings tens of dB between windows while a tone does not.
# Measured here: tone tails spread 3-5 dB, speech tails 43-61 dB.
FLAT_TAIL_DB      = 15.0    # spread below this is not speech
FLAT_TAIL_MIN     = 5.0     # the flat run must last at least this long
RUNAWAY_SECONDS   = 20.0    # only look for tone padding on clips this long


# A clip the model finished decays to near silence; one it cut off mid-word
# stops at full level. Measured across 80 raw outputs the two groups sit about
# 50 dB apart (-65 dB vs -16 dB), so the test is unambiguous.
TRUNCATION_PEAK_DB = -25.0
TRUNCATION_WINDOW  = 0.08


def ends_abruptly(path: Path, window: float = TRUNCATION_WINDOW) -> bool:
    """True when a clip stops at full level, i.e. the model cut it mid-word."""
    try:
        duration = get_audio_duration(path)
    except RuntimeError:
        return False
    if duration <= window:
        return False
    result = subprocess.run(
        ["ffmpeg", "-hide_banner", "-ss", f"{duration - window:.3f}",
         "-t", f"{window}", "-i", str(path),
         "-af", "volumedetect", "-f", "null", "-"],
        capture_output=True, text=True,
    )
    peak = re.search(r"max_volume:\s*(-?[0-9.]+)", result.stderr)
    return bool(peak) and float(peak.group(1)) > TRUNCATION_PEAK_DB


def flat_tail_start(path: Path, window: float = 0.5) -> float:
    """Start of a long, unnaturally constant tail, or None."""
    result = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(path),
         "-af", f"astats=metadata=1:reset={window},"
                "ametadata=print:key=lavfi.astats.Overall.RMS_level",
         "-f", "null", "-"],
        capture_output=True, text=True,
    )
    blob = result.stderr + result.stdout
    values = [(-120.0 if "inf" in m else float(m))
              for m in re.findall(r"RMS_level=(-?[0-9.]+|-?inf)", blob)]
    need = int(FLAT_TAIL_MIN / window)
    if len(values) < need + 1:
        return None
    for i in range(len(values) - need):
        tail = values[i:]
        if max(tail) - min(tail) <= FLAT_TAIL_DB:
            return i * window
    return None


def trim_trailing_silence(input_path: Path, output_path: Path) -> float:
    """Cut the final silence off a clip. Returns the seconds removed."""
    duration = get_audio_duration(input_path)
    start = trailing_silence_start(input_path)
    if start is None and duration > RUNAWAY_SECONDS:
        start = flat_tail_start(input_path)
    keep = None if start is None else start + SILENCE_KEEP_MS / 1000
    if keep is None or keep >= duration - 0.02 or keep <= 0.05:
        if input_path != output_path:
            shutil.copy2(input_path, output_path)
        return 0.0
    _run([
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-t", f"{keep:.3f}",
        str(output_path),
    ], f"Trim {duration - keep:.1f}s trailing silence")
    return duration - keep


def _assemble_batch(clips: list, output_path: str, label: str,
                    total_seconds: float = 0.0, start_offset: float = 0.0):
    """Mix a batch of (start_sec, path) clips with adelay into output_path.

    start_offset is subtracted from every clip's start time so sub-batch files
    contain no leading silence. The caller re-applies the offset at merge time.
    """
    run = (lambda cmd, lbl: _run_progress(cmd, lbl, total_seconds)) if total_seconds > 0 else _run
    if len(clips) == 1:
        delay_ms = max(0, int((clips[0][0] - start_offset) * 1000))
        run([
            "ffmpeg", "-y",
            "-i", str(clips[0][1]),
            "-af", f"adelay={delay_ms}:all=1",
            output_path,
        ], label)
        return
    inputs, filter_parts = [], []
    for i, (start_sec, path) in enumerate(clips):
        inputs += ["-i", str(path)]
        delay_ms = max(0, int((start_sec - start_offset) * 1000))
        filter_parts.append(f"[{i}:a]adelay={delay_ms}:all=1[a{i}]")
    mix_inputs = "".join(f"[a{i}]" for i in range(len(clips)))
    filter_parts.append(
        f"{mix_inputs}amix=inputs={len(clips)}:duration=longest:normalize=0[out]"
    )
    run(
        ["ffmpeg", "-y"] + inputs +
        ["-filter_complex", ";".join(filter_parts), "-map", "[out]", output_path],
        label,
    )


def assemble_timed_audio(clips: list, output_path: str, batch_size: int = 500):
    """Place each clip at its start time via adelay then mix.
    clips: list of (start_sec, Path).
    Processes in batches to stay under OS file-descriptor limits.
    """
    n_batches = (len(clips) - 1) // batch_size + 1
    if len(clips) <= batch_size:
        total_sec = clips[-1][0] if clips else 0.0
        _assemble_batch(clips, output_path,
                        f"Assembling {len(clips)} timed clips",
                        total_seconds=total_sec)
        return

    tmp_dir = Path(output_path).parent / "_batch_tmp"
    tmp_dir.mkdir(exist_ok=True)
    batch_files = []
    batch_offsets = []
    try:
        for b, start in enumerate(range(0, len(clips), batch_size)):
            batch = clips[start:start + batch_size]
            batch_offset = batch[0][0]          # first clip's global start time
            total_sec    = batch[-1][0] - batch_offset  # content duration only
            batch_path   = str(tmp_dir / f"batch_{b:04d}.flac")
            _assemble_batch(
                batch, batch_path,
                f"Assembling batch {b + 1}/{n_batches} ({len(batch)} clips)",
                total_seconds=total_sec,
                start_offset=batch_offset,
            )
            batch_files.append(batch_path)
            batch_offsets.append(batch_offset)
        # Merge: re-apply each batch's offset so clips land at their global times
        batch_clips = [(offset, Path(p)) for offset, p in zip(batch_offsets, batch_files)]
        _assemble_batch(batch_clips, output_path, f"Merging {len(batch_files)} batches")
    finally:
        for p in batch_files:
            Path(p).unlink(missing_ok=True)
        tmp_dir.rmdir() if tmp_dir.exists() and not any(tmp_dir.iterdir()) else None


def extract_video_fragment(video_path: str, start_sec: float, end_sec: float,
                           output_path: str):
    _run([
        "ffmpeg", "-y",
        "-ss", _ffmpeg_ts(start_sec),
        "-i", video_path,
        "-t", str(end_sec - start_sec),
        "-c", "copy",
        output_path,
    ], f"Extracting video {_ffmpeg_ts(start_sec)}–{_ffmpeg_ts(end_sec)}")


def list_subtitle_streams(video_path: str) -> list:
    """Return subtitle stream info dicts for all subtitle tracks in a video file."""
    result = subprocess.run(
        ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", video_path],
        capture_output=True, text=True,
    )
    subs = []
    for s in json.loads(result.stdout).get("streams", []):
        if s.get("codec_type") != "subtitle":
            continue
        tags = s.get("tags", {})
        subs.append({
            "index":      s["index"],
            "codec_name": s.get("codec_name", "unknown"),
            "language":   tags.get("language", ""),
            "title":      tags.get("title", ""),
        })
    return subs


def extract_subtitle_track(video_path: str, stream_index: int, output_path: str):
    """Extract one subtitle track from a video to an SRT file."""
    _run([
        "ffmpeg", "-y",
        "-i", video_path,
        "-map", f"0:{stream_index}",
        "-c:s", "subrip",
        output_path,
    ], f"Extracting subtitle stream #{stream_index}")


def get_media_duration(path: str) -> float:
    """Container duration in seconds (0.0 if unknown)."""
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "csv=p=0", path],
        capture_output=True, text=True,
    )
    try:
        return float(result.stdout.strip())
    except ValueError:
        return 0.0


def merge_voiceover(video_path: str, tts_path: str, output_path: str,
                    orig_volume: float = 0.60, tts_volume: float = 1.8):
    """Produce a video with two audio tracks:
      track 0 — Original (first audio stream from source, stream-copied)
      track 1 — Polish   (lektor mix: TTS at tts_volume + original at orig_volume)
    plus every subtitle track from the source, stream-copied.

    Done in two passes:
      1. Render the lektor mix on its own to a temporary audio file.
      2. Stream-copy everything (video + both audio tracks + all subtitles)
         into the final MKV.

    Combining both in a single ffmpeg call fails on sources with many subtitle
    tracks: the sparse subtitle streams stall the muxer while the audio filter
    graph keeps producing, and ffmpeg eventually aborts with
    "Failed to inject frame into filter network: Invalid argument".
    Splitting the work keeps the filter graph free of subtitle streams, and
    makes the final mux a pure copy (roughly 10x faster too).

    MKV rips often have audio PTS starting at a non-zero offset (e.g. 4.043s).
    Without correction amix can't output from PTS 0 until BOTH inputs have
    frames, which delays the TTS. asetpts resets the PTS to 0 and adelay
    re-inserts the offset as explicit silence, so the background audio stays at
    its correct position while amix gets frames from both inputs from PTS 0.
    """
    probe = subprocess.run(
        ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", video_path],
        capture_output=True, text=True,
    )
    streams = json.loads(probe.stdout).get("streams", [])
    audio_start = next(
        (float(s.get("start_time", 0)) for s in streams if s.get("codec_type") == "audio"),
        0.0,
    )

    if audio_start > 0.01:
        print(f"  Audio PTS offset: {audio_start:.3f}s — normalising orig_mix for lektor track")
        delay_ms = int(audio_start * 1000)
        norm = f"asetpts=PTS-STARTPTS,adelay={delay_ms}:all=1,aresample=async=1"
    else:
        norm = "aresample=async=1"
    # The lektor track is stereo on purpose. ffmpeg's AAC encoder does not
    # signal a 5.1 channel layout — it comes back as "unknown" in every
    # container — so players cannot route the channels and the track plays
    # silent. Stereo is declared correctly and plays everywhere; 192k is also
    # a good rate for two channels where it is a poor one spread over six.
    # The 5.1 source is downmixed with ffmpeg's standard coefficients, and the
    # mono TTS becomes equal-level L/R, i.e. centred.
    stereo = "aformat=channel_layouts=stereo"
    filter_complex = (
        f"[0:a:0]{norm},{stereo},volume={orig_volume}[orig_mix];"
        f"[1:a]{stereo},volume={tts_volume}[tts];"
        "[orig_mix][tts]amix=inputs=2:duration=first:normalize=0[lektor]"
    )

    out = Path(output_path)
    mix_path = out.parent / f"{out.stem}_mix.m4a"
    try:
        # ── Pass 1: render the lektor mix on its own ─────────────────────────
        _run_progress([
            "ffmpeg", "-y",
            "-i", video_path, "-i", tts_path,
            "-filter_complex", filter_complex,
            "-map", "[lektor]",
            "-c:a", "aac", "-b:a", "192k",
            str(mix_path),
        ], "Mixing lektor audio track", get_media_duration(video_path))

        # ── Pass 2: stream-copy mux — video + both audio tracks + subtitles ──
        sub_streams = [s for s in streams if s.get("codec_type") == "subtitle"]
        if sub_streams:
            print(f"  Including {len(sub_streams)} subtitle track(s) in output")

        cmd = [
            "ffmpeg", "-y",
            "-i", video_path, "-i", str(mix_path),
            "-map", "0:v", "-map", "0:a:0", "-map", "1:a:0",
        ]
        for s in sub_streams:
            cmd += ["-map", f"0:{s['index']}"]
        cmd += [
            "-c", "copy",
            # Must be an OUTPUT option (before the output path) — placed before
            # -i it silently configures the demuxer instead and does nothing.
            # Sparse subtitle tracks make the muxer give up on interleaving
            # after its default 10s window, which strands the lektor audio
            # hundreds of MB away from the video it belongs with. Players read
            # MKV sequentially, so they never reach it and the track is silent.
            # 0 means "buffer until every stream has a packet" — correct
            # interleaving, and only a couple of GB of transient buffering.
            "-max_interleave_delta", "0",
            "-metadata:s:a:0", "title=Original",
            "-metadata:s:a:1", "title=Polish",
            "-metadata:s:a:1", "language=pol",
            output_path,
        ]
        _run(cmd, "Muxing final video (2 audio tracks + subtitles)")
    finally:
        mix_path.unlink(missing_ok=True)


# ---------------------------------------------------------------------------
# ComfyUI API helpers
# ---------------------------------------------------------------------------

def _comfy_post(path: str, data: dict) -> dict:
    payload = json.dumps(data).encode()
    req = Request(
        f"{COMFY_URL}{path}",
        data=payload,
        headers={"Content-Type": "application/json"},
    )
    with urlopen(req, timeout=30) as r:
        return json.loads(r.read())


def _comfy_get(path: str) -> dict:
    with urlopen(f"{COMFY_URL}{path}", timeout=30) as r:
        return json.loads(r.read())


def upload_voice_wav(wav_path: Path) -> tuple:
    """Copy voice WAV into ComfyUI's input dir; return (filename, reference_text)."""
    dst = COMFY_INPUT_DIR / wav_path.name
    if not dst.exists() or dst.stat().st_mtime < wav_path.stat().st_mtime:
        shutil.copy2(wav_path, dst)
        print(f"  Copied {wav_path.name} → {dst}")
    txt_path = wav_path.with_suffix(".txt")
    ref_text = txt_path.read_text(encoding="utf-8").strip() if txt_path.exists() else ""
    if ref_text:
        print(f"  Reference text loaded ({len(ref_text)} chars)")
    else:
        print("  WARNING: no reference .txt found — voice cloning quality will be lower")
    return wav_path.name, ref_text


def build_text_workflow(text: str, voice_filename: str, voice_ref_text: str,
                        seed: int, temperature: float = 0.9,
                        max_new_tokens: int = 4096) -> dict:
    """ComfyUI workflow for a single text utterance with proper voice reference."""
    return {
        "1": {
            "class_type": "HiggsAudioV3EngineNode",
            "inputs": {
                "model":          "higgs-audio-v3-tts-4b",
                "device":         "auto",
                "dtype":          "auto",
                "attention":      "auto",
                "temperature":    temperature,
                "top_p":          0.95,
                "top_k":          50,
                "max_new_tokens": max_new_tokens,
            },
        },
        "2": {
            "class_type": "LoadAudio",
            "inputs": {"audio": voice_filename},
        },
        "3": {
            # Pairs audio with its transcript so the model can clone the voice correctly
            "class_type": "CharacterVoicesNode",
            "inputs": {
                "voice_name":      "none",
                "reference_text":  voice_ref_text,
                "opt_audio_input": ["2", 0],
            },
        },
        "4": {
            "class_type": "UnifiedTTSTextNode",
            "inputs": {
                "TTS_engine":        ["1", 0],
                "text":              text,
                "narrator_voice":    "none",
                "seed":              seed,
                "opt_narrator":      ["3", 0],
                "enable_audio_cache": False,
            },
        },
        "5": {
            "class_type": "SaveAudio",
            "inputs": {
                "audio":           ["4", 0],
                "filename_prefix": "audio/lektor_clip",
            },
        },
    }


def submit_workflow(workflow: dict) -> str:
    client_id = str(uuid.uuid4())
    resp = _comfy_post("/prompt", {"prompt": workflow, "client_id": client_id})
    return resp["prompt_id"]


BACKEND_RETRY_SECONDS = 3.0     # how often to look for ComfyUI while it is down


# Launching the backend is guarded: a ComfyUI that dies on startup would
# otherwise spawn a new terminal window every few seconds.
COMFY_AUTOSTART_MAX      = 3      # attempts per run
COMFY_AUTOSTART_COOLDOWN = 90.0   # seconds between attempts
_autostart = {"attempts": 0, "last": 0.0}

_TERMINALS = (
    ("gnome-terminal", lambda c: ["gnome-terminal", "--title", "ComfyUI", "--",
                                  "bash", "-lc", c]),
    ("konsole",        lambda c: ["konsole", "-p", "tabtitle=ComfyUI", "-e",
                                  "bash", "-lc", c]),
    ("xfce4-terminal", lambda c: ["xfce4-terminal", "--title", "ComfyUI",
                                  "-e", f"bash -lc {shlex.quote(c)}"]),
    ("xterm",          lambda c: ["xterm", "-T", "ComfyUI", "-e", "bash", "-lc", c]),
)


def start_comfy_backend(script: Path) -> bool:
    """Launch ComfyUI in its own terminal window. True if one was started."""
    now = time.monotonic()
    if _autostart["attempts"] >= COMFY_AUTOSTART_MAX:
        return False
    if _autostart["last"] and now - _autostart["last"] < COMFY_AUTOSTART_COOLDOWN:
        return False
    if not script.is_file():
        print(f"  Cannot start ComfyUI automatically: {script} not found.",
              file=sys.stderr)
        _autostart["attempts"] = COMFY_AUTOSTART_MAX     # do not keep complaining
        return False

    _autostart["attempts"] += 1
    _autostart["last"] = now
    # Pause on exit so a backend that fails to start leaves its error on screen
    # instead of the window vanishing.
    inner = (f"{shlex.quote(str(script))}; echo; "
             f"echo '[ComfyUI exited — press any key to close]'; read -n 1")

    if os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"):
        for name, build in _TERMINALS:
            if not shutil.which(name):
                continue
            try:
                subprocess.Popen(build(inner), start_new_session=True)
            except OSError:
                continue
            print(f"  Starting ComfyUI in a new {name} window "
                  f"(attempt {_autostart['attempts']} of {COMFY_AUTOSTART_MAX})…",
                  file=sys.stderr)
            return True

    log = script.parent / "comfyui_autostart.log"
    try:
        with open(log, "ab") as handle:
            subprocess.Popen([str(script)], stdout=handle, stderr=handle,
                             start_new_session=True)
    except OSError as exc:
        print(f"  Could not start ComfyUI: {exc}", file=sys.stderr)
        return False
    print(f"  No terminal available — started ComfyUI in the background, "
          f"logging to {log}", file=sys.stderr)
    return True


def _prompt_known(prompt_id: str) -> bool:
    """True while ComfyUI still has this prompt queued, running or in history.

    A restart wipes both, so this is how a job lost to a crash is spotted.
    If the answer cannot be determined, say yes: resubmitting a job that is
    actually still running would waste a generation.
    """
    try:
        if prompt_id in _comfy_get(f"/history/{prompt_id}"):
            return True
        queue = _comfy_get("/queue")
    except Exception:                                   # noqa: BLE001
        return True
    for key in ("queue_running", "queue_pending"):
        for item in queue.get(key, []):
            if isinstance(item, (list, tuple)) and len(item) > 1 and item[1] == prompt_id:
                return True
    return False


def run_workflow(workflow: dict, poll_interval: float = 0.25,
                 max_interval: float = 2.0, backoff_after: float = 10.0) -> dict:
    """Submit a workflow, wait for it, and return its history entry.

    Survives ComfyUI dying mid-job. Its queue and history do not persist across
    a restart, so a prompt submitted beforehand exists nowhere afterwards and
    polling it would wait for ever — the job has to be sent again. While the
    backend is down this reports once rather than once per poll, and slows to
    BACKEND_RETRY_SECONDS instead of hammering a closed port.

    The poll interval, not the GPU, governs perceived speed: a clip that
    finishes 100 ms after a poll still waits for the next one. It was a flat
    3 s, which quantised every clip to a 3 s boundary — across 2075 real clips
    91% of the gaps landed exactly on a multiple of 3 s. So polling starts
    tight and only backs off once a job looks slow.
    """
    prompt_id, offline = None, False
    start, interval, next_dot = time.monotonic(), poll_interval, 1.0

    while True:
        if prompt_id is None:                       # first submission, or a resubmit
            try:
                prompt_id = submit_workflow(workflow)
            except Exception as exc:                # noqa: BLE001
                if not offline:
                    print(f"\n  ComfyUI unreachable ({exc}).", file=sys.stderr)
                    offline = True
                    if not (COMFY_AUTOSTART and start_comfy_backend(COMFY_START_SCRIPT)):
                        print("  Waiting for it to come back — start it and this "
                              "will carry on.", file=sys.stderr)
                time.sleep(BACKEND_RETRY_SECONDS)
                continue
            start, interval, next_dot = time.monotonic(), poll_interval, 1.0
            print(f"    Waiting for {prompt_id[:8]}…", end="", flush=True)

        try:
            history = _comfy_get(f"/history/{prompt_id}")
        except Exception as exc:                    # noqa: BLE001
            if not offline:
                print(f"\n  ComfyUI went away ({exc}).", file=sys.stderr)
                offline = True
                if not (COMFY_AUTOSTART and start_comfy_backend(COMFY_START_SCRIPT)):
                    print("  Waiting for it to come back — restart it and this "
                          "will carry on.", file=sys.stderr)
            time.sleep(BACKEND_RETRY_SECONDS)
            continue

        if offline:                                 # it answered again
            offline = False
            print("  ComfyUI is back.", file=sys.stderr)
            if not _prompt_known(prompt_id):
                print("  The queue did not survive the restart — resubmitting "
                      "this clip.", file=sys.stderr)
                prompt_id = None
                continue

        if prompt_id in history:
            entry = history[prompt_id]
            status = entry.get("status", {})
            for msg_type, msg_data in status.get("messages", []):
                if msg_type == "execution_error":
                    print(f"\nComfyUI error: {msg_data}", file=sys.stderr)
                    sys.exit(1)
            if status.get("completed"):
                print(f" done ({time.monotonic() - start:.1f}s).")
                return entry

        elapsed = time.monotonic() - start
        # One dot per second regardless of poll rate, or the line fills with dots.
        if elapsed >= next_dot:
            print(".", end="", flush=True)
            next_dot += 1.0
        time.sleep(interval)
        if elapsed > backoff_after:
            interval = min(interval * 1.5, max_interval)


def extract_output_audio(history_entry: dict) -> Path:
    outputs = history_entry.get("outputs", {})
    for node_output in outputs.values():
        for audio_info in node_output.get("audio", []):
            subfolder = audio_info.get("subfolder", "")
            path = COMFY_OUTPUT_DIR / subfolder / audio_info["filename"]
            if path.exists():
                return path
    raise RuntimeError(f"Output audio not found in history: {outputs}")


def _generate_one(text: str, voice_filename: str, voice_ref_text: str,
                  seed: int, speed: float, out_dir: Path, label: str,
                  strip_leading_ms: float, temperature: float = 0.9,
                  trailing_fade_ms: float = 0, max_new_tokens: int = 4096,
                  retries: int = 2) -> tuple:
    """Submit one TTS job and return (path, duration_after_speedup, truncated).

    About one clip in twelve comes back cut off mid-word — the model stops
    generating before the line is finished. It is not predictable from the
    text, but it is obvious in the audio, so a clip that ends at full level is
    simply asked for again with a different seed.
    """
    attempt, truncated = 0, False
    while True:
        workflow = build_text_workflow(text, voice_filename, voice_ref_text, seed,
                                       temperature, max_new_tokens)
        raw_clip = extract_output_audio(run_workflow(workflow))

        truncated = ends_abruptly(raw_clip)
        if not truncated or attempt >= retries:
            break
        attempt += 1
        seed += 7919            # a distant seed, not the neighbouring sample
        print(f"    ! ends mid-word — retrying with seed {seed} "
              f"(attempt {attempt + 1} of {retries + 1})")

    if truncated:
        print("    ! still ends mid-word after "
              f"{attempt + 1} attempt{'s' if attempt else ''} — "
              "regenerate this one with --regen")

    if strip_leading_ms > 0:
        stripped = out_dir / f"{label}_s.flac"
        strip_leading_audio(raw_clip, stripped, strip_leading_ms)
        raw_clip = stripped

    raw_duration = get_audio_duration(raw_clip)
    trimmed = out_dir / f"{label}_t.flac"
    removed = trim_trailing_silence(raw_clip, trimmed)
    if removed > 0:
        raw_clip = trimmed
        # Losing most of the clip to silence means the model never stopped and
        # ran to the token ceiling. The speech is usually fine, but say so —
        # it can also mean the line was cut short, which needs a --regen.
        if removed > 20 and removed > raw_duration * 0.5:
            print(f"    ! {removed:.0f}s of trailing silence removed "
                  f"({raw_duration:.0f}s generated for {raw_duration - removed:.1f}s of "
                  f"speech) — the model did not stop on its own; check this clip")

    fast_clip = out_dir / f"{label}.flac"
    speed_up_audio(raw_clip, fast_clip, speed, fade_out_ms=trailing_fade_ms)
    trimmed.unlink(missing_ok=True)

    return fast_clip, get_audio_duration(fast_clip), truncated


def generate_clips(blocks: list, voice_filename: str, voice_ref_text: str,
                   base_seed: int, speed: float, clips_dir: Path,
                   strip_leading_ms: float = 0.0, resume_from: int = 0,
                   temperature: float = 0.9, trailing_fade_ms: float = 0,
                   max_new_tokens: int = 4096, retries: int = 2) -> list:
    """
    Generate one audio clip per block via ComfyUI, apply per-clip speed-up,
    then calculate start positions with "push forward if overlap" logic.

    If resume_from > 0, blocks 1..(resume_from-1) are loaded from existing
    clip files instead of being regenerated (crash-recovery mode).

    Returns list of (actual_start_sec, clip_path).
    """
    clips = []
    cursor = 0.0
    total = len(blocks)
    still_truncated = []

    for idx, (t_start, t_end, text) in enumerate(blocks):
        clip_n = idx + 1
        label = f"clip_{clip_n:04d}"
        canonical = clips_dir / f"{label}.flac"

        if resume_from > 0 and clip_n < resume_from:
            if not canonical.exists():
                sys.exit(f"Resume from block {resume_from}: {canonical} not found — "
                         f"lower --resume-from or regenerate from an earlier block.")
            duration = get_audio_duration(canonical)
            print(f"  [{clip_n}/{total}] skip (resuming) — {canonical.name} {duration:.2f}s")
            actual_start = max(t_start, cursor)
            clips.append((actual_start, canonical))
            cursor = actual_start + duration
            continue

        display = text.replace("\n", " ")
        print(f"  [{clip_n}/{total}] {_fmt_srt_time(t_start)}  {display[:60]!r}")

        seed = base_seed + idx
        fast_clip, duration, truncated = _generate_one(
            text, voice_filename, voice_ref_text,
            seed, speed, clips_dir, label, strip_leading_ms, temperature,
            trailing_fade_ms=trailing_fade_ms, max_new_tokens=max_new_tokens,
            retries=retries,
        )
        if truncated:
            still_truncated.append(clip_n)

        # Place at SRT time or after previous clip, whichever is later.
        # No trimming — clips play to their natural end; re-align at the next
        # natural SRT gap that is wider than the accumulated overrun.
        actual_start = max(t_start, cursor)
        clips.append((actual_start, fast_clip))
        cursor = actual_start + duration

        delay = actual_start - t_start
        delay_str = f" [+{delay:.2f}s behind SRT]" if delay > 0.05 else ""
        print(f"    → placed at {_fmt_srt_time(actual_start)}, duration {duration:.2f}s{delay_str}")

    if still_truncated:
        print(f"\n  {len(still_truncated)} clip(s) still end mid-word after retries. "
              f"To try again:")
        print(f"    --regen {' '.join(str(n) for n in still_truncated)}")

    return clips


# ---------------------------------------------------------------------------
# Rebuild / Regen helpers
# ---------------------------------------------------------------------------

def job_dir(out_dir: Path, video_path: str) -> Path:
    """Per-job working directory: <output_dir>/clips/<video filename without extension>/

    Everything that belongs to one video lives here — the clips, the processed
    tts_input.srt, the assembled audio, any extracted subtitle or video
    fragment. Only the finished <stem>_lektor.mkv is written to <output_dir>.

    Without this separation every job overwrote the previous one's clip_NNNN.flac
    and tts_input.srt, so --rebuild and --regen were only ever safe on the most
    recent run. The clips and tts_input.srt are a matched pair — block N in the
    SRT is clip_NNNN.flac — which is why the SRT lives here too rather than at
    the top level.
    """
    return out_dir / "clips" / Path(video_path).stem


def _legacy_hint(out_dir: Path, job: Path) -> str:
    """Point at the pre-per-job layout if that is what the user actually has."""
    if (out_dir / "tts_input.srt").is_file():
        return (f"\n\nA tts_input.srt exists at {out_dir} in the old shared layout. "
                f"Clips are now kept per job in {job}. Move that run's files there, "
                f"or regenerate it.")
    return ""


def _rebuild_video(args, out_dir: Path):
    """Re-assemble audio from existing clips and rebuild the final video.
    No TTS generation — use this to change volumes or fix assembly without
    waiting for the model to run again.
    """
    job     = job_dir(out_dir, args.video)
    tts_srt = job / "tts_input.srt"
    if not tts_srt.exists():
        sys.exit(f"No tts_input.srt found in {job} — run a full generation first."
                 + _legacy_hint(out_dir, job))

    blocks    = load_tts_input_srt(tts_srt)
    total     = len(blocks)
    clips_dir = job

    missing = [
        clips_dir / f"clip_{i+1:04d}.flac"
        for i in range(total)
        if not (clips_dir / f"clip_{i+1:04d}.flac").exists()
    ]
    if missing:
        sys.exit("Missing clip files — regenerate them first:\n  "
                 + "\n  ".join(str(p) for p in missing[:10]))

    clips  = []
    cursor = 0.0
    for idx, (t_start, t_end, text) in enumerate(blocks):
        clip_path = clips_dir / f"clip_{idx+1:04d}.flac"
        duration  = get_audio_duration(clip_path)
        actual_start = max(t_start, cursor)
        clips.append((actual_start, clip_path))
        cursor = actual_start + duration

    assembled_tts = job / "tts_assembled.flac"
    print(f"\n  Assembling {total} clips…")
    assemble_timed_audio(clips, str(assembled_tts))
    print(f"  Assembled: {assembled_tts} ({assembled_tts.stat().st_size // 1024} KB)")

    # TTS timestamps are relative to --start, so we need the same video source
    # the original run used: a fragment when --end was given, full video otherwise.
    if args.end:
        start_sec = _parse_mmss(args.start)
        end_sec   = _parse_mmss(args.end)
        video_fragment = job / "video_fragment.mkv"
        if not video_fragment.exists():
            extract_video_fragment(args.video, start_sec, end_sec, str(video_fragment))
        video_src = str(video_fragment)
    else:
        video_src = args.video

    stem = Path(args.video).stem
    final_output = out_dir / f"{stem}_lektor.mkv"
    print(f"\n  Building 2-track video "
          f"(TTS {int(args.tts_volume*100)}% / orig {int(args.orig_volume*100)}%)…")
    merge_voiceover(video_src, str(assembled_tts), str(final_output),
                    orig_volume=args.orig_volume, tts_volume=args.tts_volume)

    print(f"\nDone! Output: {final_output}")


def _check_job_clips(args, out_dir: Path):
    """List clips in this video's job folder that end mid-word.

    Prints them as a --regen argument list, so fixing them is a copy and paste.
    """
    job = job_dir(out_dir, args.video)
    clips = sorted(job.glob("clip_*.flac"))
    if not clips:
        sys.exit(f"No clips found in {job}." + _legacy_hint(out_dir, job))

    print(f"\n  Checking {len(clips)} clips in {job}…")
    bad = []
    for clip in clips:
        if ends_abruptly(clip):
            try:
                number = int(clip.stem.split("_")[1])
            except (IndexError, ValueError):
                continue
            bad.append(number)
            print(f"    {clip.name}  ends at full level "
                  f"({get_audio_duration(clip):.2f}s)")

    if not bad:
        print(f"\n  No truncated clips. All {len(clips)} end cleanly.")
        return
    print(f"\n  {len(bad)} of {len(clips)} clips ({100 * len(bad) / len(clips):.0f}%) "
          f"end mid-word. To regenerate them:")
    print(f"    python3 {Path(__file__).name} \"{args.video}\" --regen "
          f"{' '.join(str(n) for n in bad)}")


def _trim_job_clips(args, out_dir: Path):
    """Trim trailing silence from every clip already on disk for this video.

    Repairs a job generated before trimming existed, or one where the model ran
    to the token ceiling. Only silence is removed, so the speech is untouched
    and nothing needs regenerating — follow with --rebuild.
    """
    job = job_dir(out_dir, args.video)
    clips = sorted(job.glob("clip_*.flac"))
    if not clips:
        sys.exit(f"No clips found in {job}." + _legacy_hint(out_dir, job))

    print(f"\n  Scanning {len(clips)} clips in {job}…")
    total_removed, changed, runaways = 0.0, 0, []
    for i, clip in enumerate(clips, 1):
        before = get_audio_duration(clip)
        tmp = clip.with_suffix(".trim.flac")
        removed = trim_trailing_silence(clip, tmp)
        if removed > 0.02:
            tmp.replace(clip)
            total_removed += removed
            changed += 1
            if removed > 20:
                runaways.append((clip.name, before, removed))
            print(f"    {clip.name}  {before:7.2f}s -> {before - removed:6.2f}s "
                  f"(-{removed:.2f}s)")
        else:
            tmp.unlink(missing_ok=True)

    print(f"\n  Trimmed {changed}/{len(clips)} clips, {total_removed:.0f}s of silence removed.")
    if runaways:
        print(f"  {len(runaways)} clip(s) had run to the token ceiling:")
        for name, before, removed in runaways:
            print(f"    {name}: {before:.0f}s generated for {before - removed:.1f}s of speech")
        print("  Their speech is intact, but listen to them — if any line is cut "
              "short, regenerate it with --regen.")
    print("\n  Now run --rebuild to reassemble with the corrected timing.")


def _regen_clips(args, out_dir: Path):
    """Regenerate specific clips by tts_input block number.
    Only replaces FLAC files — run --rebuild afterwards to reassemble audio and video.
    """
    job     = job_dir(out_dir, args.video)
    tts_srt = job / "tts_input.srt"
    if not tts_srt.exists():
        sys.exit(f"No tts_input.srt found in {job} — run a full generation first."
                 + _legacy_hint(out_dir, job))

    blocks    = load_tts_input_srt(tts_srt)
    total     = len(blocks)
    regen_set = set(args.regen)
    clips_dir = job
    clips_dir.mkdir(parents=True, exist_ok=True)

    bad = [n for n in regen_set if n < 1 or n > total]
    if bad:
        sys.exit(f"Block number(s) out of range (tts_input.srt has {total} blocks): {bad}")

    # Set up voice
    wav_path = Path(args.voice)
    if not wav_path.is_file():
        wav_path = VOICE_PROMPTS_DIR / f"{args.voice}.wav"
    if not wav_path.is_file():
        sys.exit(f"Voice WAV not found: {wav_path}")
    voice_filename, voice_ref_text = upload_voice_wav(wav_path)

    base_seed = args.seed if args.seed != 0 else random.randint(1, 2**31 - 1)
    print(f"\n  Regenerating blocks {sorted(regen_set)} with seed {base_seed}…")

    for idx, (t_start, t_end, text) in enumerate(blocks):
        clip_n = idx + 1
        if clip_n not in regen_set:
            continue
        label = f"clip_{clip_n:04d}"
        print(f"\n  REGEN [{clip_n}/{total}] {text[:70]!r}")
        seed = base_seed + idx
        clip_path, duration, truncated = _generate_one(
            text, voice_filename, voice_ref_text,
            seed, args.speed, clips_dir, label,
            args.strip_leading_ms, args.temperature,
            trailing_fade_ms=args.trailing_fade_ms,
            max_new_tokens=args.max_tokens,
            retries=args.truncation_retries,
        )
        print(f"    → saved as {clip_path.name} ({duration:.2f}s)")

    print(f"\nDone. Run --rebuild to reassemble audio and rebuild the video.")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    global COMFY_URL, COMFY_AUTOSTART, COMFY_START_SCRIPT

    # --settings is resolved before the main parser is built, because the file's
    # contents become that parser's defaults.
    pre = argparse.ArgumentParser(add_help=False)
    pre.add_argument("--settings", default=str(SETTINGS_FILE), metavar="FILE")
    pre_args, _ = pre.parse_known_args()
    settings_path = Path(pre_args.settings)
    cfg = load_settings(settings_path)

    parser = argparse.ArgumentParser(
        description="Generate AI voiceover via ComfyUI + Higgs Audio V3, per-subtitle.",
        epilog=f"Defaults are read from {SETTINGS_FILE.name} next to this script; "
               f"any flag given on the command line overrides it.",
    )
    parser.add_argument("--settings", default=str(SETTINGS_FILE), metavar="FILE",
                        help=f"Settings file to read defaults from "
                             f"(default: {SETTINGS_FILE.name})")
    parser.add_argument("srt",   nargs="?", default=None,
                        help="Input SRT subtitle file (omit with --embedded-subs)")
    parser.add_argument("video", nargs="?", default=None,
                        help="Input video file (MKV, MP4, …)")
    parser.add_argument("--start", default="00:00", metavar="MM:SS",
                        help="Fragment start time (default: 00:00)")
    parser.add_argument("--end",   default=None, metavar="MM:SS",
                        help="Fragment end time")
    parser.add_argument("--voice", default=cfg["voice"],
                        help=f"Voice prompt name or WAV path (default: {cfg['voice']})")
    parser.add_argument("--speed", type=float, default=cfg["speed"], metavar="FACTOR",
                        help=f"Per-clip playback speed via atempo (default: {cfg['speed']})")
    parser.add_argument("--seed", type=int, default=0,
                        help="Base generation seed (default: 0 = random each run)")
    parser.add_argument("--orig-volume", type=float, default=cfg["orig_volume"], metavar="0-1",
                        help=f"Original audio volume in mix (default: {cfg['orig_volume']})")
    parser.add_argument("--tts-volume", type=float, default=cfg["tts_volume"], metavar="0-2",
                        help=f"Voiceover volume in mix (default: {cfg['tts_volume']})")
    parser.add_argument("--output-dir", default=cfg["output_dir"],
                        help=f"Output directory (default: {cfg['output_dir']}/)")
    parser.add_argument("--comfy-url", default=cfg["comfy_url"],
                        help=f"ComfyUI URL (default: {cfg['comfy_url']})")
    parser.add_argument("--skip-tts", metavar="AUDIO_FILE",
                        help="Skip TTS and use this pre-assembled audio file")
    parser.add_argument("--temperature", type=float, default=cfg["temperature"], metavar="T",
                        help=f"Sampling temperature (default: {cfg['temperature']}; "
                             f"lower = more deterministic)")
    parser.add_argument("--strip-leading-ms", type=float, default=cfg["strip_leading_ms"],
                        metavar="MS",
                        help=f"Strip N ms from the start of each TTS clip to remove "
                             f"leading breath/noise artifact "
                             f"(default: {cfg['strip_leading_ms']}, try 80–150)")
    parser.add_argument("--trailing-fade-ms", type=float, default=cfg["trailing_fade_ms"],
                        metavar="MS",
                        help=f"Fade out the last N ms of each TTS clip to suppress "
                             f"tail pops/noise (default: {cfg['trailing_fade_ms']}, "
                             f"set 0 to disable)")
    parser.add_argument("--max-tokens", type=int, default=cfg["max_tokens"], metavar="N",
                        help=f"Max audio tokens the TTS model may generate per clip "
                             f"(default: {cfg['max_tokens']}; raise if long clips are "
                             f"cut off mid-word)")
    parser.add_argument("--resume-from", type=int, default=0, metavar="N",
                        help="Resume a crashed run: skip blocks 1..(N-1) by loading "
                             "their existing clip files, generate from block N onwards.")
    parser.add_argument("--rebuild", action="store_true",
                        help="Re-assemble audio from existing clips and rebuild the "
                             "final video without any TTS generation. "
                             "Use this to change --orig-volume / --tts-volume.")
    parser.add_argument("--regen", nargs="+", type=int, metavar="N",
                        help="Regenerate specific clip(s) by tts_input block number "
                             "(same number as in clips/clip_NNNN.flac and tts_input.srt). "
                             "Only replaces the selected FLAC files; run --rebuild "
                             "afterwards to reassemble audio and rebuild the video.")
    parser.add_argument("--no-comfy-autostart", action="store_true",
                        help="Do not launch ComfyUI automatically when it is not "
                             "answering. By default it is started via the script named "
                             "by [comfy] start_command in settings.ini.")
    parser.add_argument("--truncation-retries", type=int,
                        default=cfg["truncation_retries"], metavar="N",
                        help=f"How many extra attempts to make when the model returns a "
                             f"clip cut off mid-word (default: {cfg['truncation_retries']}; "
                             f"0 disables retrying).")
    parser.add_argument("--check-clips", action="store_true",
                        help="Report which clips in this video's job folder end mid-word, "
                             "as a ready-made --regen list, then stop.")
    parser.add_argument("--trim-silence", action="store_true",
                        help="Trim trailing silence from the clips already generated for "
                             "this video, then stop. Repairs runs where the model failed "
                             "to stop and padded a clip out to --max-tokens, which pushes "
                             "every later clip out of sync. Follow with --rebuild.")
    parser.add_argument("--embedded-sub-index", type=int, default=None, metavar="N",
                        help="With --embedded-subs, pick this ffmpeg stream index instead "
                             "of prompting. Used by ui.py and for scripted runs.")
    parser.add_argument("--embedded-subs", action="store_true",
                        help="Pick a subtitle track embedded in the video file instead of "
                             "providing an external SRT. Lists available tracks and prompts "
                             "for a choice. The video path may be given as the first "
                             "positional argument (no SRT needed).")
    args = parser.parse_args()

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

    if settings_path.is_file():
        print(f"  Settings: {settings_path}")
    else:
        print(f"  Settings: {settings_path.name} not found — using built-in defaults")

    # ── Resolve positional ambiguity for modes that don't need an SRT file ────
    # When the user writes `script.py video.mkv --embedded-subs` (or --rebuild /
    # --regen without an SRT), argparse puts the lone positional into args.srt.
    # Detect this and move it to args.video.
    needs_srt = not (args.embedded_subs or args.rebuild or args.trim_silence
                     or args.check_clips or bool(args.regen))
    if args.video is None and args.srt is not None and not needs_srt:
        args.video = args.srt
        args.srt = None
    if args.video is None:
        parser.error("A video file is required.")
    if needs_srt and args.srt is None:
        parser.error("Provide an SRT subtitle file, or use --embedded-subs to pick one "
                     "from the video.")

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ── Maintenance: report clips the model cut short, then stop ─────────────
    if args.check_clips:
        _check_job_clips(args, out_dir)
        return

    # ── Maintenance: trim silence from existing clips, then stop ─────────────
    if args.trim_silence:
        _trim_job_clips(args, out_dir)
        return

    # ── Rebuild mode: reassemble + remux from existing clips, no TTS ─────────
    if args.rebuild:
        _rebuild_video(args, out_dir)
        return

    # ── Regen mode: replace specific clips, then rebuild audio + video ────────
    if args.regen:
        _regen_clips(args, out_dir)
        return

    # Everything belonging to this video lives together, so successive jobs
    # cannot overwrite each other's clips or tts_input.srt.
    clips_dir = job_dir(out_dir, args.video)
    clips_dir.mkdir(parents=True, exist_ok=True)
    print(f"  Job folder: {clips_dir}")

    start_sec = _parse_mmss(args.start)
    end_sec   = _parse_mmss(args.end) if args.end else None

    # ── Step 1: Parse subtitles ───────────────────────────────────────────────
    print(f"\n[1/4] Parsing subtitles {args.start}–{args.end or 'end'}…")

    if args.embedded_subs:
        sub_list = list_subtitle_streams(args.video)
        if not sub_list:
            sys.exit(f"No subtitle tracks found in {Path(args.video).name}.")
        if args.embedded_sub_index is not None:
            # Non-interactive selection by ffmpeg stream index, for the GUI and
            # for scripted batch runs. Interactive behaviour is unchanged when
            # this is omitted.
            chosen = next((s for s in sub_list
                           if s["index"] == args.embedded_sub_index), None)
            if chosen is None:
                avail = ", ".join(str(s["index"]) for s in sub_list)
                sys.exit(f"No subtitle stream #{args.embedded_sub_index} in "
                         f"{Path(args.video).name}. Available: {avail}")
            print(f"  Using subtitle stream #{chosen['index']} "
                  f"({chosen['codec_name']}{' ' + chosen['language'] if chosen['language'] else ''})")
        else:
            print(f"  Available subtitle tracks in {Path(args.video).name}:")
            for i, s in enumerate(sub_list):
                lang   = f" [{s['language']}]" if s['language'] else ""
                title  = f" — {s['title']}" if s['title'] else ""
                compat = ("" if s['codec_name'].lower() in _TEXT_SUB_CODECS
                          else " (image-based, cannot be used for TTS)")
                print(f"    [{i + 1}] Stream #{s['index']}: {s['codec_name']}{lang}{title}{compat}")
            while True:
                try:
                    choice = int(input("  Select track number: ").strip()) - 1
                    if 0 <= choice < len(sub_list):
                        break
                    print(f"  Enter a number from 1 to {len(sub_list)}.")
                except (ValueError, EOFError, KeyboardInterrupt):
                    sys.exit("\nAborted.")
            chosen = sub_list[choice]
        if chosen['codec_name'].lower() not in _TEXT_SUB_CODECS:
            sys.exit(f"Stream #{chosen['index']} ({chosen['codec_name']}) is image-based "
                     "and cannot be converted to text for TTS.")
        embedded_srt = clips_dir / "embedded_sub.srt"
        extract_subtitle_track(args.video, chosen['index'], str(embedded_srt))
        srt_path = str(embedded_srt)
    else:
        srt_path = args.srt

    if end_sec is None:
        end_sec = 99 * 3600
    blocks = parse_subtitle_blocks(srt_path, start_sec, end_sec)
    if not blocks:
        sys.exit("No subtitle blocks found in the specified range.")
    blocks = split_dialogue_blocks(blocks)
    print(f"      {len(blocks)} entries after dialogue splitting.")

    srt_preview = clips_dir / "tts_input.srt"
    srt_preview.write_text(blocks_to_srt(blocks), encoding="utf-8")
    print(f"      Saved to {srt_preview} for inspection.")

    # ── Step 2: Generate TTS per block ────────────────────────────────────────
    if args.skip_tts:
        assembled_tts = Path(args.skip_tts)
        print(f"\n[2/4] Using existing TTS audio: {assembled_tts}")
    else:
        print(f"\n[2/4] Generating TTS ({len(blocks)} clips, speed {args.speed}×)…")

        wav_path = Path(args.voice)
        if not wav_path.is_file():
            wav_path = 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}")

        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 = upload_voice_wav(wav_path)

        clips = 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)

        assembled_tts = clips_dir / "tts_assembled.flac"
        assemble_timed_audio(clips, str(assembled_tts))
        print(f"  Assembled: {assembled_tts} ({assembled_tts.stat().st_size // 1024} KB)")
        print(f"  Clips saved to {clips_dir}")

    # ── Step 3: Extract video fragment ───────────────────────────────────────
    print(f"\n[3/4] Extracting video fragment…")
    real_end = _parse_mmss(args.end) if args.end else None
    if real_end is not None:
        video_fragment = clips_dir / "video_fragment.mkv"
        extract_video_fragment(args.video, start_sec, real_end, str(video_fragment))
    else:
        video_fragment = Path(args.video)

    # ── Step 4: Mix voiceover with video ─────────────────────────────────────
    stem = Path(args.video).stem
    final_output = out_dir / f"{stem}_lektor.mkv"
    print(f"\n[4/4] Building 2-track video (Original + Polish lektor "
          f"TTS {int(args.tts_volume*100)}% / orig {int(args.orig_volume*100)}%)…")
    merge_voiceover(str(video_fragment), str(assembled_tts), str(final_output),
                    orig_volume=args.orig_volume, tts_volume=args.tts_volume)

    print(f"\nDone! Output: {final_output}")


if __name__ == "__main__":
    main()