"""
Karaoke Generator — two-step Gradio app

  Tab 1 · Transcribe
    Audio in → faster-whisper → timestamped JSON (download & edit externally).

  Tab 2 · Generate Video
    Audio + refined JSON + styling → karaoke MP4.

JSON schema (one object per line/segment):
[
  {
    "start": 1.23,
    "end":   3.45,
    "text":  "He Karunanidhi Swami",
    "words": [
      {"word": "He",           "start": 1.23, "end": 1.50},
      {"word": "Karunanidhi", "start": 1.50, "end": 2.20},
      {"word": "Swami",       "start": 2.20, "end": 3.45}
    ]
  },
  ...
]

All models, cache and outputs live inside the project directory.
"""

import json
import os
import glob
import queue
import shutil
import subprocess
import threading
from pathlib import Path

import gradio as gr

# --------------------------------------------------------------------------- #
# Project-local storage
# --------------------------------------------------------------------------- #
PROJECT_DIR = Path(__file__).parent.resolve()

# On HuggingFace Spaces the repo is read-only; use /tmp for all runtime files.
# Locally, keep everything inside the project folder.
ON_HF_SPACES = bool(os.environ.get("SPACE_ID"))
_RUNTIME = Path("/tmp/karaoke") if ON_HF_SPACES else PROJECT_DIR

WHISPER_DIR = _RUNTIME / "models" / "whisper"
MODELS_DIR  = _RUNTIME / "models" / "audio-separator"
HF_MODELS_DIR = _RUNTIME / "models" / "hf"
CACHE_DIR   = _RUNTIME / "cache"
OUTPUTS_DIR = _RUNTIME / "outputs"
WORK_DIR    = _RUNTIME / "work"

for _d in (WHISPER_DIR, MODELS_DIR, HF_MODELS_DIR, CACHE_DIR, OUTPUTS_DIR, WORK_DIR):
    _d.mkdir(parents=True, exist_ok=True)

os.environ["WHISPER_DOWNLOAD_ROOT"] = str(WHISPER_DIR)
os.environ["TRANSFORMERS_CACHE"]    = str(HF_MODELS_DIR)
os.environ["HF_HOME"]               = str(HF_MODELS_DIR)

# --------------------------------------------------------------------------- #
# Constants
# --------------------------------------------------------------------------- #
VIDEO_W, VIDEO_H = 1280, 720

FONT_CHOICES = [
    "Arial", "Arial Black", "Verdana", "Tahoma", "Trebuchet MS",
    "Georgia", "Times New Roman", "Courier New", "Impact", "Comic Sans MS",
]
WHISPER_MODELS  = ["tiny", "base", "small", "medium", "large"]
LANGUAGES       = {
    "Auto-detect": None,
    "English":     "en",
    "Hindi":       "hi",
    "Gujarati":    "gu",
}

# Fine-tuned HuggingFace models — keyed by (language_choice, model_name)
HF_MODELS = {
    ("Gujarati", "small"): {
        "repo":  "vasista22/whisper-gujarati-small",
        "lang":  "gu",
        "label": "Gujarati fine-tuned (vasista22/whisper-gujarati-small)",
    },
}
REMOVAL_METHODS = [
    "Skip (use original audio — instant)",
    "Fast (ffmpeg phase cancel — seconds, stereo only)",
    "AI — MDX-Net (best quality, very slow on CPU)",
]


# --------------------------------------------------------------------------- #
# Color helper
# --------------------------------------------------------------------------- #
def hex_to_ass(color: str) -> str:
    color = (color or "").strip().lstrip("#")
    if len(color) == 3:
        color = "".join(c * 2 for c in color)
    if len(color) != 6:
        return "&H00FFFFFF"
    r, g, b = color[0:2], color[2:4], color[4:6]
    return f"&H00{b}{g}{r}".upper()


# --------------------------------------------------------------------------- #
# JSON ↔ internal segment format
# --------------------------------------------------------------------------- #
def segments_to_json(segments: list[dict]) -> str:
    """Convert internal segments to the user-facing JSON string."""
    out = []
    for seg in segments:
        out.append({
            "start": round(seg["start"], 3),
            "end":   round(seg["end"],   3),
            "text":  " ".join(w[0] for w in seg["words"]).strip(),
            "words": [
                {"word": w[0], "start": round(w[1], 3), "end": round(w[2], 3)}
                for w in seg["words"]
            ],
        })
    return json.dumps(out, ensure_ascii=False, indent=2)


