"""Dubbing pipeline: extract audio -> transcribe -> translate -> TTS -> mux."""
import os
import subprocess
from pathlib import Path

MAX_ATEMPO = 1.65  # never speed speech up more than this (1.5->1.8 on 2026-10-26 to fix ~1.7s drift, then 1.8->1.65 same night: 1.8 sounded too fast; 1.65 fits every ratio measured on the Oct 25 test video, 1.61-1.67, with ~0.04s total drift)
MAX_VSTRETCH = 1.40  # one scene window may be slowed at most this much so its speech can stay natural (2026-11-07, his idea: stretch only the part that would drift, not the whole video)


def run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError((p.stderr or "ffmpeg failed").strip()[-400:])
    return p


def probe_duration(path):
    p = run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=noprint_wrappers=1:nokey=1", str(path)])
    return float(p.stdout.strip())


def probe_size(path):
    p = run(["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=width,height",
             "-of", "csv=p=0", str(path)])
    w, h = p.stdout.strip().split(",")[:2]
    return int(w), int(h)


def blur_region(src, dst, box):
    """Blur one rectangle of the picture for the whole video.

    Used to hide subtitles burned into the source video: the user marks
    the subtitle area once (box = x, y, w, h in pixels) and that area is
    gaussian-blurred on every frame. Audio is re-encoded so the output
    is a clean mp4 the rest of the pipeline can treat as the source.
    """
    x, y, w, h = box
    fc = (f"[0:v]split[a][b];[b]crop={w}:{h}:{x}:{y},gblur=sigma=28[bb];"
          f"[a][bb]overlay={x}:{y}[v]")
    run(["ffmpeg", "-y", "-v", "error", "-i", str(src),
         "-filter_complex", fc, "-map", "[v]", "-map", "0:a?",
         "-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
         "-c:a", "aac", "-b:a", "160k", str(dst)])


def has_video_stream(path):
    p = run(["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=codec_type", "-of", "csv=p=0", str(path)])
    return "video" in p.stdout


def has_audio_stream(path):
    p = run(["ffprobe", "-v", "error", "-select_streams", "a:0",
             "-show_entries", "stream=codec_type", "-of", "csv=p=0", str(path)])
    return "audio" in p.stdout


def extract_audio(src, dst):
    run(["ffmpeg", "-y", "-v", "error", "-i", str(src),
         "-vn", "-ac", "1", "-ar", "16000", str(dst)])


from gemini_tts import QuotaExhausted
from gemini_tts import configure as _tts_configure
from gemini_tts import synthesize as _gemini_tts
from nara_tts import NARA_VOICES, NARA_VOICE_FOR

# Edge-TTS voice name -> Gemini prebuilt voice (male/female pairs)
GEMINI_VOICE_FOR = {    "my-MM-NilarNeural": "Aoede",    # female Burmese
    "my-MM-ThihaNeural": "Charon",   # male Burmese
    "en-US-AvaNeural": "Aoede",      # female English
    "en-US-AndrewNeural": "Charon",  # male English
}

# When Edge-TTS recovers, the recovery watcher touches this flag file and the
# TTS order flips back to Edge-first automatically (no redeploy needed).
EDGE_PRIMARY_FLAG = Path(__file__).parent / ".edge_primary"


def _edge_first():
    try:
        return EDGE_PRIMARY_FLAG.exists()
    except OSError:
        return False


# Per-job voice engine choice from the UI ('edge' | 'gemini' | None).
# None keeps the standing flag order. app.py sets this at job start;
# the worker queue is serial, so a module-level setting is safe.
_ENGINE = None


# Per-job Nara voice choice from the UI ('shoon' | 'thazin' |
# 'ko-banyar' | 'htet-myat'). None keeps the gender map below
# (female -> shoon, male -> ko-banyar). Same serial-queue rule as _ENGINE.
_NARA_VOICE = None


def set_nara_voice(name):
    global _NARA_VOICE
    _NARA_VOICE = name if name in NARA_VOICES else None


