"""
Mazinger Dubber — Main Gradio UI
Cinematic film-editing control panel aesthetic.
DaVinci Resolve meets a broadcast dubbing suite.
"""

from __future__ import annotations

import re
from typing import Generator

import gradio as gr

from theme import MazingerTheme, CUSTOM_CSS
from pipeline_runner import run_pipeline, STAGE_NAMES

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VOICE_THEMES = [
    "narrator-m",
    "narrator-f",
    "young-m",
    "young-f",
    "deep-m",
    "deep-f",
    "warm-m",
    "warm-f",
    "news-m",
    "news-f",
    "storyteller-m",
    "storyteller-f",
    "kid-m",
    "kid-f",
    "teen-m",
    "teen-f",
]

VOICE_PROFILES = [
    "abubakr",
    "daheeh-v1",
    "3b1b",
    "italian-v1",
    "morgan-freeman",
    "trump-v1",
]

# Languages supported by voice themes (TTS generation)
THEME_LANGUAGES = [
    "Chinese", "English", "French", "German", "Italian",
    "Japanese", "Korean", "Portuguese", "Russian", "Spanish",
]

# All languages supported by translation (34 from mazinger repo)
LANGUAGES = [
    "Arabic", "Bengali", "Chinese (Simplified)", "Chinese (Traditional)",
    "Czech", "Danish", "Dutch", "English", "Finnish", "French",
    "German", "Greek", "Hebrew", "Hindi", "Hungarian", "Indonesian",
    "Italian", "Japanese", "Korean", "Malay", "Norwegian", "Persian",
    "Polish", "Portuguese", "Romanian", "Russian", "Spanish",
    "Swedish", "Thai", "Turkish", "Ukrainian", "Urdu", "Vietnamese",
]

# 3-letter abbreviations matching STAGE_NAMES order (10 stages)
_STAGE_ABBREVS = ["DWN", "STT", "THB", "DSC", "REV", "TRN", "SEG", "TTS", "ASM", "SUB"]

# ---------------------------------------------------------------------------
# Pipeline HTML builder
# ---------------------------------------------------------------------------

def build_pipeline_html(active: int = -1, done_up_to: int = -1) -> str:
    """
    Render a horizontal hex-node pipeline visualizer.
    active     — index of the currently running stage (amber pulse)
    done_up_to — stages with index < done_up_to are marked done (teal)
    """
    nodes_html = []
    stage_count = len(STAGE_NAMES)

    for i, name in enumerate(STAGE_NAMES):
        abbrev = _STAGE_ABBREVS[i] if i < len(_STAGE_ABBREVS) else name[:3].upper()

        # Determine node state
        if i < done_up_to:
            node_cls = "pipeline-node done"
        elif i == active:
            node_cls = "pipeline-node active"
        else:
            node_cls = "pipeline-node"

        node_html = f"""
        <div style="display:flex;flex-direction:column;align-items:center;gap:0.35rem;">
          <div class="{node_cls}" title="{name}">
            <span style="
              display:flex;align-items:center;justify-content:center;
              width:100%;height:100%;
              font-family:'JetBrains Mono',monospace;
              font-size:0.55rem;font-weight:700;
              color:{'#0D0F12' if (i < done_up_to or i == active) else '#5A5A6A'};
              letter-spacing:0.04em;
            ">{abbrev}</span>
          </div>
          <div class="pipeline-label">{name}</div>
        </div>"""
        nodes_html.append(node_html)

        # Connector after every node except the last
        if i < stage_count - 1:
            connector_cls = "pipeline-connector done" if i < done_up_to else "pipeline-connector"
            nodes_html.append(
                f'<div class="{connector_cls}" style="margin-bottom:1.1rem;"></div>'
            )

    inner = "\n".join(nodes_html)
    return f'<div class="pipeline-container">{inner}</div>'


# ---------------------------------------------------------------------------
# YouTube thumbnail helper
# ---------------------------------------------------------------------------

_YT_RE = re.compile(
    r"(?:youtube\.com/(?:watch\?v=|shorts/)|youtu\.be/)([A-Za-z0-9_-]{11})"
)


def fetch_youtube_thumbnail(url: str) -> str:
    """Return HTML img tag for YouTube thumbnail (client-side, no server download)."""
    if not url:
        return ""
    match = _YT_RE.search(url)
    if match:
        vid_id = match.group(1)
        return (
            f'<img src="https://img.youtube.com/vi/{vid_id}/hqdefault.jpg" '
            f'style="width:100%;max-height:200px;object-fit:cover;border-radius:8px;" />'
        )
    return ""


# ---------------------------------------------------------------------------
# Voice mode toggle helpers
# ---------------------------------------------------------------------------