def json_to_segments(json_str: str) -> list[dict]:
    """Parse user-edited JSON back into internal segment format."""
    data = json.loads(json_str)
    segments = []
    for item in data:
        words = [
            [w["word"], float(w["start"]), float(w["end"])]
            for w in item.get("words", [])
        ]
        if not words:
            # plain text line — treat as single word
            words = [[item["text"], float(item["start"]), float(item["end"])]]
        segments.append({
            "start": float(item["start"]),
            "end":   float(item["end"]),
            "words": words,
        })
    return segments


# --------------------------------------------------------------------------- #
# faster-whisper transcription
# --------------------------------------------------------------------------- #
def get_audio_duration(audio_path: str) -> float:
    """Return duration in seconds via ffprobe."""
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", audio_path],
        capture_output=True, text=True,
    )
    try:
        return float(r.stdout.strip())
    except ValueError:
        return 0.0


def run_transcribe_hf(audio_path: str, hf_info: dict,
                      segment_q: queue.Queue) -> None:
    """
    Transcribe using a HuggingFace transformers pipeline (e.g. fine-tuned Whisper).
    Puts the same progress/done/error dicts as run_transcribe into segment_q.
    Words are grouped into line segments by silence gaps (>0.6 s) or max 8 words.
    """
    try:
        import torch
        from transformers import pipeline as hf_pipeline

        repo  = hf_info["repo"]
        lang  = hf_info["lang"]
        cache = str(HF_MODELS_DIR)

        segment_q.put({"type": "progress", "pct": 0.02,
                        "text": f"Loading {repo} (first run downloads ~500 MB)…"})

        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        pipe = hf_pipeline(
            task="automatic-speech-recognition",
            model=repo,
            chunk_length_s=30,
            device=device,
            model_kwargs={"cache_dir": cache},
        )
        pipe.model.config.forced_decoder_ids = (
            pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
        )

        segment_q.put({"type": "progress", "pct": 0.10,
                        "text": "Model loaded. Transcribing…"})

        # Use chunk-level timestamps (return_timestamps=True) — works on all
        # fine-tuned Whisper models without needing alignment_heads in the config.
        # Word timestamps ("word" level) require alignment_heads which many
        # fine-tuned models don't export.
        result = pipe(str(audio_path), return_timestamps=True)
        chunks = result.get("chunks", [])

        segment_q.put({"type": "progress", "pct": 0.80,
                        "text": f"Got {len(chunks)} chunks — distributing word timings…"})

        # Each chunk has a (start, end) timestamp and a text string.
        # We split the text into words and distribute timings evenly within
        # the chunk. The user refines the JSON externally anyway.
        segments: list[dict] = []
        for i, ch in enumerate(chunks):
            text = ch["text"].strip()
            if not text:
                continue
            ts    = ch.get("timestamp") or (0.0, 0.0)
            c_s   = float(ts[0] or 0.0)
            c_e   = float(ts[1] or c_s + 1.0)
            words = text.split()
            if not words:
                continue
            dur   = (c_e - c_s) / len(words)
            word_entries = [
                [w, round(c_s + j * dur, 3), round(c_s + (j + 1) * dur, 3)]
                for j, w in enumerate(words)
            ]
            seg = {"start": c_s, "end": c_e, "words": word_entries}
            segments.append(seg)
            pct = min(0.80 + 0.15 * (i + 1) / max(len(chunks), 1), 0.95)
            segment_q.put({"type": "progress", "pct": pct,
                            "text": f"[{c_s:.1f}s] {text}"})

        segment_q.put({"type": "done", "segments": segments})

    except Exception as e:
        segment_q.put({"type": "error", "msg": str(e)})


def run_transcribe_hf_with_fallback(audio_path: str, hf_info: dict,
                                    fallback_model: str, fallback_lang,
                                    segment_q: queue.Queue) -> None:
    """Try the HF fine-tuned model; on any error fall back to faster-whisper."""
    # Collect everything from the HF attempt into a temp queue first.
    tmp_q: queue.Queue = queue.Queue()
    run_transcribe_hf(audio_path, hf_info, tmp_q)

    # Replay messages, watching for an error.
    hf_failed = False
    while not tmp_q.empty():
        msg = tmp_q.get_nowait()
        if msg["type"] == "error":
            hf_failed = True
            segment_q.put({"type": "progress", "pct": 0.05,
                            "text": f"Fine-tuned model failed ({msg['msg'][:120]}). "
                                    f"Falling back to faster-whisper '{fallback_model}'…"})
            break
        segment_q.put(msg)

    if hf_failed:
        run_transcribe(audio_path, fallback_model, fallback_lang, segment_q)