def valid_nara_voice(name):
    """Return a valid Nara voice id, defaulting to shoon."""
    return name if name in NARA_VOICES else "shoon"


def set_engine(name):
    global _ENGINE
    _ENGINE = name if name in ("edge", "gemini", "nara") else None


def configure_tts(keys):
    """Configure Gemini TTS key rotation. Call once at startup."""
    _tts_configure(keys)


_model = None


def transcribe(wav, lang_hint):
    """Return [{start, end, text}]. lang_hint: 'en' or 'my'."""
    global _model
    from faster_whisper import WhisperModel
    if _model is None:
        _model = WhisperModel("base", device="cpu", compute_type="int8")
    segments, _info = _model.transcribe(
        str(wav), language=lang_hint, vad_filter=True,
        vad_parameters=dict(min_silence_duration_ms=400))
    return [{"start": float(s.start), "end": float(s.end), "text": s.text.strip()}
            for s in segments if s.text and s.text.strip()]


def merge_segments(segs, max_gap=0.6, max_len=10.0):
    """Merge close whisper segments into speakable utterances."""
    merged = []
    for s in segs:
        if (merged and s["start"] - merged[-1]["end"] <= max_gap
                and s["end"] - merged[-1]["start"] <= max_len):
            merged[-1]["end"] = s["end"]
            merged[-1]["text"] = merged[-1]["text"] + " " + s["text"]
        else:
            merged.append(dict(s))
    return merged


def _synth_gemini(text, voice, lang, out_path):
    _gemini_tts(text, GEMINI_VOICE_FOR.get(voice, "Aoede"), lang, out_path)


def _synth_edge(text, voice, out_path):
    import asyncio
    import edge_tts
    import edge_tts.communicate as _comm
    import edge_tts.constants as _C

    # 2026-10: Microsoft's classic speech.platform.bing.com endpoint answers
    # 403 in many regions. Use the newer api.msedgeservices.com endpoint for
    # the same voices (token as Ocp-Apim-Subscription-Key, not
    # TrustedClientToken). Patched at call time so library updates can't
    # silently revert it.
    _new_wss = (
        "wss://api.msedgeservices.com/tts/cognitiveservices/websocket/v1"
        f"?Ocp-Apim-Subscription-Key={_C.TRUSTED_CLIENT_TOKEN}"
    )
    _comm.WSS_URL = _new_wss
    _C.WSS_URL = _new_wss

    async def _go():
        await edge_tts.Communicate(text, voice).save(str(out_path))
    asyncio.run(_go())


def _synth_nara(text, voice, out_path):
    from nara_tts import synthesize as _nara_synthesize
    out_path = Path(out_path)
    wav = out_path.with_suffix(".wav")
    voice_id = _NARA_VOICE or NARA_VOICE_FOR.get(voice, "shoon")
    _nara_synthesize(text, voice_id, wav)
    # build_dub_track probes u{i}.mp3 with ffprobe and later runs atempo
    # through ffmpeg, so hand it a real mp3, not a renamed wav.
    run(["ffmpeg", "-y", "-v", "error", "-i", str(wav), str(out_path)])
    wav.unlink(missing_ok=True)


def synthesize(text, voice, out_path):
    """TTS one utterance.

    Order: the per-job engine chosen in the UI goes first (set_engine);
    with no choice, Nara first (the standing default since 2026-11-13),
    then Edge, then Gemini. The other engines are the fallback, then gTTS
    (last resort, robotic single voice). Nara first falls back to Edge,
    then Gemini, so a local-model failure never kills a job outright.
    """
    lang = "my" if voice.startswith("my-MM") else "en"
    err = None
    if _ENGINE == "nara":
        order = ("nara", "edge", "gemini")
    elif _ENGINE in ("edge", "gemini"):
        order = (_ENGINE, "gemini" if _ENGINE == "edge" else "edge")
    else:
        order = ("nara", "edge", "gemini")
    for which in order:
        try:
            if which == "edge":
                _synth_edge(text, voice, out_path)
            elif which == "nara":
                _synth_nara(text, voice, out_path)
            else:
                _synth_gemini(text, voice, lang, out_path)
            return
        except QuotaExhausted:
            # Don't silently fall through on quota - the user should know to wait.
            # Edge-TTS/gTTS can't cover a sustained quota outage anyway.
            raise
        except Exception as e:  # noqa: BLE001
            err = e
            print(f"[tts] {which} failed ({type(e).__name__}), trying next")
    # fallback: gTTS (free, single voice per language — no male/female choice)
    try:
        from gtts import gTTS
        gTTS(text, lang=lang).save(str(out_path))
        print(f"[tts] used gTTS fallback (previous: {type(err).__name__})")
    except Exception:
        raise err


