#!/usr/bin/env python3
"""
lektor_ai — Convert SRT subtitles to AI voiceover using Boson AI TTS (higgs-audio-v3),
then mix the result with a video file's original audio.

Usage:
  python3 lektor_ai.py input/film.srt input/film.mkv --start 04:07 --end 05:45

Requires a free Boson AI API key from https://boson.ai/workspace
Set it via:  export BOSON_API_KEY=bai-xxxx
"""

import argparse
import base64
import os
import re
import subprocess
import sys
from pathlib import Path

# ---------------------------------------------------------------------------
# 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"Unrecognised 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:
    """Accept HH:MM:SS or MM:SS."""
    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 _clean_text(lines: list) -> str:
    """Strip dialogue dashes and return plain text."""
    out = []
    for line in lines:
        line = line.strip()
        line = re.sub(r"^-\s*", "", line)
        if line:
            out.append(line)
    return "\n".join(out)


def parse_subtitle_blocks(srt_path: str, start_sec: float, end_sec: float) -> list:
    """
    Return list of (start_relative, end_relative, plain_text) for each subtitle
    block whose start falls within [start_sec, end_sec].
    Times are relative to start_sec.
    """
    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


# ---------------------------------------------------------------------------
# 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 extract_video_fragment(video_path: str, start_sec: float, end_sec: float,
                           output_path: str):
    duration = end_sec - start_sec
    _run([
        "ffmpeg", "-y",
        "-ss", _ffmpeg_ts(start_sec),
        "-i", video_path,
        "-t", str(duration),
        "-c", "copy",
        output_path,
    ], f"Extracting video {_ffmpeg_ts(start_sec)}–{_ffmpeg_ts(end_sec)}")


