"""
🎬 Movie Dubber — Any Language → Hindi
Single-file Streamlit app for Hugging Face Spaces (CPU, free tier)

ALL code is in this one file — no utils/ folder needed.
This is the most reliable way to avoid import errors on HF Spaces.

Pipeline:
  Upload MP4 → Extract Audio → Transcribe → Translate → TTS → Merge → Download
"""

import os
import gc
import math
import tempfile
import subprocess

import streamlit as st

# ─── Page config (MUST be first Streamlit call) ───────────────────────────────
st.set_page_config(
    page_title="🎬 Movie Dubber — Any Language → Hindi",
    page_icon="🎬",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
.stProgress > div > div { background-color: #FF4B4B; }
.seg-card {
    background: #1a1a2e;
    border-left: 4px solid #FF4B4B;
    border-radius: 8px;
    padding: 10px 14px;
    margin: 6px 0;
}
.spk { background:#FF4B4B; color:#fff; border-radius:10px;
        padding:2px 9px; font-size:0.78em; font-weight:700; }
.ts  { color:#999; font-size:0.82em; }
</style>
""", unsafe_allow_html=True)


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 1 — AUDIO UTILITIES
# ═══════════════════════════════════════════════════════════════════════════════

def extract_audio(video_path: str) -> str:
    """
    Pull audio from video using ffmpeg.
    Returns path to a 16 kHz mono WAV — ideal for Whisper.
    """
    out = os.path.join(tempfile.gettempdir(), "extracted_audio.wav")
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-vn",                   # no video
        "-acodec", "pcm_s16le",  # uncompressed WAV
        "-ar", "16000",          # 16 kHz
        "-ac", "1",              # mono
        out,
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return out


def detect_silence_segments(audio_path: str):
    """Return list of (start_sec, end_sec) for silent / music regions."""
    from pydub import AudioSegment
    from pydub.silence import detect_nonsilent

    audio  = AudioSegment.from_wav(audio_path)
    speech = detect_nonsilent(audio, min_silence_len=700, silence_thresh=-40)

    total_ms = len(audio)
    silence, prev = [], 0
    for s, e in speech:
        if s > prev + 200:
            silence.append((prev / 1000, s / 1000))
        prev = e
    if prev < total_ms - 200:
        silence.append((prev / 1000, total_ms / 1000))
    return silence


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 2 — TRANSCRIPTION
# ═══════════════════════════════════════════════════════════════════════════════

def transcribe_audio(audio_path: str, model_size: str = "base", progress_cb=None):
    """
    Transcribe audio with faster-whisper (INT8 quantised, CPU-optimised).
    Returns list of { start, end, text, speaker, lang }.
    """
    from faster_whisper import WhisperModel

    if progress_cb:
        progress_cb(5, f"Loading Whisper '{model_size}'…")

    model = WhisperModel(model_size, device="cpu", compute_type="int8")

    if progress_cb:
        progress_cb(20, "Transcribing…")

    raw_segs, info = model.transcribe(
        audio_path,
        beam_size=3,
        vad_filter=True,
        vad_parameters={"min_silence_duration_ms": 500},
        word_timestamps=True,
    )

    if progress_cb:
        progress_cb(50, f"Language detected: {info.language}")

    segments = []
    for i, seg in enumerate(raw_segs):
        text = seg.text.strip()
        if not text:
            continue
        speaker = _heuristic_speaker(segments, seg.start)
        segments.append({
            "start":   round(seg.start, 2),
            "end":     round(seg.end, 2),
            "text":    text,
            "speaker": speaker,
            "lang":    info.language,
        })
        if progress_cb and i % 5 == 0:
            progress_cb(min(90, 50 + i * 2), f"Segment {i+1} done…")

    del model
    gc.collect()
    return segments


def _heuristic_speaker(existing, current_start):
    """Flip speaker ID when gap > 1.5 s (lightweight alternative to pyannote)."""
    if not existing:
        return "1"
    if current_start - existing[-1]["end"] > 1.5:
        return "2" if existing[-1]["speaker"] == "1" else "1"
    return existing[-1]["speaker"]


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 3 — TRANSLATION (Helsinki-NLP MarianMT)
# ═══════════════════════════════════════════════════════════════════════════════

_DIRECT_HI = {"en-hi": "Helsinki-NLP/opus-mt-en-hi"}
_TO_EN = {
    "de": "Helsinki-NLP/opus-mt-de-en",
    "fr": "Helsinki-NLP/opus-mt-fr-en",
    "es": "Helsinki-NLP/opus-mt-es-en",
    "it": "Helsinki-NLP/opus-mt-it-en",
    "pt": "Helsinki-NLP/opus-mt-pt-en",
    "ja": "Helsinki-NLP/opus-mt-ja-en",
    "zh": "Helsinki-NLP/opus-mt-zh-en",
    "ko": "Helsinki-NLP/opus-mt-ko-en",
    "ar": "Helsinki-NLP/opus-mt-ar-en",
    "ru": "Helsinki-NLP/opus-mt-ru-en",
    "tr": "Helsinki-NLP/opus-mt-tr-en",
    "nl": "Helsinki-NLP/opus-mt-nl-en",
    "pl": "Helsinki-NLP/opus-mt-pl-en",
}
_mt_cache = {}


def _load_mt(name):
    from transformers import MarianMTModel, MarianTokenizer
    if name not in _mt_cache:
        tok = MarianTokenizer.from_pretrained(name)
        mdl = MarianMTModel.from_pretrained(name)
        mdl.eval()
        _mt_cache[name] = (tok, mdl)
    return _mt_cache[name]


def _mt(text, model_name):
    tok, mdl = _load_mt(model_name)
    inp = tok([text], return_tensors="pt", padding=True,
              truncation=True, max_length=512)
    out = mdl.generate(**inp, num_beams=2, max_length=512)
    return tok.decode(out[0], skip_special_tokens=True).strip()


def translate_to_hindi(text: str, source_lang: str = "en") -> str:
    """Translate text → Hindi. Pivots through English if no direct model."""
    if not text.strip() or source_lang == "hi":
        return text
    src = source_lang.lower()
    # Direct path
    if f"{src}-hi" in _DIRECT_HI:
        try:
            return _mt(text, _DIRECT_HI[f"{src}-hi"])
        except Exception:
            pass
    # Pivot: src → English
    if src != "en" and src in _TO_EN:
        try:
            text = _mt(text, _TO_EN[src])
        except Exception:
            pass
    # English → Hindi
    try:
        return _mt(text, _DIRECT_HI["en-hi"])
    except Exception:
        return f"[Translation failed: {text}]"


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 4 — HINDI TTS (gTTS — free, instant, no GPU)
# ═══════════════════════════════════════════════════════════════════════════════

# ── Voice options ─────────────────────────────────────────────────────────────
# gTTS uses Google's TTS via HTTP. Different TLDs give subtly different voices.
# We label them Male/Female as a UX convenience; the underlying difference is
# accent/server — gTTS doesn't expose true gender control, but the variation
# is clearly audible and useful for multi-speaker dubbing.
VOICE_OPTIONS = {
    "Male 1 (Standard)":    "com",       # US/international server
    "Male 2 (Indian)":      "co.in",     # Indian English server
    "Female 1 (UK)":        "co.uk",     # UK server — higher pitch
    "Female 2 (Australia)": "com.au",    # Australian server
    "Default":              "com",       # fallback
}

# Legacy alias kept so any existing code using SPEAKER_OPTIONS still works
SPEAKER_OPTIONS = VOICE_OPTIONS


def _text_to_wav(text: str, tld: str, out_path: str):
    """Hindi text → WAV via gTTS (MP3 intermediate)."""
    from gtts import gTTS
    from pydub import AudioSegment

    if not text.strip():
        AudioSegment.silent(500).export(out_path, format="wav")
        return

    mp3 = out_path.replace(".wav", ".mp3")
    gTTS(text=text, lang="hi", tld=tld, slow=False).save(mp3)
    AudioSegment.from_mp3(mp3).export(out_path, format="wav")
    try:
        os.remove(mp3)
    except Exception:
        pass


def _fit_to_duration(speech, target_ms):
    """Speed-fit TTS audio to match original segment duration."""
    from pydub import AudioSegment
    dur = len(speech)
    if not dur or not target_ms:
        return speech
    ratio = dur / target_ms
    if ratio > 1.3:
        mult = min(ratio, 1.8)
        speech = speech._spawn(
            speech.raw_data,
            overrides={"frame_rate": int(speech.frame_rate * mult)},
        ).set_frame_rate(speech.frame_rate)
    elif dur < target_ms:
        speech = speech + AudioSegment.silent(target_ms - dur)
    return speech


def generate_hindi_speech(
    segments,
    speaker="com",
    speaker_voice_map=None,
    speech_volume=1.2,
    progress_cb=None,
):
    """
    Generate full dubbed audio track.

    Supports per-speaker voices via speaker_voice_map:
        { "1": "com", "2": "co.uk", ... }

    If speaker_voice_map is None, all segments use `speaker` (tld string).
    Each segment → TTS with its speaker's voice → speed-fit → overlaid at
    correct timestamp on a silent base track.
    """
    from pydub import AudioSegment

    total_ms = int(max(s["end"] for s in segments) * 1000) + 2000
    track    = AudioSegment.silent(duration=total_ms)
    tmp, n   = tempfile.gettempdir(), len(segments)

    for i, seg in enumerate(segments):
        if progress_cb:
            progress_cb(int((i + 1) / n * 90), f"Synthesising {i+1}/{n}…")

        hindi = seg.get("hindi_text", "").strip()
        if not hindi:
            continue

        # Pick voice: per-speaker map → fallback to global speaker tld
        spk_id = str(seg.get("speaker", "1"))
        if speaker_voice_map and spk_id in speaker_voice_map:
            tld = speaker_voice_map[spk_id]
        else:
            tld = speaker

        start_ms  = int(seg["start"] * 1000)
        target_ms = int(seg["end"] * 1000) - start_ms
        wav = os.path.join(tmp, f"seg_{i:04d}.wav")

        try:
            _text_to_wav(hindi, tld, wav)
        except Exception:
            continue

        speech = AudioSegment.from_wav(wav)
        if speech_volume != 1.0:
            speech = speech + (20 * math.log10(max(speech_volume, 0.01)))
        speech = _fit_to_duration(speech, target_ms)
        track  = track.overlay(speech, position=start_ms)

        try:
            os.remove(wav)
        except Exception:
            pass
        gc.collect()

    out = os.path.join(tmp, "hindi_dubbed_audio.wav")
    track.export(out, format="wav")
    return out


# ── Transcription download helper ─────────────────────────────────────────────

def build_transcription_txt(segments: list) -> str:
    """
    Format transcribed segments into a readable .txt file.

    Output format per segment:
        [MM:SS - MM:SS] Speaker N:
        Original text

    Easy to read, copy, and share.
    """
    lines = []
    for seg in segments:
        start_m, start_s = divmod(int(seg["start"]), 60)
        end_m,   end_s   = divmod(int(seg["end"]),   60)
        ts = f"[{start_m:02d}:{start_s:02d} - {end_m:02d}:{end_s:02d}]"
        lines.append(f"{ts} Speaker {seg['speaker']}:")
        lines.append(seg["text"])
        lines.append("")   # blank line between segments
    return "\n".join(lines)


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 5 — VIDEO MERGE (ffmpeg, no video re-encoding)
# ═══════════════════════════════════════════════════════════════════════════════

def merge_audio_video(video_path, dubbed_audio_path, bgm_volume=0.15):
    """
    Mix Hindi speech with faint original BGM, then mux into video.
    Video stream is COPIED (no re-encode) → fast even on CPU.
    """
    out = os.path.join(tempfile.gettempdir(), "hindi_dubbed_output.mp4")
    bgm = str(round(bgm_volume, 3))

    result = subprocess.run([
        "ffmpeg", "-y",
        "-i", video_path,
        "-i", dubbed_audio_path,
        "-filter_complex",
        f"[0:a]volume={bgm}[b];[b][1:a]amix=inputs=2:duration=first[mix]",
        "-map", "0:v", "-map", "[mix]",
        "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest",
        out,
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    if result.returncode != 0:
        # Fallback: replace audio entirely (skip BGM blend)
        fallback = os.path.join(tempfile.gettempdir(), "dubbed_fallback.mp4")
        subprocess.run([
            "ffmpeg", "-y",
            "-i", video_path, "-i", dubbed_audio_path,
            "-map", "0:v", "-map", "1:a",
            "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest",
            fallback,
        ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return fallback

    return out


# ═══════════════════════════════════════════════════════════════════════════════
# SECTION 6 — SMART AUTO-ASSIGN
# Splits a Hindi dubbing text block and maps lines → segments
# Rule-based only — no ML models, fully CPU-safe
# ═══════════════════════════════════════════════════════════════════════════════

import re

# Emotion keyword lists for simple rule-based detection
_EMOTION_RULES = {
    "😠 Angry":  ["गुस्सा", "क्रोध", "चिल्लाओ", "मत करो", "बंद करो", "नहीं", "कभी नहीं"],
    "💕 Soft":   ["प्यार", "शांत", "धीरे", "माफ", "दुख", "रो", "आँसू"],
    "😄 Happy":  ["खुशी", "हाहाहा", "बढ़िया", "शाबाश", "वाह", "मज़ा", "जीत"],
    "😐 Neutral": [],  # default
}

def _detect_emotion(text: str) -> str:
    """Return emotion label for a Hindi line — pure keyword matching, no model."""
    for label, keywords in _EMOTION_RULES.items():
        if any(kw in text for kw in keywords):
            return label
    return "😐 Neutral"


def _split_hindi_text(raw: str) -> list:
    """
    Split Hindi paragraph into sentences.
    Splits on: newline, Hindi danda (।), ? !
    Strips whitespace and drops empty lines.
    """
    normalized = re.sub(r"[।?!\n]+", "\n", raw)
    return [ln.strip() for ln in normalized.splitlines() if ln.strip()]


def _merge_short_lines(lines: list, target: int) -> list:
    """If we have MORE lines than segments, merge the shortest adjacent pairs."""
    lines = list(lines)
    while len(lines) > target:
        # Find index of shortest line that has a neighbour
        idx = min(range(len(lines) - 1), key=lambda i: len(lines[i]))
        merged = lines[idx] + " " + lines[idx + 1]
        lines = lines[:idx] + [merged] + lines[idx + 2:]
    return lines


def _split_long_lines(lines: list, target: int) -> list:
    """If we have FEWER lines than segments, split the longest line at midpoint."""
    lines = list(lines)
    while len(lines) < target:
        idx = max(range(len(lines)), key=lambda i: len(lines[i]))
        line = lines[idx]
        mid = len(line) // 2

        # Find a natural break near the middle (comma or space)
        split_at = -1
        for offset in range(mid):
            for pos in [mid - offset, mid + offset]:
                if 0 < pos < len(line) and line[pos] in "،, ":
                    split_at = pos
                    break
            if split_at != -1:
                break
        if split_at == -1:
            split_at = mid  # fallback: hard split at middle

        p1 = line[:split_at].strip()
        p2 = line[split_at:].strip().lstrip("،, ")
        if not p1 or not p2:
            break  # can't split further

        lines = lines[:idx] + [p1, p2] + lines[idx + 1:]
    return lines


def auto_assign_hindi_text(segments: list, hindi_text: str) -> list:
    """
    Split a block of Hindi dubbing text and assign each line
    to the matching segment, preserving speaker + timestamps.

    Steps:
      1. Split text into lines on newline / danda / ? / !
      2. Adjust line count to match segment count:
         - Too many lines → merge shortest adjacent pairs
         - Too few lines  → split longest lines at midpoint
      3. Assign one line per segment
      4. Detect emotion per line (keyword-based, no model)
    """
    lines = _split_hindi_text(hindi_text)
    n = len(segments)
    if not lines:
        return segments

    if len(lines) > n:
        lines = _merge_short_lines(lines, n)
    elif len(lines) < n:
        lines = _split_long_lines(lines, n)

    # Safety pad
    while len(lines) < n:
        lines.append("")

    return [
        {**seg, "hindi_text": lines[i] if i < len(lines) else "",
         "emotion": _detect_emotion(lines[i] if i < len(lines) else "")}
        for i, seg in enumerate(segments)
    ]


# ═══════════════════════════════════════════════════════════════════════════════
# STREAMLIT UI
# ═══════════════════════════════════════════════════════════════════════════════

# Session state
for k in ["audio_path", "segments", "translated_segments",
          "dubbed_audio_path", "output_video_path", "speaker_voice_map"]:
    if k not in st.session_state:
        st.session_state[k] = None

# ── Sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.title("⚙️ Settings")

    st.subheader("🎙️ Transcription")
    whisper_model = st.selectbox(
        "Whisper model",
        ["tiny", "base", "small"],
        index=1,
        help="tiny=fastest | base=balanced ✅ | small=most accurate",
    )

    st.subheader("🎚️ Audio Mix")
    bgm_volume    = st.slider("Background music volume", 0.0, 1.0, 0.15, 0.05)
    speech_volume = st.slider("Hindi speech volume",     0.5, 2.0, 1.2,  0.1)

    st.subheader("ℹ️ Stack")
    st.caption("faster-whisper · Helsinki-NLP · gTTS · ffmpeg\n100% free · CPU · Open-source")

# ── Main ─────────────────────────────────────────────────────────────────────
st.title("🎬 Movie Dubber")
st.markdown("**Upload any video → Get Hindi dubbed output — 100% free**")

# STEP 1 — Upload
st.header("Step 1 — Upload Video")
uploaded = st.file_uploader(
    "MP4 / AVI / MKV / MOV (keep under ~200 MB for CPU tier)",
    type=["mp4", "avi", "mkv", "mov"],
)

if uploaded:
    video_path = os.path.join(tempfile.gettempdir(), uploaded.name)
    with open(video_path, "wb") as f:
        f.write(uploaded.read())

    st.video(video_path)
    st.success(f"✅ {uploaded.name} ({os.path.getsize(video_path)/1e6:.1f} MB)")

    # STEP 2 — Extract Audio
    st.header("Step 2 — Extract Audio")
    if st.button("🔊 Extract Audio", use_container_width=True):
        with st.spinner("Running ffmpeg…"):
            bar  = st.progress(0, "Extracting…")
            apth = extract_audio(video_path)
            bar.progress(70, "Detecting silence…")
            sil  = detect_silence_segments(apth)
            bar.progress(100, "Done!")
            st.session_state.audio_path = apth

        st.audio(apth)
        st.success("✅ Audio extracted!")
        if sil:
            with st.expander(f"🔇 {len(sil)} silence segment(s)"):
                for s, e in sil[:10]:
                    st.caption(f"  {s:.1f}s → {e:.1f}s")

    # STEP 3 — Transcribe
    if st.session_state.audio_path:
        st.header("Step 3 — Transcribe Speech")
        st.info(f"Whisper **{whisper_model}** on CPU — ~1–3 min/min of audio. ☕")

        if st.button("📝 Transcribe", use_container_width=True):
            bar  = st.progress(0, "Loading…")
            segs = transcribe_audio(
                st.session_state.audio_path,
                model_size=whisper_model,
                progress_cb=lambda p, t: bar.progress(p, t),
            )
            bar.progress(100, "Done!")
            st.session_state.segments = segs
            st.session_state.translated_segments = [{**s, "hindi_text": ""} for s in segs]
            st.session_state.speaker_voice_map = None  # reset on new transcription
            st.success(f"✅ {len(segs)} segments found!")

        # ── Download transcription as TXT ─────────────────────────────────────
        if st.session_state.segments:
            txt_content = build_transcription_txt(st.session_state.segments)
            st.download_button(
                label="⬇️ Download Transcription (.txt)",
                data=txt_content.encode("utf-8"),
                file_name="transcription.txt",
                mime="text/plain",
                use_container_width=True,
            )

        # ── Speaker voice assignment ──────────────────────────────────────────
        if st.session_state.segments:
            st.subheader("🎙️ Assign Voices to Speakers")
            st.caption(
                "Each detected speaker can have a different Hindi voice. "
                "Changes take effect when you generate Hindi speech."
            )

            # Find unique speaker IDs in order of first appearance
            seen, unique_speakers = set(), []
            for seg in st.session_state.segments:
                spk = str(seg.get("speaker", "1"))
                if spk not in seen:
                    seen.add(spk)
                    unique_speakers.append(spk)

            voice_names = list(VOICE_OPTIONS.keys())

            # Default assignments: alternate Male 1 / Female 1 for first two,
            # then Default for any extras
            defaults = ["Male 1 (Standard)", "Female 1 (UK)", "Male 2 (Indian)",
                        "Female 2 (Australia)", "Default"]

            voice_map = {}  # spk_id → tld string
            cols = st.columns(min(len(unique_speakers), 3))
            for idx, spk_id in enumerate(unique_speakers):
                col = cols[idx % len(cols)]
                with col:
                    default_voice = defaults[idx] if idx < len(defaults) else "Default"
                    chosen = st.selectbox(
                        f"Speaker {spk_id} voice",
                        voice_names,
                        index=voice_names.index(default_voice),
                        key=f"voice_spk_{spk_id}",
                    )
                    voice_map[spk_id] = VOICE_OPTIONS[chosen]

            if st.button("💾 Save Voice Assignments", use_container_width=True):
                st.session_state.speaker_voice_map = voice_map
                st.success(
                    "✅ Voice assignments saved! "
                    + " | ".join(
                        f"Speaker {k} → {[n for n,v in VOICE_OPTIONS.items() if v==tld][0]}"
                        for k, tld in voice_map.items()
                    )
                )

    # STEP 4 — Translate
    if st.session_state.segments:
        st.header("Step 4 — Translate to Hindi")

        tab_auto, tab_machine = st.tabs([
            "✍️ Write Hindi Yourself (Smart Assign)",
            "🤖 Auto Machine-Translate",
        ])

        # ── Tab A: Smart assign from user's own Hindi block ──────────────────
        with tab_auto:
            st.markdown(
                "Write your full Hindi dubbing script in **one box** — "
                "the app will split it and assign each line to the right speaker automatically."
            )
            st.markdown(
                "**Supported separators:** newline · Hindi danda `।` · `?` · `!`"
            )

            # Show original lines for reference
            with st.expander("📄 Show original transcribed lines for reference"):
                for seg in st.session_state.segments:
                    st.caption(
                        f"[{seg['start']:.1f}s → {seg['end']:.1f}s] "
                        f"Speaker {seg['speaker']}: {seg['text']}"
                    )

            hindi_block = st.text_area(
                "Write full Hindi dubbing text here",
                height=250,
                placeholder=(
                    "अरे सुनो, तुम कहाँ जा रहे हो?\n"
                    "मैं बाजार जा रहा हूँ।\n"
                    "ठीक है, जल्दी आना।"
                ),
                key="hindi_block_input",
            )

            n_segs  = len(st.session_state.segments)
            n_lines = len(_split_hindi_text(hindi_block)) if hindi_block.strip() else 0

            # Live line-count feedback
            if hindi_block.strip():
                if n_lines == n_segs:
                    st.success(f"✅ {n_lines} lines detected — perfect match with {n_segs} segments!")
                elif n_lines > n_segs:
                    st.warning(
                        f"⚠️ {n_lines} lines, {n_segs} segments — "
                        f"{n_lines - n_segs} extra line(s) will be merged automatically."
                    )
                else:
                    st.warning(
                        f"⚠️ {n_lines} lines, {n_segs} segments — "
                        f"{n_segs - n_lines} missing line(s) will be split automatically."
                    )

            if st.button("🎯 Auto Assign to Speakers", use_container_width=True,
                         disabled=not hindi_block.strip()):
                assigned = auto_assign_hindi_text(
                    st.session_state.segments, hindi_block
                )
                st.session_state.translated_segments = assigned
                st.success("✅ Hindi text assigned to all speakers!")

                # Preview result
                st.markdown("#### Assignment preview")
                for seg in assigned:
                    emotion = seg.get("emotion", "")
                    st.markdown(
                        f'<div class="seg-card">'
                        f'<span class="spk">Speaker {seg["speaker"]}</span> '
                        f'<span class="ts">{seg["start"]:.1f}s → {seg["end"]:.1f}s</span> '
                        f'<span style="font-size:0.8em;color:#aaa;">{emotion}</span><br>'
                        f'{seg["hindi_text"]}'
                        f'</div>',
                        unsafe_allow_html=True,
                    )

        # ── Tab B: Machine translation (original behaviour) ───────────────────
        with tab_machine:
            st.info(
                "Automatically translates each segment using Helsinki-NLP MarianMT. "
                "Free, offline, CPU-only."
            )
            if st.button("🌏 Auto Translate All Segments", use_container_width=True):
                segs = st.session_state.segments
                bar  = st.progress(0, "Loading translation model…")
                out  = []
                for i, seg in enumerate(segs):
                    bar.progress(int((i+1)/len(segs)*100), f"Translating {i+1}/{len(segs)}…")
                    out.append({**seg, "hindi_text": translate_to_hindi(
                        seg["text"], source_lang=seg.get("lang", "en")
                    )})
                    gc.collect()
                st.session_state.translated_segments = out
                st.success("✅ Translation done!")

    # STEP 5 — Edit
    if st.session_state.translated_segments and any(
        s.get("hindi_text") for s in st.session_state.translated_segments
    ):
        st.header("Step 5 — Review & Edit Translations")
        st.caption("Fix Hindi text before generating speech.")

        edited = []
        for i, seg in enumerate(st.session_state.translated_segments):
            c1, c2 = st.columns(2)
            with c1:
                st.markdown(
                    f'<div class="seg-card">'
                    f'<span class="spk">Speaker {seg.get("speaker","?")}</span> '
                    f'<span class="ts">{seg["start"]:.1f}s→{seg["end"]:.1f}s</span><br>'
                    f'<b>Original:</b> {seg["text"]}</div>',
                    unsafe_allow_html=True,
                )
            with c2:
                new = st.text_area(
                    f"Hindi #{i+1}", value=seg.get("hindi_text",""),
                    key=f"hi_{i}", height=100, label_visibility="collapsed",
                )
            edited.append({**seg, "hindi_text": new})

        if st.button("💾 Save Edits", use_container_width=True):
            st.session_state.translated_segments = edited
            st.success("✅ Saved!")

        # STEP 6 — TTS
        st.header("Step 6 — Generate Hindi Speech")
        if st.button("🎤 Generate Hindi Audio", use_container_width=True):
            bar  = st.progress(0, "Starting…")
            dpth = generate_hindi_speech(
                st.session_state.translated_segments,
                speaker="com",                                # global fallback tld
                speaker_voice_map=st.session_state.speaker_voice_map,
                speech_volume=speech_volume,
                progress_cb=lambda p, t: bar.progress(p, t),
            )
            bar.progress(100, "Done!")
            st.session_state.dubbed_audio_path = dpth
            st.audio(dpth)
            st.success("✅ Hindi audio ready!")

        # STEP 7 — Merge
        if st.session_state.dubbed_audio_path:
            st.header("Step 7 — Merge & Download")
            if st.button("🎬 Create Dubbed Video", use_container_width=True):
                with st.spinner("Merging…"):
                    out_path = merge_audio_video(
                        video_path=video_path,
                        dubbed_audio_path=st.session_state.dubbed_audio_path,
                        bgm_volume=bgm_volume,
                    )
                    st.session_state.output_video_path = out_path

                st.video(out_path)
                st.success("🎉 Done!")
                with open(out_path, "rb") as f:
                    st.download_button(
                        "⬇️ Download Hindi Dubbed Video",
                        data=f,
                        file_name="hindi_dubbed.mp4",
                        mime="video/mp4",
                        use_container_width=True,
                    )

else:
    st.markdown("""
    ---
    ### How it works

    | Step | What happens |
    |------|-------------|
    | 1️⃣ Upload | Any MP4 / AVI / MKV video |
    | 2️⃣ Extract | ffmpeg pulls the audio track |
    | 3️⃣ Transcribe | Whisper converts speech → text with timestamps |
    | 4️⃣ Translate | MarianMT translates → natural Hindi |
    | 5️⃣ Edit | Fix any translation manually |
    | 6️⃣ Speak | gTTS generates Hindi audio per segment |
    | 7️⃣ Merge | ffmpeg mixes Hindi audio back into video |

    > ⚠️ On CPU: transcription takes ~1–3 min per minute of video. Grab a chai ☕
    """)