def build_dub_track(utterances, voice, tmpdir, total_dur, progress=None):
    """TTS each utterance, time-fit, mix into one dub track.

    Returns (dub_wav_path, warp, cues) where warp is a list of
    (start, end, factor) video segments for mux(): segments with
    factor > 1 are slowed down in the picture. cues is the placed
    utterances on the OUTPUT timeline ({start, end, text}) — the
    subtitle writers use it so burned-in text follows the dubbed speech
    exactly, stretch and cascade included.

    Finish-the-sentence rule (2026-10-25, user order): an utterance is
    NEVER trimmed; it plays to completion and the next one cascades after
    it, so no sentence is ever cut off mid-word.

    Local-stretch rule (2026-11-07, user order): speech stays at natural
    speed wherever it can. Each utterance gets the window up to the next
    utterance's start (so it can use the natural pause after it). If the
    natural speech is longer than its window, ONLY that window of the
    picture is slowed down — by the ratio it needs, capped at
    MAX_VSTRETCH — instead of rushing the voice. Windows that fit are
    untouched and play at normal speed. If one line is so long that even
    the stretch cap is not enough, that line alone is sped up (atempo,
    capped at MAX_ATEMPO) and, in the extreme, cascades like before.
    """
    tmpdir = Path(tmpdir)
    # 1) synthesize everything at natural speed and measure it
    raw = []
    for i, u in enumerate(utterances):
        mp3 = tmpdir / f"u{i}.mp3"
        synthesize(u["text"], voice, mp3)
        raw.append(probe_duration(mp3))
        if progress:
            progress(i + 1, len(utterances))

    # 2) per-window fit: stretch factor f and speech speed a per utterance
    n = len(utterances)
    warp = []          # (start, end, factor) covering the whole timeline
    starts = []        # utterance start times on the warped timeline
    speeds = []        # atempo per utterance (1.0 = natural)
    t_new = 0.0        # cursor on the warped timeline
    if utterances and utterances[0]["start"] > 0:
        warp.append((0.0, utterances[0]["start"], 1.0))
        t_new = utterances[0]["start"]
    for i, u in enumerate(utterances):
        win_end = utterances[i + 1]["start"] if i + 1 < n else total_dur
        window = max(win_end - u["start"], 0.5)
        need = raw[i] / window
        f = min(need, MAX_VSTRETCH) if need > 1.0 else 1.0
        a = need / f if need > 1.0 else 1.0
        warp.append((u["start"], win_end, f))
        starts.append(t_new)
        speeds.append(min(a, MAX_ATEMPO))
        t_new += window * f
    out_dur = t_new

    # 3) render the placed segments and mix them. Placement cascades: a
    # line never starts before the previous one has finished.
    placed = []
    cues = []
    cursor = 0.0
    for i, start in enumerate(starts):
        start = max(start, cursor)
        fitted = raw[i] / speeds[i]
        cursor = start + fitted
        cues.append({"start": start, "end": cursor,
                     "text": utterances[i]["text"]})
        out_dur = max(out_dur, cursor)
        mp3 = tmpdir / f"u{i}.mp3"
        start_ms = int(round(start * 1000))
        wav_i = tmpdir / f"p{i}.wav"
        af = []
        if speeds[i] > 1.001:
            af.append(f"atempo={speeds[i]:.3f}")
        af.append(f"adelay={start_ms}|{start_ms}")
        run(["ffmpeg", "-y", "-v", "error", "-i", str(mp3),
             "-af", ",".join(af), "-ar", "44100", "-ac", "2", str(wav_i)])
        placed.append(wav_i)
    dub = tmpdir / "dub.wav"
    inputs = []
    for w in placed:
        inputs += ["-i", str(w)]
    run(["ffmpeg", "-y", "-v", "error", *inputs,
         "-filter_complex",
         f"amix=inputs={len(placed)}:dropout_transition=0:normalize=0,"
         f"apad=whole_dur={out_dur:.2f}",
         "-ar", "44100", "-ac", "2", str(dub)])
    return dub, warp, cues