def _voice_mode_change(mode: str) -> tuple[dict, dict, dict, dict]:
    """Show/hide voice control components based on selected mode."""
    show_theme   = mode == "theme"
    show_profile = mode == "profile"
    show_clone   = mode == "clone"
    return (
        gr.update(visible=show_theme),    # voice_theme dropdown
        gr.update(visible=show_profile),  # voice_profile dropdown
        gr.update(visible=show_clone),    # voice_sample audio
        gr.update(visible=show_clone),    # voice_script textbox
    )


# ---------------------------------------------------------------------------
# Main pipeline processor (generator)
# ---------------------------------------------------------------------------

def process_video(
    source_url: str,
    source_upload,
    target_lang: str,
    source_lang: str,
    voice_mode: str,
    voice_theme: str,
    voice_profile: str,
    voice_sample,
    voice_script: str,
    output_type: str,
    embed_subs: bool,
    sub_font: str,
    sub_size: int,
    slice_start: str,
    slice_end: str,
    sub_position: str,
    sub_color: str,
    sub_bg_alpha: float,
    sub_outline_width: int,
    sub_bold: bool,
    sub_line_spacing: int,
    sub_source: str,
    asr_review: bool,
    tempo_mode: str,
    max_tempo: float,
    words_per_second: float,
    duration_budget: float,
    translate_technical: bool,
) -> Generator:
    """
    Drive run_pipeline() and yield 6 UI updates per stage event:
      [pipeline_html, log_text, video_output, audio_output, srt_preview, srt_download]
    """
    # Resolve source: uploaded file takes priority over URL
    source = source_url or ""
    if source_upload is not None:
        source = source_upload if isinstance(source_upload, str) else source_upload.name

    if not source:
        yield (
            build_pipeline_html(),
            "[ERROR] Please provide a video URL or upload a file.",
            None, None, "", None,
        )
        return

    log_lines: list[str] = []
    video_path = None
    # Validate theme + language compatibility
    if voice_mode == "theme" and target_lang not in THEME_LANGUAGES:
        yield (
            build_pipeline_html(),
            f"[ERROR] Voice themes don't support {target_lang}. Use auto-clone, a voice profile, or custom clone instead.",
            None, None, "", None,
        )
        return

    audio_path = None
    srt_content = ""
    srt_path = None

    # Convert slider defaults to None when user hasn't changed them
    wps_val = words_per_second if words_per_second != 0 else None
    db_val = duration_budget if duration_budget != 0 else None

    for stage_idx, log_msg, result in run_pipeline(
        source=source,
        target_language=target_lang,
        source_language=source_lang,
        voice_mode=voice_mode,
        voice_theme=voice_theme,
        voice_profile=voice_profile,
        voice_sample=voice_sample,
        voice_script=voice_script,
        output_type=output_type,
        embed_subtitles=embed_subs,
        subtitle_font=sub_font,
        subtitle_font_size=sub_size,
        slice_start=slice_start,
        slice_end=slice_end,
        subtitle_position=sub_position,
        subtitle_color=sub_color,
        subtitle_bg_alpha=sub_bg_alpha,
        subtitle_outline_width=sub_outline_width,
        subtitle_bold=sub_bold,
        subtitle_line_spacing=sub_line_spacing,
        subtitle_source=sub_source,
        asr_review=asr_review,
        tempo_mode=tempo_mode,
        max_tempo=max_tempo,
        words_per_second=wps_val,
        duration_budget=db_val,
        translate_technical_terms=translate_technical,
    ):
        log_lines.append(log_msg)

        # Collect final results
        if result.get("final_path"):
            fpath = result["final_path"]
            if fpath.endswith((".mp4", ".mkv", ".avi", ".mov", ".webm")):
                video_path = fpath
            else:
                audio_path = fpath
        if result.get("srt_path"):
            srt_path = result["srt_path"]
            try:
                srt_content = open(srt_path).read()
            except Exception:
                pass

        is_done = result.get("done", False)
        done_up_to = stage_idx - 1 if not is_done else len(STAGE_NAMES)
        active = stage_idx - 1 if not is_done and stage_idx <= len(STAGE_NAMES) else -1

        yield (
            build_pipeline_html(active=active, done_up_to=done_up_to),
            "\n".join(log_lines),
            video_path,
            audio_path,
            srt_content,
            srt_path,
        )


# ---------------------------------------------------------------------------
# Gradio Blocks layout
# ---------------------------------------------------------------------------