def run_transcribe(audio_path: str, model_name: str, language_code,
                   segment_q: queue.Queue) -> list[dict]:
    """
    Transcribes audio with faster-whisper.
    Puts progress dicts into segment_q as each segment arrives:
        {"type": "progress", "pct": float, "text": str}
        {"type": "done", "segments": list}
        {"type": "error", "msg": str}
    """
    try:
        from faster_whisper import WhisperModel

        segment_q.put({"type": "progress", "pct": 0.02,
                        "text": f"Loading faster-whisper '{model_name}' (int8, CPU)…"})
        model = WhisperModel(
            model_name,
            device="cpu",
            compute_type="int8",
            download_root=str(WHISPER_DIR),
        )

        duration = get_audio_duration(audio_path)
        lang_label = language_code or "auto"
        segment_q.put({"type": "progress", "pct": 0.05,
                        "text": f"Model loaded. Transcribing (language={lang_label})…"})

        segments_iter, info = model.transcribe(
            audio_path,
            language=language_code,
            word_timestamps=True,
            vad_filter=True,
            vad_parameters={"min_silence_duration_ms": 500},
        )
        detected = info.language if not language_code else language_code
        segment_q.put({"type": "progress", "pct": 0.08,
                        "text": f"Language: {detected}. Processing audio…"})

        segments = []
        for seg in segments_iter:          # ← work happens here, one segment at a time
            words = []
            if seg.words:
                for w in seg.words:
                    words.append([w.word.strip(), float(w.start), float(w.end)])
            else:
                words.append([seg.text.strip(), float(seg.start), float(seg.end)])

            if any(w[0] for w in words):
                segments.append({
                    "start": float(seg.start),
                    "end":   float(seg.end),
                    "words": words,
                })
                # stream each segment as it arrives
                pct = min(0.08 + 0.90 * (seg.end / duration), 0.97) if duration else 0.5
                line_text = " ".join(w[0] for w in words).strip()
                segment_q.put({"type": "progress", "pct": pct,
                                "text": f"[{seg.start:.1f}s] {line_text}"})

        segment_q.put({"type": "done", "segments": segments})

    except Exception as e:
        segment_q.put({"type": "error", "msg": str(e)})


# --------------------------------------------------------------------------- #
# Vocal removal
# --------------------------------------------------------------------------- #
def remove_vocals_fast(src: str, dst: str):
    r = subprocess.run(
        ["ffmpeg", "-y", "-i", src, "-af", "pan=stereo|c0=c0-c1|c1=c1-c0", dst],
        capture_output=True, text=True,
    )
    if r.returncode != 0 or not Path(dst).exists():
        raise RuntimeError(f"ffmpeg vocal removal failed:\n{r.stderr[-600:]}")


def remove_vocals_ai(src: str, out_dir: str, log_fn) -> str:
    from audio_separator.separator import Separator
    import logging
    log_fn("Loading MDX-Net model…")
    sep = Separator(
        log_level=logging.WARNING,
        model_file_dir=str(MODELS_DIR),
        output_dir=out_dir,
        output_format="wav",
    )
    sep.load_model(model_filename="UVR_MDXNET_KARA_2.onnx")
    log_fn("Separating (CPU — slow)…")
    stems = sep.separate(src)
    if not stems or len(stems) < 2:
        raise RuntimeError("Separation failed — song may be too long for available RAM.")
    return stems[0]  # instrumental


# --------------------------------------------------------------------------- #
# ASS subtitle builder
# --------------------------------------------------------------------------- #
def _ass_time(sec: float) -> str:
    sec = max(0.0, sec)
    h, rem = divmod(int(sec), 3600)
    m, s   = divmod(rem, 60)
    cs     = int(round((sec - int(sec)) * 100)) % 100
    return f"{h}:{m:02d}:{s:02d}.{cs:02d}"


