"""
PlotWeaver — Live Commentary Translation Platform
===================================================
Event management, multi-language dubbing, live streaming.
"""

import os
import time
import tempfile
import numpy as np
import re
import soundfile as sf
import gradio as gr
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

from languages import LANGUAGES, LANGUAGE_GROUPS, ALL_LANGUAGE_NAMES, QWEN_VOICES
from tts_engine import synthesize_chunked
from qwen_engine import dub_video_qwen, translate_chunk_qwen
from pipeline import (
    load_models, transcribe, translate_text, translate_sentence,
    split_into_sentences, extract_audio_from_video, get_media_duration,
    stretch_audio_to_duration, mux_video_audio, tts_pipe_local,
)
import pipeline

# Load all models at startup
load_models()


# =============================================================================
# Helper functions
# =============================================================================

def get_voices_for_language(lang_name):
    """Get available voices for a language based on its engine."""
    config = LANGUAGES.get(lang_name, {})
    engine = config.get("tts_engine", "local")
    if engine == "qwen":
        return QWEN_VOICES
    elif engine == "yourvoic" and config.get("yourvoic_voices"):
        return config["yourvoic_voices"]
    elif engine == "local":
        return ["Default (local model)"]
    return ["Peter"]


def full_pipeline_audio(audio_input, target_language):
    """Full pipeline: English audio → target language audio."""
    if audio_input is None:
        return None, "Please upload or record audio."

    lang_config = LANGUAGES.get(target_language)
    if not lang_config:
        return None, f"Language '{target_language}' not configured."

    sample_rate, audio_array = audio_input
    audio_array = audio_array.astype(np.float32)
    if audio_array.ndim > 1:
        audio_array = audio_array.mean(axis=1)
    if audio_array.max() > 1.0 or audio_array.min() < -1.0:
        max_val = max(abs(audio_array.max()), abs(audio_array.min()))
        if max_val > 0:
            audio_array = audio_array / max_val

    log = []
    total_start = time.time()

    # ASR
    t0 = time.time()
    english = transcribe(audio_array, sample_rate)
    log.append(f"**ASR** ({time.time()-t0:.2f}s)\n{english}")
    if not english:
        return None, "ASR returned empty text."

    # MT
    t0 = time.time()
    nllb_code = lang_config["nllb"]
    translated, en_sents, tgt_sents = translate_text(english, nllb_code, fast=False)
    log.append(f"\n**Translation** ({time.time()-t0:.2f}s)")
    for e, t in zip(en_sents, tgt_sents):
        log.append(f"  EN: {e}\n  {target_language.upper()}: {t}")
    if not translated:
        return None, "Translation returned empty."

    # TTS
    t0 = time.time()
    audio_out, sr_out = synthesize_chunked(
        translated, lang_config, tts_pipe=pipeline.tts_pipe_local
    )
    log.append(f"\n**TTS** ({time.time()-t0:.2f}s) = {len(audio_out)/sr_out:.1f}s audio")

    total = time.time() - total_start
    log.append(f"\n**Total: {total:.2f}s**")

    return (sr_out, audio_out), "\n".join(log)


def full_pipeline_text(english_text, target_language, voice_name):
    """Text-only pipeline: English text → target language audio."""
    if not english_text or not english_text.strip():
        return None, "Please enter English text."

    lang_config = LANGUAGES.get(target_language)
    if not lang_config:
        return None, f"Language '{target_language}' not configured."

    log = []
    total_start = time.time()

    # MT
    t0 = time.time()
    nllb_code = lang_config["nllb"]
    translated, en_sents, tgt_sents = translate_text(english_text.strip(), nllb_code, fast=False)
    log.append(f"**Translation** ({time.time()-t0:.2f}s)")
    for e, t in zip(en_sents, tgt_sents):
        log.append(f"  EN: {e}\n  {target_language.upper()}: {t}")
    if not translated:
        return None, "Translation returned empty."

    # TTS
    t0 = time.time()
    audio_out, sr_out = synthesize_chunked(
        translated, lang_config, tts_pipe=pipeline.tts_pipe_local
    )
    log.append(f"\n**TTS** ({time.time()-t0:.2f}s) = {len(audio_out)/sr_out:.1f}s audio")

    total = time.time() - total_start
    log.append(f"\n**Total: {total:.2f}s**")

    return (sr_out, audio_out), "\n".join(log)