def _assemble_batch(clips: list, output_path: str, label: str):
    """Mix a batch of (start_sec, path) clips with absolute adelay into output_path."""
    if len(clips) == 1:
        delay_ms = int(clips[0][0] * 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 = int(start_sec * 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 audio clip at its correct timestamp using adelay, then mix.
    clips: list of (start_sec, audio_path).
    Processes in batches to stay under OS file-descriptor limits.
    """
    if len(clips) <= batch_size:
        _assemble_batch(clips, output_path, f"Assembling {len(clips)} timed clips")
        return

    tmp_dir = Path(output_path).parent / "_batch_tmp"
    tmp_dir.mkdir(exist_ok=True)
    batch_files = []
    try:
        for b, start in enumerate(range(0, len(clips), batch_size)):
            batch = clips[start:start + batch_size]
            batch_path = str(tmp_dir / f"batch_{b:04d}.flac")
            _assemble_batch(
                batch, batch_path,
                f"Assembling batch {b + 1}/{(len(clips) - 1) // batch_size + 1} ({len(batch)} clips)",
            )
            batch_files.append(batch_path)
        # merge batch files — they already carry absolute timing, so delay=0
        batch_clips = [(0.0, Path(p)) for p in 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 merge_voiceover(video_path: str, tts_path: str, output_path: str,
                    orig_volume: float = 0.75):
    _run([
        "ffmpeg", "-y",
        "-i", video_path,
        "-i", tts_path,
        "-filter_complex",
        (
            f"[0:a]volume={orig_volume}[orig];"
            "[1:a]volume=1.0[tts];"
            "[orig][tts]amix=inputs=2:duration=first:normalize=0[out]"
        ),
        "-map", "0:v",
        "-map", "[out]",
        "-c:v", "copy",
        "-c:a", "aac",
        "-b:a", "192k",
        output_path,
    ], "Merging voiceover with video")


# ---------------------------------------------------------------------------
# Boson AI TTS (higgs-audio-v3, OpenAI-compatible API)
# ---------------------------------------------------------------------------

_LOCAL_API_BASE    = "http://localhost:8000/v1"
_BOSON_API_BASE    = "https://api.boson.ai/v1"
_BOSON_MODEL       = "higgs-audio-v3-tts"
_VLLM_MODEL        = "bosonai/higgs-audio-v3-tts-4b"
_VOICE_PROMPTS_DIR = "/home/lukas/higgs-audio/examples/voice_prompts"


def _normalize_tts_text(text: str) -> str:
    text = text.replace("\n", " ").strip()
    text = re.sub(r"\s+", " ", text)
    text = text.replace("(", " ").replace(")", " ")
    if text and text[-1] not in ".!?,;\"'":
        text += "."
    return text.strip()


def _load_voice_prompt(voice: str) -> tuple:
    """Return (ref_audio_b64, ref_text) for the given voice name or wav path."""
    wav_path = Path(voice)
    if not wav_path.exists():
        wav_path = Path(_VOICE_PROMPTS_DIR) / f"{voice}.wav"
    txt_path = wav_path.with_suffix(".txt")

    if not wav_path.exists():
        raise FileNotFoundError(f"Voice prompt WAV not found: {wav_path}")

    ref_audio_b64 = base64.b64encode(wav_path.read_bytes()).decode()
    ref_text = txt_path.read_text(encoding="utf-8").strip() if txt_path.exists() else ""
    return ref_audio_b64, ref_text


def generate_tts_per_block(
    blocks: list,
    download_dir: str,
    voice: str = "gudowski",
    api_key: str = "",
    local: bool = True,
) -> list:
    """
    Generate one audio clip per subtitle block.
    local=True  → calls local vllm-omni server at localhost:8095 (no API key needed)
    local=False → calls Boson AI cloud API (needs BOSON_API_KEY)
    blocks: list of (start_sec, end_sec, plain_text)
    Returns list of (start_sec, audio_Path).
    """
    try:
        from openai import OpenAI
    except ImportError:
        sys.exit("openai package not found. Run: pip install openai")

    if local:
        client = OpenAI(api_key="local", base_url=_LOCAL_API_BASE)
        model = _VLLM_MODEL
        print(f"  Using local vllm-omni server at {_LOCAL_API_BASE}")
    else:
        if not api_key:
            api_key = os.environ.get("BOSON_API_KEY", "")
        if not api_key:
            sys.exit(
                "Boson AI API key required.\n"
                "Get a free key at https://boson.ai/workspace then:\n"
                "  export BOSON_API_KEY=bai-xxxx"
            )
        client = OpenAI(api_key=api_key, base_url=_BOSON_API_BASE)
        model = _BOSON_MODEL
        print(f"  Using Boson AI cloud API")

    ref_audio_b64, ref_text = _load_voice_prompt(voice)
    print(f"  Voice: {voice}  (ref text: {ref_text[:60]!r}…)")

    dl_dir = Path(download_dir).resolve()
    dl_dir.mkdir(parents=True, exist_ok=True)

    results = []
    total = len(blocks)

    for idx, (start_sec, end_sec, text) in enumerate(blocks):
        text = _normalize_tts_text(text)
        print(f"  [{idx+1}/{total}] {_fmt_srt_time(start_sec)}  {text[:70]!r}")

        response = client.audio.speech.create(
            model=model,
            input=text,
            voice="alloy",
            response_format="mp3",
            extra_body={
                "ref_audio": f"data:audio/wav;base64,{ref_audio_b64}",
                "ref_text": ref_text,
            },
        )

        out_path = dl_dir / f"clip_{idx+1:04d}.mp3"
        out_path.write_bytes(response.content)
        results.append((start_sec, out_path))

    return results


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

def main():
    parser = argparse.ArgumentParser(
        description="Generate per-subtitle AI voiceover using Boson AI TTS (higgs-audio-v3)."
    )
    parser.add_argument("srt", help="Input SRT subtitle file")
    parser.add_argument("video", 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",
                        help="Voice prompt name or WAV path (default: gudowski)")
    parser.add_argument("--cloud", action="store_true",
                        help="Use Boson AI cloud API instead of local vllm-omni server")
    parser.add_argument("--api-key", default="",
                        help="Boson AI API key for --cloud mode (or set BOSON_API_KEY)")
    parser.add_argument("--orig-volume", type=float, default=0.75, metavar="0-1",
                        help="Original audio volume in mix (default: 0.75)")
    parser.add_argument("--output-dir", default="output",
                        help="Output directory (default: output/)")
    parser.add_argument("--skip-tts", metavar="AUDIO_FILE",
                        help="Skip TTS and use this pre-assembled audio file")
    args = parser.parse_args()

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

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

    # ── Step 1: Parse subtitle blocks ─────────────────────────────────────────
    print(f"\n[1/4] Parsing subtitles {args.start}–{args.end or 'end'}…")
    if end_sec is None:
        end_sec = 99 * 3600
    blocks = parse_subtitle_blocks(args.srt, start_sec, end_sec)
    if not blocks:
        sys.exit("No subtitle blocks found in the specified range.")
    print(f"      {len(blocks)} subtitle block(s).")

    # ── 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 via Boson AI)…")
        clips = generate_tts_per_block(
            blocks,
            str(out_dir),
            voice=args.voice,
            api_key=args.api_key,
            local=not args.cloud,
        )
        assembled_tts = out_dir / "tts_assembled.wav"
        print(f"\n  Assembling timed audio…")
        assemble_timed_audio(clips, str(assembled_tts))

    # ── 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.mp4"
    print(f"\n[4/4] Mixing (TTS 100%, original at {int(args.orig_volume * 100)}%)…")
    merge_voiceover(str(video_fragment), str(assembled_tts), str(final_output),
                    orig_volume=args.orig_volume)

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


if __name__ == "__main__":
    main()