#!/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 json
import random
import re
import shutil
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path
from urllib.request import urlopen, Request

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

COMFY_URL         = "http://127.0.0.1:8189"
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("/home/lukas/Programy/lektor_ai/gudowski")

# 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")


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"]


def wait_for_completion(prompt_id: str, poll_interval: int = 3) -> dict:
    print(f"    Waiting for {prompt_id[:8]}…", end="", flush=True)
    while True:
        try:
            history = _comfy_get(f"/history/{prompt_id}")
        except Exception as e:
            print(f"\n  Poll error: {e}", file=sys.stderr)
            time.sleep(poll_interval)
            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(" done.")
                return entry

        print(".", end="", flush=True)
        time.sleep(poll_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) -> tuple:
    """Submit one TTS job and return (path, duration_after_speedup)."""
    workflow = build_text_workflow(text, voice_filename, voice_ref_text, seed,
                                   temperature, max_new_tokens)
    prompt_id = submit_workflow(workflow)
    history_entry = wait_for_completion(prompt_id)
    raw_clip = extract_output_audio(history_entry)

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

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

    return fast_clip, get_audio_duration(fast_clip)


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

    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 = _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,
        )

        # 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}")

    return clips


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

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.
    """
    tts_srt = out_dir / "tts_input.srt"
    if not tts_srt.exists():
        sys.exit(f"No tts_input.srt found in {out_dir} — run a full generation first.")

    blocks    = load_tts_input_srt(tts_srt)
    total     = len(blocks)
    clips_dir = out_dir / "clips"

    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 = out_dir / "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 = out_dir / "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 _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.
    """
    tts_srt = out_dir / "tts_input.srt"
    if not tts_srt.exists():
        sys.exit(f"No tts_input.srt found in {out_dir} — run a full generation first.")

    blocks    = load_tts_input_srt(tts_srt)
    total     = len(blocks)
    regen_set = set(args.regen)
    clips_dir = out_dir / "clips"
    clips_dir.mkdir(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 = _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,
        )
        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
    parser = argparse.ArgumentParser(
        description="Generate AI voiceover via ComfyUI + Higgs Audio V3, per-subtitle."
    )
    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="gudowski_short",
                        help="Voice prompt name or WAV path (default: gudowski_short)")
    parser.add_argument("--speed", type=float, default=1.2, metavar="FACTOR",
                        help="Per-clip playback speed via atempo (default: 1.2)")
    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=0.60, metavar="0-1",
                        help="Original audio volume in mix (default: 0.60)")
    parser.add_argument("--tts-volume", type=float, default=1.8, metavar="0-2",
                        help="Voiceover volume in mix (default: 1.8)")
    parser.add_argument("--output-dir", default="output",
                        help="Output directory (default: output/)")
    parser.add_argument("--comfy-url", default=COMFY_URL,
                        help=f"ComfyUI URL (default: {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=0.9, metavar="T",
                        help="Sampling temperature (default: 0.9; lower = more deterministic)")
    parser.add_argument("--strip-leading-ms", type=float, default=0.0, metavar="MS",
                        help="Strip N ms from the start of each TTS clip to remove "
                             "leading breath/noise artifact (default: 0, try 80–150)")
    parser.add_argument("--trailing-fade-ms", type=float, default=80.0, metavar="MS",
                        help="Fade out the last N ms of each TTS clip to suppress "
                             "tail pops/noise (default: 80, set 0 to disable)")
    parser.add_argument("--max-tokens", type=int, default=4096, metavar="N",
                        help="Max audio tokens the TTS model may generate per clip "
                             "(default: 4096; raise if long clips are 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("--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

    # ── 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 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)
    clips_dir = out_dir / "clips"
    clips_dir.mkdir(exist_ok=True)

    # ── 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

    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}.")
        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 = out_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 = out_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)

        assembled_tts = out_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 = out_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()