with gr.Blocks(
    theme=MazingerTheme(),
    css=CUSTOM_CSS,
    title="Mazinger Dubber",
) as demo:

    # ── Header ──────────────────────────────────────────────────────────────
    gr.HTML("""
    <div class="studio-header">
      <img src="https://raw.githubusercontent.com/bakrianoo/mazinger/refs/heads/master/docs/assets/main-logo-refined.png" alt="Mazinger" style="height:80px;margin-bottom:0.5rem;" />
      <p>End-to-end video dubbing studio</p>
    </div>
    """)

    # ── Row: Source + Voice ──────────────────────────────────────────────────
    with gr.Row(equal_height=False):

        # Source column
        with gr.Column(scale=1, elem_classes=["card-surface"]):
            gr.Markdown("### SOURCE")
            source_url = gr.Textbox(
                label="Video URL",
                placeholder="YouTube URL or direct video/audio link...",
                lines=1,
                max_lines=2,
            )
            gr.Markdown("<center style='color:var(--text-secondary);font-size:0.85rem'>YouTube links download via proxy — or upload directly</center>")
            source_upload = gr.File(
                label="Upload video or audio",
                file_types=["video", "audio"],
            )
            thumbnail_preview = gr.HTML(value="")

        # Voice column
        with gr.Column(scale=1, elem_classes=["card-surface"]):
            gr.Markdown("### VOICE")
            voice_mode = gr.Radio(
                choices=["auto", "theme", "profile", "clone"],
                value="auto",
                label="Voice mode",
                info="Auto: clones voice from source. Theme: pre-built voices. Profile: celebrity voices. Clone: your own sample.",
            )
            voice_theme = gr.Dropdown(
                choices=VOICE_THEMES,
                value="narrator-m",
                label="Voice theme (supports: EN, ES, FR, DE, IT, PT, RU, JA, KO, ZH)",
                visible=False,
            )
            voice_profile = gr.Dropdown(
                choices=VOICE_PROFILES,
                value="abubakr",
                label="Voice profile",
                visible=False,
            )
            voice_sample = gr.Audio(
                label="Voice sample audio",
                type="filepath",
                visible=False,
            )
            voice_script = gr.Textbox(
                label="Transcript of the voice sample (required for Qwen TTS)",
                placeholder="Type exactly what is said in the audio sample...",
                lines=3,
                visible=False,
            )

    # ── Row: Language & Output settings ─────────────────────────────────────
    with gr.Row():
        with gr.Column(elem_classes=["card-surface"]):
            with gr.Row():
                target_lang = gr.Dropdown(
                    choices=LANGUAGES,
                    value="Arabic",
                    label="Target language",
                )
                source_lang = gr.Dropdown(
                    choices=["auto"] + LANGUAGES,
                    value="auto",
                    label="Source language",
                )
                output_type = gr.Radio(
                    choices=["audio", "video"],
                    value="audio",
                    label="Output type",
                )
            with gr.Row():
                embed_subs = gr.Checkbox(
                    value=True,
                    label="Burn subtitles into video",
                )
                asr_review = gr.Checkbox(
                    value=False,
                    label="Review transcription (LLM fixes typos/punctuation)",
                )
                translate_technical = gr.Checkbox(
                    value=False,
                    label="Translate technical terms",
                    info="When off, keeps technical terms (Python, API, etc.) in English",
                )

    # ── Advanced settings ─────────────────────────────────────────────────
    with gr.Accordion("Subtitle styling", open=False, elem_classes=["card-surface"]):
        with gr.Row():
            sub_font = gr.Dropdown(
                choices=[
                    "Noto Sans Arabic",
                    "Cairo",
                    "IBM Plex Sans Arabic",
                    "Tajawal",
                    "Amiri",
                    "Noto Sans",
                    "Roboto",
                    "Open Sans",
                    "Montserrat",
                    "Lato",
                    "Arial",
                ],
                value="Noto Sans Arabic",
                label="Subtitle font",
                allow_custom_value=True,
            )
            sub_size = gr.Slider(
                minimum=12,
                maximum=48,
                value=24,
                step=1,
                label="Font size (px)",
            )
            sub_position = gr.Dropdown(
                choices=["bottom", "top", "center"],
                value="bottom",
                label="Position",
            )
            sub_color = gr.Dropdown(
                choices=["white", "yellow", "cyan", "green", "red", "blue", "magenta"],
                value="white",
                label="Color",
            )
        with gr.Row():
            sub_bg_alpha = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.6,
                step=0.05,
                label="Background opacity",
            )
            sub_outline_width = gr.Slider(
                minimum=0,
                maximum=5,
                value=1,
                step=1,
                label="Outline width (px)",
            )
            sub_bold = gr.Checkbox(
                value=False,
                label="Bold",
            )
            sub_line_spacing = gr.Slider(
                minimum=0,
                maximum=24,
                value=8,
                step=1,
                label="Line spacing (px)",
            )
        with gr.Row():
            sub_source = gr.Radio(
                choices=["translated", "original"],
                value="translated",
                label="Subtitle source",
                info="Which subtitles to burn: translated or original language",
            )

    with gr.Accordion("Audio & tempo", open=False, elem_classes=["card-surface"]):
        with gr.Row():
            tempo_mode = gr.Dropdown(
                choices=["auto", "dynamic", "fixed", "off"],
                value="auto",
                label="Tempo mode",
                info="auto: speed up overflows. dynamic: speed up + slow down. fixed: constant rate. off: no adjustment.",
            )
            max_tempo = gr.Slider(
                minimum=1.0,
                maximum=2.0,
                value=1.5,
                step=0.05,
                label="Max tempo (speed-up limit)",
            )

    with gr.Accordion("Translation tuning", open=False, elem_classes=["card-surface"]):
        with gr.Row():
            words_per_second = gr.Slider(
                minimum=0,
                maximum=5.0,
                value=0,
                step=0.1,
                label="Words per second (0 = auto-estimate)",
                info="Target speech rate for word budget. 0 lets the pipeline auto-estimate from source.",
            )
            duration_budget = gr.Slider(
                minimum=0,
                maximum=1.0,
                value=0,
                step=0.05,
                label="Duration budget (0 = default 0.85)",
                info="Fraction of time to fill with speech. Lower = more breathing room.",
            )

    with gr.Accordion("Trim & slice", open=False, elem_classes=["card-surface"]):
        with gr.Row():
            slice_start = gr.Textbox(
                label="Start time",
                placeholder="00:01:00",
                value="",
            )
            slice_end = gr.Textbox(
                label="End time",
                placeholder="00:04:00",
                value="",
            )

    # ── Dub button ───────────────────────────────────────────────────────────
    dub_btn = gr.Button(
        "DUB THIS VIDEO",
        elem_classes=["dub-btn"],
        size="lg",
        variant="primary",
    )

    # ── Pipeline visualizer ──────────────────────────────────────────────────
    pipeline_display = gr.HTML(value=build_pipeline_html())

    # ── Logs ─────────────────────────────────────────────────────────────────
    with gr.Accordion("Logs", open=False, elem_classes=["log-area"]):
        log_output = gr.Textbox(
            label="",
            lines=12,
            max_lines=30,
            interactive=False,
            show_label=False,
            show_copy_button=True,
        )

    # ── Result section ───────────────────────────────────────────────────────
    # NOTE: Do NOT wrap outputs in gr.Group() — Gradio 5.x has a bug where
    # output components inside gr.Group break API schema introspection,
    # causing click handlers to silently do nothing.
    gr.Markdown("### RESULT")
    with gr.Row():
        video_output = gr.Video(
            label="Dubbed video",
        )
        audio_output = gr.Audio(
            label="Dubbed audio",
        )
    with gr.Row():
        srt_preview = gr.Textbox(
            label="Subtitles (SRT)",
            lines=10,
            interactive=False,
            show_copy_button=True,
        )
        srt_download = gr.File(
            label="Download SRT",
            file_types=[".srt"],
        )

    # ── Footer ───────────────────────────────────────────────────────────────
    gr.HTML("""
    <div class="studio-footer">
      mazinger by <a href="https://huggingface.co/bakrianoo" target="_blank">@bakrianoo</a>
      &nbsp;&middot;&nbsp;
      space by <a href="https://huggingface.co/HeshamHaroon" target="_blank">@HeshamHaroon</a>
    </div>
    """)

    # ── Event wiring ─────────────────────────────────────────────────────────

    # YouTube thumbnail on URL change
    source_url.change(
        fn=fetch_youtube_thumbnail,
        inputs=source_url,
        outputs=thumbnail_preview,
    )

    # Voice mode toggle
    voice_mode.change(
        fn=_voice_mode_change,
        inputs=voice_mode,
        outputs=[voice_theme, voice_profile, voice_sample, voice_script],
    )

    # Main dub pipeline
    dub_btn.click(
        fn=process_video,
        inputs=[
            source_url,
            source_upload,
            target_lang,
            source_lang,
            voice_mode,
            voice_theme,
            voice_profile,
            voice_sample,
            voice_script,
            output_type,
            embed_subs,
            sub_font,
            sub_size,
            slice_start,
            slice_end,
            sub_position,
            sub_color,
            sub_bg_alpha,
            sub_outline_width,
            sub_bold,
            sub_line_spacing,
            sub_source,
            asr_review,
            tempo_mode,
            max_tempo,
            words_per_second,
            duration_budget,
            translate_technical,
        ],
        outputs=[
            pipeline_display,
            log_output,
            video_output,
            audio_output,
            srt_preview,
            srt_download,
        ],
    )


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------

demo.queue()

if __name__ == "__main__":
    demo.launch(ssr_mode=False)