def dub_video(video_path, target_languages, dub_voice, chunk_seconds, progress=gr.Progress()):
    """
    Dub a video into one or more target languages.
    Routes to Qwen Omni for global languages, local pipeline for African languages.
    """
    if video_path is None:
        return None, "Please upload a video."

    if not target_languages:
        return None, "Please select at least one target language."

    results_log = []
    output_videos = []

    for lang_name in target_languages:
        lang_config = LANGUAGES.get(lang_name)
        if not lang_config:
            results_log.append(f"**{lang_name}**: not configured, skipped")
            continue

        engine = lang_config.get("tts_engine", "local")
        results_log.append(f"\n{'='*50}")
        results_log.append(f"**Dubbing: {lang_name}** (engine: {engine})")
        results_log.append(f"{'='*50}")

        try:
            if engine == "qwen":
                # Qwen Omni: end-to-end speech-to-speech (best for global languages)
                qwen_lang_name = lang_config.get("qwen_name", lang_name)
                voice = dub_voice if dub_voice in QWEN_VOICES else "Ethan"
                out_video, log_text = dub_video_qwen(
                    video_path, qwen_lang_name, voice=voice,
                    chunk_seconds=chunk_seconds, progress_fn=progress,
                )
                results_log.append(log_text)
                if out_video:
                    output_videos.append(out_video)

            else:
                # Local/YourVoic pipeline: ASR → NLLB → TTS
                work_dir = tempfile.mkdtemp(prefix=f"dub_{lang_name}_")
                extracted_audio = os.path.join(work_dir, "audio.wav")
                tgt_audio_raw = os.path.join(work_dir, "tgt_raw.wav")
                tgt_audio_aligned = os.path.join(work_dir, "tgt_aligned.wav")
                output_video = os.path.join(work_dir, f"dubbed_{lang_name}.mp4")

                progress(0.05, desc=f"{lang_name}: extracting audio...")
                extract_audio_from_video(video_path, extracted_audio)
                video_duration = get_media_duration(video_path)
                results_log.append(f"Video: {video_duration:.1f}s")

                audio_array, sr = sf.read(extracted_audio, dtype="float32")
                if audio_array.ndim > 1:
                    audio_array = audio_array.mean(axis=1)

                progress(0.15, desc=f"{lang_name}: transcribing...")
                t0 = time.time()
                english = transcribe(audio_array, sr)
                results_log.append(f"ASR: {time.time()-t0:.1f}s")
                if not english:
                    results_log.append("ASR empty — skipped")
                    continue

                progress(0.4, desc=f"{lang_name}: translating...")
                t0 = time.time()
                nllb_code = lang_config["nllb"]
                translated, _, _ = translate_text(english, nllb_code, fast=True)
                results_log.append(f"MT: {time.time()-t0:.1f}s")
                if not translated:
                    results_log.append("Translation empty — skipped")
                    continue

                progress(0.65, desc=f"{lang_name}: synthesizing...")
                t0 = time.time()
                tgt_audio, tgt_sr = synthesize_chunked(
                    translated, lang_config, tts_pipe=pipeline.tts_pipe_local
                )
                sf.write(tgt_audio_raw, tgt_audio, tgt_sr)
                tgt_duration = len(tgt_audio) / tgt_sr
                results_log.append(f"TTS: {time.time()-t0:.1f}s ({tgt_duration:.1f}s audio)")

                progress(0.85, desc=f"{lang_name}: aligning...")
                MAX_STRETCH = 1.2
                stretch_ratio = tgt_duration / video_duration

                if stretch_ratio <= MAX_STRETCH:
                    if abs(stretch_ratio - 1.0) > 0.02:
                        stretch_audio_to_duration(tgt_audio_raw, tgt_audio_aligned, video_duration)
                    else:
                        import shutil
                        shutil.copy(tgt_audio_raw, tgt_audio_aligned)
                    extend_video = False
                    final_duration = video_duration
                else:
                    import shutil
                    shutil.copy(tgt_audio_raw, tgt_audio_aligned)
                    extend_video = True
                    final_duration = tgt_duration
                    results_log.append(f"Audio longer ({stretch_ratio:.1f}x) — extending video")

                progress(0.95, desc=f"{lang_name}: combining...")
                mux_video_audio(
                    video_path, tgt_audio_aligned, output_video,
                    extend_video=extend_video, target_duration=final_duration
                )
                output_videos.append(output_video)

        except Exception as e:
            logger.exception(f"Dubbing {lang_name} failed")
            results_log.append(f"Error: {str(e)}")

    progress(1.0, desc="Done!")
    final_video = output_videos[0] if output_videos else None
    return final_video, "\n".join(results_log)