def build_ass(segments, font_name, font_size,
              primary_hex, highlight_hex, outline_hex, outline_w):
    header = (
        f"[Script Info]\nScriptType: v4.00+\n"
        f"PlayResX: {VIDEO_W}\nPlayResY: {VIDEO_H}\n"
        "WrapStyle: 2\nScaledBorderAndShadow: yes\n\n"
        "[V4+ Styles]\n"
        "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, "
        "OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, "
        "ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, "
        "Alignment, MarginL, MarginR, MarginV, Encoding\n"
        f"Style: Default,{font_name},{font_size},"
        f"{hex_to_ass(highlight_hex)},{hex_to_ass(primary_hex)},"
        f"{hex_to_ass(outline_hex)},&H64000000,"
        f"-1,0,0,0,100,100,0,0,1,{outline_w},1,2,60,60,80,1\n\n"
        "[Events]\n"
        "Format: Layer, Start, End, Style, MarginL, MarginR, MarginV, Effect, Text\n"
    )
    events = []
    for seg in segments:
        body = "".join(
            f"{{\\kf{max(1, int(round((we - ws) * 100)))}}}"
            f"{txt.replace('{','(').replace('}',')')} "
            for txt, ws, we in seg["words"]
        ).strip()
        events.append(
            f"Dialogue: 0,{_ass_time(seg['start'])},{_ass_time(seg['end'])},"
            f"Default,,0,0,0,,{body}"
        )
    return header + "\n".join(events) + "\n"


# --------------------------------------------------------------------------- #
# Tab 1 — Transcribe
# --------------------------------------------------------------------------- #
def do_transcribe(audio_file, whisper_model, language_choice, progress=gr.Progress()):
    if audio_file is None:
        raise gr.Error("Please upload an audio file.")

    lang_code = LANGUAGES.get(language_choice)
    log_lines: list[str] = []

    hf_info = HF_MODELS.get((language_choice, whisper_model))
    seg_q: queue.Queue = queue.Queue()

    if hf_info:
        log_lines.append(f"Trying fine-tuned model: {hf_info['label']}")
        t = threading.Thread(
            target=run_transcribe_hf_with_fallback,
            args=(str(audio_file), hf_info, whisper_model, lang_code, seg_q),
            daemon=True,
        )
    else:
        t = threading.Thread(
            target=run_transcribe,
            args=(str(audio_file), whisper_model, lang_code, seg_q),
            daemon=True,
        )
    t.start()

    segments = None
    while True:
        try:
            msg = seg_q.get(timeout=0.5)
        except queue.Empty:
            if not t.is_alive():
                break
            yield None, None, "\n".join(log_lines)
            continue

        if msg["type"] == "progress":
            log_lines.append(msg["text"])
            progress(msg["pct"], desc=msg["text"][:60])
            yield None, None, "\n".join(log_lines)

        elif msg["type"] == "done":
            segments = msg["segments"]
            break

        elif msg["type"] == "error":
            raise gr.Error(f"Transcription failed: {msg['msg']}")

    t.join()
    if segments is None:
        raise gr.Error("Transcription ended without results.")

    json_str  = segments_to_json(segments)
    json_path = OUTPUTS_DIR / "lyrics.json"
    json_path.write_text(json_str, encoding="utf-8")

    log_lines.append(f"✓ Done — {len(segments)} segments saved to lyrics.json")
    progress(1.0, desc="Done!")
    yield json_str, str(json_path), "\n".join(log_lines)