# ---- subtitles (burn-in ASS + downloadable SRT) ----

def _ts_ass(t):
    t = max(0.0, t)
    cs = int(round(t * 100))
    return (f"{cs // 360000}:{(cs // 6000) % 60:02d}:"
            f"{(cs // 100) % 60:02d}.{cs % 100:02d}")


def _ts_srt(t):
    t = max(0.0, t)
    ms = int(round(t * 1000))
    return (f"{ms // 3600000:02d}:{(ms // 60000) % 60:02d}:"
            f"{(ms // 1000) % 60:02d},{ms % 1000:03d}")


def _clean_cues(cues):
    """Sort, drop empties, and keep cues from overlapping each other."""
    out = [c for c in sorted(cues, key=lambda c: c["start"])
           if c["text"] and c["text"].strip()]
    for i, c in enumerate(out):
        c["end"] = max(c["end"], c["start"] + 0.4)
        if i + 1 < len(out) and c["end"] > out[i + 1]["start"]:
            c["end"] = max(out[i + 1]["start"], c["start"] + 0.2)
    return out


def write_ass(cues, path, w, h, bg="box"):
    """Write an ASS file for libass burn-in.

    bg='box' draws a semi-transparent black box behind the text (stays
    readable over any picture); bg='none' is plain text with an outline
    and shadow. Font size and margins scale with the frame height so
    portrait phone videos and landscape clips both look right.
    """
    cues = _clean_cues(cues)
    fs = max(30, round(h * 0.046))
    mv = round(h * 0.028)
    ml = round(w * 0.04)
    if bg == "box":
        style = (f"Style: Dub, Noto Sans Myanmar,{fs},&H00FFFFFF,"
                 f"&H000000FF,&H00000000,&H80141414,0,0,0,0,100,100,0,0,"
                 f"3,{max(6, round(fs * 0.22))},0,2,{ml},{ml},{mv},1")
    else:
        style = (f"Style: Dub, Noto Sans Myanmar,{fs},&H00FFFFFF,"
                 f"&H000000FF,&H00000000,&H7D000000,0,0,0,0,100,100,0,0,"
                 f"1,{max(2, round(fs * 0.09))},1,2,{ml},{ml},{mv},1")
    lines = ["[Script Info]", "ScriptType: v4.00+", f"PlayResX: {w}",
             f"PlayResY: {h}", "WrapStyle: 0", "ScaledBorderAndShadow: yes",
             "", "[V4+ Styles]",
             ("Format: Name, Fontname, Fontsize, PrimaryColour, "
              "SecondaryColour, OutlineColour, BackColour, Bold, Italic, "
              "Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, "
              "BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, "
              "MarginV, Encoding"),
             style, "", "[Events]",
             "Format: Layer, Start, End, Style, Name, MarginL, MarginR, "
             "MarginV, Effect, Text"]
    for c in cues:
        # ASS control chars would break the line; Burmese text rarely has
        # them, but strip rather than risk a corrupt subtitle file
        text = (c["text"].replace("\\", "").replace("{", "").replace("}", "")
                .replace("\r", " ").replace("\n", " "))
        lines.append(f"Dialogue: 0,{_ts_ass(c['start'])},{_ts_ass(c['end'])},"
                     f"Dub,,0,0,0,,{text}")
    Path(path).write_text("\n".join(lines) + "\n", encoding="utf-8")