def update_voices(language):
    """Update voice dropdown when language changes."""
    voices = get_voices_for_language(language)
    return gr.update(choices=voices, value=voices[0])


# =============================================================================
# Gradio UI
# =============================================================================

EXAMPLES = [
    "And it's a brilliant goal from the striker!",
    "The referee has shown a yellow card. Corner kick for the home team.",
    "What a save by the goalkeeper! The match is heading into injury time.",
    "He dribbles past two defenders and shoots! The ball hits the back of the net!",
]

CSS = """
.main-header { text-align: center; margin-bottom: 0.5rem; }
.main-header h1 { font-size: 1.8rem; font-weight: 700; margin: 0; }
.main-header p { color: #666; font-size: 0.95rem; }
.lang-group-label { font-weight: 600; font-size: 0.85rem; color: #888; text-transform: uppercase; letter-spacing: 0.05em; margin-top: 0.5rem; }
"""

with gr.Blocks(
    title="PlotWeaver — Live Commentary Translation",
    theme=gr.themes.Soft(),
    css=CSS,
) as demo:

    gr.HTML("""
    <div class="main-header">
        <h1>PlotWeaver</h1>
        <p>Live commentary translation platform &mdash; English to 40+ languages</p>
        <p style="font-size:0.8rem; color:#999">ASR (Whisper) &rarr; MT (NLLB-200) &rarr; TTS (YourVoic + local models)</p>
    </div>
    """)

    with gr.Tabs():

        # ====== TAB 1: EVENT MANAGEMENT ======
        with gr.TabItem("Event Management"):
            gr.Markdown("### Create new event")
            gr.Markdown("Configure your live broadcast event with target languages and input source.")

            with gr.Row():
                with gr.Column(scale=2):
                    event_name = gr.Textbox(
                        label="Event name",
                        placeholder="e.g. Premier League: Arsenal vs. Chelsea",
                    )
                    with gr.Row():
                        start_time = gr.Textbox(label="Start time", placeholder="08:30 PM")
                        end_time = gr.Textbox(label="End time", placeholder="10:30 PM")
                        event_date = gr.Textbox(label="Date", placeholder="2026-06-06")

                    gr.Markdown("#### Input source")
                    input_method = gr.Radio(
                        choices=["RTMP Stream", "WebRTC (Browser)", "Direct Audio Feed"],
                        value="RTMP Stream",
                        label="Input method",
                    )

                    gr.Markdown("#### Target languages")
                    gr.Markdown("Select languages for simultaneous broadcast. Additional languages consume more stream minutes.")

                    # Language checkboxes grouped by category
                    target_langs = gr.CheckboxGroup(
                        choices=ALL_LANGUAGE_NAMES,
                        label="Languages",
                        value=["Yoruba"],
                    )

                with gr.Column(scale=1):
                    gr.Markdown("#### Estimate summary")
                    estimate_display = gr.Markdown(
                        value="**Event:** Not configured\n\n**Languages:** 1 selected\n\n**Estimated duration:** --\n\n**Total estimate:** --"
                    )
                    create_event_btn = gr.Button("Create Event", variant="primary", size="lg")
                    event_status = gr.Markdown("")

            def update_estimate(name, langs, start, end):
                n_langs = len(langs) if langs else 0
                lang_list = ", ".join(langs) if langs else "None"
                return (
                    f"**Event:** {name or 'Not set'}\n\n"
                    f"**Languages:** {n_langs} selected\n\n"
                    f"{lang_list}\n\n"
                    f"**Input:** Configured\n\n"
                    f"**Rate:** 1x (Standard)"
                )

            for inp in [event_name, target_langs, start_time, end_time]:
                inp.change(
                    fn=update_estimate,
                    inputs=[event_name, target_langs, start_time, end_time],
                    outputs=[estimate_display],
                )

            def create_event(name, langs):
                if not name:
                    return "Please enter an event name."
                if not langs:
                    return "Please select at least one language."
                return f"Event **{name}** created with {len(langs)} languages: {', '.join(langs)}"

            create_event_btn.click(
                fn=create_event,
                inputs=[event_name, target_langs],
                outputs=[event_status],
            )

        # ====== TAB 2: LIVE STUDIO ======
        with gr.TabItem("Live Studio"):
            gr.Markdown("### Live streaming translation")
            gr.Markdown("Record or stream English commentary and hear it translated in real-time.")

            with gr.Row():
                studio_language = gr.Dropdown(
                    choices=ALL_LANGUAGE_NAMES,
                    value="Yoruba",
                    label="Target language",
                )
                studio_voice = gr.Dropdown(
                    choices=get_voices_for_language("Yoruba"),
                    value=get_voices_for_language("Yoruba")[0],
                    label="Voice",
                )

            studio_language.change(
                fn=update_voices,
                inputs=[studio_language],
                outputs=[studio_voice],
            )

            with gr.Row():
                with gr.Column():
                    studio_audio_in = gr.Audio(
                        label="English commentary (upload or record)",
                        type="numpy",
                        sources=["upload", "microphone"],
                    )
                    studio_translate_btn = gr.Button("Translate", variant="primary", size="lg")

                with gr.Column():
                    studio_audio_out = gr.Audio(label="Translated audio", type="numpy", autoplay=True)
                    studio_log = gr.Markdown(label="Pipeline log")

            studio_translate_btn.click(
                fn=full_pipeline_audio,
                inputs=[studio_audio_in, studio_language],
                outputs=[studio_audio_out, studio_log],
            )

        # ====== TAB 3: VIDEO DUBBING ======
        with gr.TabItem("Video Dubbing"):
            gr.Markdown("### Video dubbing (English → multi-language)")
            gr.Markdown(
                "Upload a video with English commentary and get back a dubbed version. "
                "**Global languages** (Arabic, French, Spanish, etc.) use Qwen Omni for best quality. "
                "**African languages** (Yoruba, Hausa, etc.) use the local Whisper → NLLB → MMS-TTS pipeline."
            )

            with gr.Row():
                with gr.Column():
                    dub_video_in = gr.Video(label="Upload English video", sources=["upload"])
                    dub_languages = gr.CheckboxGroup(
                        choices=ALL_LANGUAGE_NAMES,
                        label="Target languages",
                        value=["Yoruba"],
                    )
                    with gr.Row():
                        dub_voice = gr.Dropdown(
                            choices=QWEN_VOICES,
                            value="Ethan",
                            label="Voice (for Qwen languages)",
                            info="Applies to Arabic, French, Spanish, etc. Local languages use default voice.",
                        )
                        dub_chunk_slider = gr.Slider(
                            minimum=30, maximum=300, value=120, step=10,
                            label="Chunk duration (seconds)",
                            info="Shorter = more API calls but less timeout risk.",
                        )
                    dub_btn = gr.Button("Dub Video", variant="primary", size="lg")

                with gr.Column():
                    dub_video_out = gr.Video(label="Dubbed video (download from player)")
                    dub_log = gr.Markdown(
                        label="Processing log",
                        value="Upload a video and select languages to start."
                    )

            dub_btn.click(
                fn=dub_video,
                inputs=[dub_video_in, dub_languages, dub_voice, dub_chunk_slider],
                outputs=[dub_video_out, dub_log],
            )

        # ====== TAB 4: TEXT TRANSLATION ======
        with gr.TabItem("Text \u2192 Audio"):
            gr.Markdown("### Text to translated speech")
            gr.Markdown("Type English text, choose a language, and hear the translated audio.")

            with gr.Row():
                text_language = gr.Dropdown(
                    choices=ALL_LANGUAGE_NAMES,
                    value="Yoruba",
                    label="Target language",
                )
                text_voice = gr.Dropdown(
                    choices=get_voices_for_language("Yoruba"),
                    value=get_voices_for_language("Yoruba")[0],
                    label="Voice",
                )

            text_language.change(
                fn=update_voices,
                inputs=[text_language],
                outputs=[text_voice],
            )

            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="English text",
                        placeholder="Type English football commentary here...",
                        lines=4,
                    )
                    text_btn = gr.Button("Translate to speech", variant="primary", size="lg")
                    gr.Examples(
                        examples=[[e] for e in EXAMPLES],
                        inputs=[text_input],
                        label="Example commentary",
                    )

                with gr.Column():
                    text_audio_out = gr.Audio(label="Translated audio", type="numpy", autoplay=True)
                    text_log = gr.Markdown(label="Pipeline log")

            text_btn.click(
                fn=full_pipeline_text,
                inputs=[text_input, text_language, text_voice],
                outputs=[text_audio_out, text_log],
            )

        # ====== TAB 5: RECORDINGS ======
        with gr.TabItem("Recordings & Clips"):
            gr.Markdown("### Recordings management")
            gr.Markdown(
                "Past dubbed recordings will appear here. "
                "This feature is coming soon — for now, use Video Dubbing to create new recordings "
                "and download them from the player."
            )

        # ====== TAB 6: VOICE MODELS ======
        with gr.TabItem("Voice Models"):
            gr.Markdown("### Voice model library")
            gr.Markdown("Browse available voices for each language.")

            voice_lang_select = gr.Dropdown(
                choices=ALL_LANGUAGE_NAMES,
                value="Yoruba",
                label="Select language",
            )
            voice_info = gr.Markdown()

            def show_voice_info(lang):
                config = LANGUAGES.get(lang, {})
                engine = config.get("tts_engine", "unknown")
                voices = config.get("yourvoic_voices", [])

                info = f"### {lang}\n\n"
                if engine == "qwen":
                    info += f"**Engine:** Qwen 3.5 Omni (end-to-end speech-to-speech)\n\n"
                    info += f"This is the highest quality option. Qwen handles ASR + translation + TTS in a single API call, "
                    info += f"preserving tone, emotion, and pacing from the original speaker.\n\n"
                    info += f"**Available voices ({len(QWEN_VOICES)}):** {', '.join(QWEN_VOICES[:10])}... and {len(QWEN_VOICES)-10} more\n\n"
                    info += f"All voices support all Qwen languages."
                elif engine == "yourvoic":
                    info += f"**Engine:** YourVoic API (TTS) + NLLB-200 (translation)\n\n"
                    info += f"**YourVoic language:** `{config.get('yourvoic_lang', 'N/A')}`\n\n"
                    info += f"**Available voices:** {', '.join(voices) if voices else 'Peter (default)'}"
                else:
                    info += f"**Engine:** Local pipeline (Whisper ASR + NLLB MT + MMS-TTS)\n\n"
                    info += f"**NLLB code:** `{config.get('nllb', 'N/A')}`\n\n"
                    info += "Uses locally fine-tuned models on GPU. Voice selection not available."

                return info

            voice_lang_select.change(fn=show_voice_info, inputs=[voice_lang_select], outputs=[voice_info])
            demo.load(fn=show_voice_info, inputs=[voice_lang_select], outputs=[voice_info])

    gr.Markdown("""
---
**PlotWeaver** by PlotweaverAI | Models:
[ASR](https://huggingface.co/PlotweaverAI/whisper-small-de-en) |
[MT](https://huggingface.co/PlotweaverAI/nllb-200-distilled-600M-african-6lang) |
[TTS](https://huggingface.co/PlotweaverAI/yoruba-mms-tts-new) |
[YourVoic API](https://yourvoic.com)
""")


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