# --------------------------------------------------------------------------- #
# Tab 2 — Generate Video
# --------------------------------------------------------------------------- #
def do_generate(
    audio_file,
    json_input,       # either uploaded file path OR pasted text
    background_image,
    removal_method,
    font_file,
    font_name,
    font_size,
    primary_color,
    highlight_color,
    outline_color,
    outline_width,
    progress=gr.Progress(),
):
    if audio_file is None:
        raise gr.Error("Please upload an audio file.")
    if not json_input or not json_input.strip():
        raise gr.Error("Please paste your lyrics JSON (from Tab 1, after editing).")
    if shutil.which("ffmpeg") is None:
        raise gr.Error("ffmpeg not found on PATH.")

    # Parse JSON
    try:
        segments = json_to_segments(json_input)
    except Exception as e:
        raise gr.Error(f"Invalid JSON: {e}")

    if not segments:
        raise gr.Error("JSON parsed but contains no segments.")

    job_dir = WORK_DIR / f"job_{os.getpid()}"
    job_dir.mkdir(parents=True, exist_ok=True)

    log_lines = []
    def log(msg): log_lines.append(msg)
    def status(): return "\n".join(log_lines)

    method_key = removal_method.split("(")[0].strip().lower()

    # ---- vocal removal ---- #
    if method_key == "skip":
        backing = str(audio_file)
        log("Vocal removal: skipped.")
        progress(0.10, desc="Using original audio…")
        yield None, status()

    elif method_key == "fast":
        log("Vocal removal: ffmpeg phase cancel…")
        progress(0.05, desc="Removing vocals…")
        yield None, status()
        out_wav = str(job_dir / "instrumental.wav")
        try:
            remove_vocals_fast(str(audio_file), out_wav)
        except RuntimeError as e:
            raise gr.Error(str(e))
        backing = out_wav
        log("  Done.")
        progress(0.20, desc="Vocal removal done.")
        yield None, status()

    else:  # AI
        log("Vocal removal: MDX-Net (CPU — slow)…")
        progress(0.03, desc="AI vocal separation…")
        yield None, status()
        q2: queue.Queue = queue.Queue()
        res2, err2 = [], []
        def _run_ai():
            try:
                res2.append(remove_vocals_ai(str(audio_file), str(job_dir),
                                             lambda m: q2.put(m)))
            except Exception as e:
                err2.append(str(e))
        t2 = threading.Thread(target=_run_ai, daemon=True)
        t2.start()
        pct = 0.03
        while t2.is_alive():
            try:
                log(f"  {q2.get(timeout=2)}")
            except queue.Empty:
                pass
            pct = min(pct + 0.004, 0.45)
            progress(pct, desc="AI separation…")
            yield None, status()
        t2.join()
        if err2:
            raise gr.Error(f"AI separation failed: {err2[0]}")
        backing = res2[0]
        log(f"  Instrumental: {Path(backing).name}")
        progress(0.50, desc="Separation done.")
        yield None, status()

    # ---- render ---- #
    log("Rendering karaoke video…")
    progress(0.80, desc="Rendering with ffmpeg…")
    yield None, status()

    fonts_dir, effective_font = None, font_name
    if font_file is not None:
        fonts_dir = job_dir / "fonts"
        fonts_dir.mkdir(exist_ok=True)
        shutil.copy(str(font_file), fonts_dir)
        effective_font = Path(font_file).stem

    ass_path = job_dir / "lyrics.ass"
    ass_path.write_text(
        build_ass(segments, effective_font, int(font_size),
                  primary_color, highlight_color, outline_color, float(outline_width)),
        encoding="utf-8",
    )

    bg_input = (["-loop", "1", "-i", str(background_image)]
                if background_image else
                ["-f", "lavfi", "-i", f"color=c=black:s={VIDEO_W}x{VIDEO_H}"])
    scale = (f"scale={VIDEO_W}:{VIDEO_H}:force_original_aspect_ratio=increase,"
             f"crop={VIDEO_W}:{VIDEO_H},setsar=1"
             if background_image else "setsar=1")

    ass_esc = str(ass_path).replace("\\", "/").replace(":", "\\:")
    vf = f"{scale},ass='{ass_esc}'"
    if fonts_dir:
        fd_esc = str(fonts_dir).replace("\\", "/").replace(":", "\\:")
        vf = f"{scale},ass='{ass_esc}':fontsdir='{fd_esc}'"

    out_video = OUTPUTS_DIR / f"karaoke_{os.getpid()}.mp4"
    r = subprocess.run([
        "ffmpeg", "-y", *bg_input, "-i", backing,
        "-vf", vf,
        "-c:v", "libx264", "-tune", "stillimage", "-pix_fmt", "yuv420p",
        "-c:a", "aac", "-b:a", "192k", "-shortest",
        str(out_video),
    ], capture_output=True, text=True)

    if r.returncode != 0 or not out_video.exists():
        raise gr.Error(f"ffmpeg failed:\n{r.stderr[-2000:]}")

    log(f"Saved: {out_video.name}")
    progress(1.0, desc="Done!")
    yield str(out_video), status()


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
def build_ui():
    with gr.Blocks(title="Karaoke Generator") as demo:
        gr.Markdown(
            "# Karaoke Generator\n"
            f"> Files stored in: `{_RUNTIME}`"
        )

        with gr.Tabs():
            # ================================================================ #
            # TAB 1 — Transcribe
            # ================================================================ #
            with gr.Tab("Step 1 — Transcribe"):
                gr.Markdown(
                    "Upload your audio, run Whisper to get timestamped lyrics, "
                    "then **download the JSON and refine it** (fix wrong words, "
                    "split/merge lines, correct timings) — paste it into Gemini or "
                    "edit manually. Then go to **Step 2** to generate the video."
                )
                with gr.Row():
                    with gr.Column():
                        t1_audio   = gr.Audio(label="Audio file", type="filepath")
                        with gr.Row():
                            t1_model = gr.Dropdown(WHISPER_MODELS, value="small",
                                                   label="Whisper model")
                            t1_lang  = gr.Dropdown(list(LANGUAGES.keys()),
                                                   value="Auto-detect",
                                                   label="Language")
                        t1_hf_notice = gr.Markdown(visible=False)
                        gr.Markdown(
                            "- **tiny** ~15 s · **small** ~1 min · **medium** ~3 min (all on CPU)\n"
                            "- Pinning the language (Hindi/Gujarati) improves accuracy\n"
                            "- Fine-tuned models auto-selected when available (e.g. Gujarati + small)"
                        )
                        t1_run = gr.Button("Transcribe", variant="primary")

                    with gr.Column():
                        t1_status   = gr.Textbox(label="Status", lines=5, interactive=False)
                        t1_download = gr.File(label="Download lyrics.json")
                        t1_json     = gr.Textbox(
                            label="Lyrics JSON — review / copy to edit externally",
                            lines=25,
                            interactive=False,
                        )

                def update_hf_notice(lang, model):
                    info = HF_MODELS.get((lang, model))
                    if info:
                        return gr.Markdown(
                            value=f"> ✨ **Fine-tuned model will be used:** `{info['repo']}`  \n"
                                  f"> Trained specifically on Gujarati speech — better accuracy than generic Whisper.",
                            visible=True,
                        )
                    return gr.Markdown(visible=False)

                t1_lang.change(update_hf_notice, inputs=[t1_lang, t1_model],
                               outputs=[t1_hf_notice])
                t1_model.change(update_hf_notice, inputs=[t1_lang, t1_model],
                                outputs=[t1_hf_notice])

                t1_run.click(
                    do_transcribe,
                    inputs=[t1_audio, t1_model, t1_lang],
                    outputs=[t1_json, t1_download, t1_status],
                )

            # ================================================================ #
            # TAB 2 — Generate Video
            # ================================================================ #
            with gr.Tab("Step 2 — Generate Video"):
                gr.Markdown(
                    "Paste your **edited JSON** below (or upload the file), "
                    "choose vocal removal + styling, and render the karaoke video."
                )
                with gr.Row():
                    # ---- left: inputs ---- #
                    with gr.Column():
                        t2_audio  = gr.Audio(label="Audio file", type="filepath")
                        t2_json   = gr.Textbox(
                            label="Lyrics JSON (paste refined JSON here)",
                            lines=14,
                            interactive=True,
                            placeholder='[{"start": 0.0, "end": 3.0, "text": "...", "words": [...]}]',
                        )
                        t2_bg     = gr.Image(label="Background image (optional)",
                                             type="filepath")

                        gr.Markdown("### Vocal Removal")
                        t2_removal = gr.Radio(
                            REMOVAL_METHODS, value=REMOVAL_METHODS[1],
                            label="Method",
                        )

                        gr.Markdown("### Font & Style")
                        with gr.Row():
                            t2_font_name = gr.Dropdown(FONT_CHOICES, value="Arial Black",
                                                       label="Font family")
                            t2_font_size = gr.Slider(24, 120, value=56, step=2,
                                                     label="Font size")
                        t2_font_file = gr.File(label="Custom font (.ttf/.otf)",
                                               file_types=[".ttf", ".otf"])
                        t2_outline   = gr.Slider(0, 8, value=2.5, step=0.5,
                                                 label="Outline thickness")
                        with gr.Row():
                            t2_primary   = gr.ColorPicker(value="#FFFFFF",
                                                          label="Upcoming text")
                            t2_highlight = gr.ColorPicker(value="#FFD400",
                                                          label="Sung / highlight")
                            t2_outline_c = gr.ColorPicker(value="#000000",
                                                          label="Outline color")

                        t2_run = gr.Button("Generate Karaoke Video", variant="primary")

                    # ---- right: outputs ---- #
                    with gr.Column():
                        t2_video  = gr.Video(label="Karaoke video")
                        t2_status = gr.Textbox(label="Status", lines=10,
                                               interactive=False)

                t2_run.click(
                    do_generate,
                    inputs=[
                        t2_audio, t2_json, t2_bg, t2_removal,
                        t2_font_file, t2_font_name, t2_font_size,
                        t2_primary, t2_highlight, t2_outline_c, t2_outline,
                    ],
                    outputs=[t2_video, t2_status],
                )

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.queue()   # required for streaming generators
    demo.launch(
        server_name="0.0.0.0",   # needed for HF Spaces / Docker
        server_port=7860,
    )