def write_srt(cues, path):
    """Write the same cues as a plain .srt for download."""
    cues = _clean_cues(cues)
    lines = []
    for i, c in enumerate(cues, 1):
        lines.append(str(i))
        lines.append(f"{_ts_srt(c['start'])} --> {_ts_srt(c['end'])}")
        lines.append(c["text"].replace("\r", " ").strip())
        lines.append("")
    Path(path).write_text("\n".join(lines), encoding="utf-8")


def cleanup_segments(job_dir):
    """Drop the per-utterance TTS intermediates once the mux is done.

    build_dub_track leaves u*.mp3 (raw TTS) and p*.wav (time-fitted
    segments) in the job dir; on a 3-minute video that is ~1GB of files
    that nothing serves afterwards — /r/<jid> only hands out dubbed.mp4.
    Set DUBLAB_KEEP_SEGMENTS=1 to keep them for debugging.
    """
    if os.environ.get("DUBLAB_KEEP_SEGMENTS"):
        return 0
    d = Path(job_dir)
    removed = 0
    for pat in ("u*.mp3", "p*.wav"):
        for f in d.glob(pat):
            try:
                f.unlink()
                removed += 1
            except OSError:
                pass
    return removed


WATERMARK_TEXT = "made with dublab.kbko.qzz.io"
_WM_FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"


def mux(src_video, dub_wav, out_mp4, warp=None, ass_path=None):
    # Burned-in credit watermark, bottom-right (always on): re-encode video
    # with drawtext; scale guard keeps even dims for yuv420p.
    wm = (f"scale=trunc(iw/2)*2:trunc(ih/2)*2,"
          f"drawtext=fontfile={_WM_FONT}:text='{WATERMARK_TEXT}':"
          "fontsize=h/40:fontcolor=white@0.72:"
          "shadowcolor=black@0.55:shadowx=1:shadowy=1:"
          "x=w-text_w-h/40:y=h-text_h-h/56,"
          "format=yuv420p")
    if ass_path:
        # Speech-timed subtitles go on LAST, after the stretch join, so
        # the cue times (written for the output timeline) line up, and
        # above the watermark layer. The path is app-generated under the
        # job dir and contains no filter-special characters.
        wm += f",ass={ass_path}"
    tail = ["-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
            "-map", "0:v:0", "-map", "1:a:0",
            "-c:a", "aac", "-b:a", "128k",
            # No -shortest: the finish-the-sentence rule lets the tail of
            # the last utterance run a little past the video end;
            # -shortest would cut it off again.
            "-movflags", "+faststart", str(out_mp4)]
    stretched = [w for w in (warp or []) if w[2] > 1.001]
    if not stretched:
        run(["ffmpeg", "-y", "-v", "error", "-i", str(src_video),
             "-i", str(dub_wav), "-vf", wm, *tail])
        return
    # Local stretch: cut the picture at the warp boundaries, slow only
    # the segments whose speech needed more time, then join them back.
    segs = warp
    n = len(segs)
    fc = ["[0:v]split=" + str(n) + "".join(f"[s{i}]" for i in range(n))]
    for i, (a, b, f) in enumerate(segs):
        fc.append(f"[s{i}]trim=start={a:.3f}:end={b:.3f},"
                  f"setpts=(PTS-STARTPTS)*{f:.4f},settb=AVTB[v{i}]")
    fc.append("".join(f"[v{i}]" for i in range(n))
              + f"concat=n={n}:v=1:a=0[vcat]")
    fc.append(f"[vcat]{wm}[vout]")
    tail[tail.index("0:v:0")] = "[vout]"
    run(["ffmpeg", "-y", "-v", "error", "-i", str(src_video),
         "-i", str(dub_wav), "-filter_complex", ";".join(fc), *tail])