"""
Apni Awaaz 🎙️ — Dub English video into the Hindi people actually speak.
Built for the Build Small Hackathon (June 2026).
"""

import gradio as gr
import spaces
import torch
import edge_tts
import asyncio
import subprocess
import tempfile
import os
from pathlib import Path
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    pipeline,
    BitsAndBytesConfig,
)

# ╔══════════════════════════════════════════════════════════════╗
# ║  THE PROMPT — this is the soul of the entire project        ║
# ╚══════════════════════════════════════════════════════════════╝

SYSTEM_PROMPT = """You are a dubbing translator. You translate English dialogue into the Hindi that real people actually speak at home in North India — not the stiff, Sanskritized Hindi of Doordarshan or official dubs.

RULES:
1. Use everyday Hindustani — the natural Hindi-Urdu mix people really speak.
2. NEVER use Sanskritized/शुद्ध words when a simpler one exists:
   - "प्राप्त करना" → "मिलना" / "पाना"
   - "आवश्यक" → "ज़रूरी"
   - "अत्यंत" → "बहुत" / "काफ़ी"
   - "उपयोग" → "इस्तेमाल"
   - "विचार करना" → "सोचना"
   - "संपन्न करना" → "करना" / "निपटाना"
   - "प्रतीक्षा" → "इंतज़ार"
   - "शीघ्र" → "जल्दी"
   - "अनुमति" → "इजाज़त"
   - "कृपया" → drop it or say "please"
   - "अवश्य" → "ज़रूर"
   - "उचित" → "सही" / "ठीक"
3. Keep English words Indians naturally keep: phone, office, meeting, tension, problem, time, chance, try, plan, sure, okay, sorry, thanks, bus, train, college, hospital, doctor, ticket, report, file.
4. Match the speaker's register. Casual stays casual, serious stays serious — but never sound like a newsreader.
5. Use natural fillers where they fit: "यार", "अरे", "बस", "ना", "वो", "मतलब", "basically".
6. Natural contractions: "कर लेंगे" not "कर लिया जाएगा", "हो जाएगा" not "संपन्न हो जाएगा".
7. Keep it CONCISE. Dubbed Hindi should be roughly the same length as the English. Don't pad.

EXAMPLES:
EN: "I need to get this done before the deadline"
❌ "मुझे समय-सीमा से पूर्व यह कार्य संपन्न करना आवश्यक है"
✅ "deadline से पहले ये निपटाना पड़ेगा"

EN: "That's a really good point, I hadn't thought about that"
❌ "यह एक अत्यंत उत्तम विचार है, मैंने इस पर विचार नहीं किया था"
✅ "अच्छी बात बोली, मेरे दिमाग़ में आया ही नहीं"

EN: "We should probably reconsider our approach"
❌ "हमें अपनी कार्यप्रणाली पर पुनर्विचार करना चाहिए"
✅ "लगता है अपना तरीका बदलना पड़ेगा"

EN: "I'm really sorry, I completely forgot about our meeting"
❌ "मुझे अत्यंत खेद है, मैं हमारी बैठक के विषय में पूर्णतः विस्मृत हो गया"
✅ "sorry यार, meeting पूरी तरह भूल गया"

EN: "Can you give me a moment? I need to think about this"
❌ "क्या आप मुझे कुछ क्षण प्रदान कर सकते हैं? मुझे इस विषय पर विचार करना है"
✅ "एक second दे, सोचने दे"

EN: "The situation is getting worse and we need to act fast"
❌ "स्थिति बिगड़ती जा रही है और हमें शीघ्र कार्रवाई करनी चाहिए"
✅ "हालात ख़राब हो रहे हैं, जल्दी कुछ करना पड़ेगा"

EN: "I don't think that's going to work. Let me try something else."
❌ "मुझे नहीं लगता कि यह कार्य करेगा। मुझे कोई अन्य विकल्प आज़माने दीजिए।"
✅ "ये नहीं चलेगा। कुछ और try करता हूँ।"

EN: "Look, I understand your concern, but we don't have a choice here"
❌ "देखिए, मैं आपकी चिंता समझता हूँ, परंतु हमारे पास यहाँ कोई विकल्प नहीं है"
✅ "देख, तेरी tension समझता हूँ, पर कोई चारा नहीं है"

Translate ONLY the given English text. Output ONLY the Hindi. No commentary."""


# ╔══════════════════════════════════════════════════════════════╗
# ║  MODEL LOADING                                              ║
# ╚══════════════════════════════════════════════════════════════╝

# -- Globals (loaded once, reused) --
whisper_pipe = None
llm_model = None
llm_tokenizer = None


def load_whisper():
    """Load Whisper on CPU. ZeroGPU moves it when @spaces.GPU fires."""
    global whisper_pipe
    if whisper_pipe is None:
        print("⏳ Loading Whisper...")
        whisper_pipe = pipeline(
            "automatic-speech-recognition",
            model="openai/whisper-medium",
            torch_dtype=torch.float16,
            device="cpu",
        )
        print("✅ Whisper loaded (CPU, will move to GPU at runtime)")
    return whisper_pipe


def load_llm():
    """
    Load Qwen 2.5 7B in 4-bit.
    Called inside @spaces.GPU so device_map="auto" lands on the A100.
    """
    global llm_model, llm_tokenizer
    if llm_model is None:
        print("⏳ Loading Qwen 2.5 7B...")
        model_id = "Qwen/Qwen2.5-7B-Instruct"

        bnb_cfg = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
        )
        llm_tokenizer = AutoTokenizer.from_pretrained(model_id)
        llm_model = AutoModelForCausalLM.from_pretrained(
            model_id,
            quantization_config=bnb_cfg,
            device_map="auto",
        )
        print("✅ Qwen loaded")
    return llm_model, llm_tokenizer


# Pre-download weights at startup (stays on CPU, fast re-load later)
load_whisper()


# ╔══════════════════════════════════════════════════════════════╗
# ║  PIPELINE STEPS                                             ║
# ╚══════════════════════════════════════════════════════════════╝


def extract_audio(video_path: str, out_path: str) -> str:
    subprocess.run(
        [
            "ffmpeg", "-i", video_path,
            "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
            out_path, "-y",
        ],
        check=True, capture_output=True,
    )
    return out_path


def get_duration(path: str) -> float:
    r = subprocess.run(
        ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
         "-of", "csv=p=0", path],
        capture_output=True, text=True,
    )
    return float(r.stdout.strip())


def transcribe(audio_path: str) -> list[dict]:
    """→ [{"timestamp": (start, end), "text": "..."}]"""
    pipe = load_whisper()
    result = pipe(
        audio_path,
        return_timestamps=True,
        chunk_length_s=30,
        generate_kwargs={"language": "en"},
    )
    return result["chunks"]


def translate_segment(text: str) -> str:
    model, tok = load_llm()
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": text},
    ]
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tok(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.3,
            do_sample=True,
            top_p=0.9,
        )
    resp = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    return resp.strip().split("\n")[0]  # first line only, no runaway generation


async def _tts(text: str, path: str, voice: str):
    comm = edge_tts.Communicate(text, voice)
    await comm.save(path)


def hindi_tts(text: str, path: str, voice: str = "hi-IN-MadhurNeural"):
    asyncio.run(_tts(text, path, voice))
    return path


def adjust_speed(in_path: str, out_path: str, target_sec: float) -> str:
    """Stretch/squeeze audio to fit the target duration (pitch-preserved)."""
    dur = get_duration(in_path)
    if dur <= 0 or target_sec <= 0:
        return in_path
    ratio = dur / target_sec
    ratio = max(0.5, min(2.0, ratio))          # atempo range
    subprocess.run(
        ["ffmpeg", "-i", in_path, "-filter:a", f"atempo={ratio:.4f}",
         "-y", out_path],
        check=True, capture_output=True,
    )
    return out_path


def stitch_and_merge(
    segments: list[dict],
    video_path: str,
    total_dur: float,
    tmpdir: str,
) -> str:
    """
    Build the dubbed audio track and merge it back onto the video.
    Uses pydub for clean overlay at exact timestamps.
    """
    from pydub import AudioSegment

    # silent canvas
    base = AudioSegment.silent(duration=int(total_dur * 1000), frame_rate=24000)

    for seg in segments:
        tts_file = seg["tts_path"]
        start_ms = int(seg["start"] * 1000)
        try:
            chunk = AudioSegment.from_file(tts_file)
            base = base.overlay(chunk, position=start_ms)
        except Exception as e:
            print(f"⚠️  overlay failed for segment at {seg['start']:.1f}s: {e}")

    dubbed_wav = os.path.join(tmpdir, "dubbed_track.wav")
    base.export(dubbed_wav, format="wav")

    # merge onto video (keep original video stream, replace audio)
    out_mp4 = os.path.join(tmpdir, "output.mp4")
    subprocess.run(
        [
            "ffmpeg",
            "-i", video_path,
            "-i", dubbed_wav,
            "-c:v", "copy",
            "-map", "0:v:0",
            "-map", "1:a:0",
            "-shortest",
            "-y", out_mp4,
        ],
        check=True, capture_output=True,
    )
    return out_mp4


# ╔══════════════════════════════════════════════════════════════╗
# ║  MAIN PIPELINE                                              ║
# ╚══════════════════════════════════════════════════════════════╝


@spaces.GPU(duration=300)
def dub_video(video_path: str, voice_gender: str, progress=gr.Progress()):
    if video_path is None:
        raise gr.Error("Upload a video first!")

    # ── move Whisper to the ZeroGPU A100 ──
    pipe = load_whisper()
    pipe.model.to("cuda")
    pipe.device = torch.device("cuda")

    # ── load LLM (first call downloads + quantises onto GPU) ──
    load_llm()

    voice = "hi-IN-MadhurNeural" if voice_gender == "Male" else "hi-IN-SwaraNeural"
    tmpdir = tempfile.mkdtemp(prefix="apni_")

    # 1 ── extract audio
    progress(0.05, desc="🎵 Extracting audio…")
    raw_audio = extract_audio(video_path, os.path.join(tmpdir, "raw.wav"))
    total_dur = get_duration(raw_audio)

    # safety: reject clips > 3 min to stay within GPU budget
    if total_dur > 180:
        raise gr.Error("Please keep clips under 3 minutes for now.")

    # 2 ── transcribe
    progress(0.15, desc="👂 Listening to English…")
    chunks = transcribe(raw_audio)
    if not chunks:
        raise gr.Error("Couldn't detect any speech. Try a clearer clip.")

    # 3 ── translate + TTS each segment
    translated = []
    n = len(chunks)
    for i, ch in enumerate(chunks):
        frac = 0.2 + 0.6 * (i / n)
        progress(frac, desc=f"🗣️ Dubbing segment {i + 1}/{n}…")

        start, end = ch["timestamp"]
        if start is None or end is None:
            continue
        seg_dur = end - start
        if seg_dur <= 0:
            continue

        # translate
        hindi = translate_segment(ch["text"])

        # TTS
        tts_raw = os.path.join(tmpdir, f"tts_{i}.mp3")
        hindi_tts(hindi, tts_raw, voice)

        # speed-adjust to fit original segment window
        tts_adj = os.path.join(tmpdir, f"tts_adj_{i}.wav")
        adjust_speed(tts_raw, tts_adj, seg_dur)

        translated.append({
            "start": start,
            "end": end,
            "en": ch["text"],
            "hi": hindi,
            "tts_path": tts_adj,
        })

    # 4 ── stitch + merge
    progress(0.85, desc="🎬 Stitching final video…")
    output_video = stitch_and_merge(translated, video_path, total_dur, tmpdir)

    # 5 ── build comparison log
    log_lines = []
    for s in translated:
        log_lines.append(
            f"[{s['start']:.1f}s → {s['end']:.1f}s]\n"
            f"  🇬🇧  {s['en']}\n"
            f"  🇮🇳  {s['hi']}"
        )
    log = "\n\n".join(log_lines)

    return output_video, log


# ╔══════════════════════════════════════════════════════════════╗
# ║  GRADIO UI                                                  ║
# ╚══════════════════════════════════════════════════════════════╝

CSS = """
.main-title {
    text-align: center;
    margin-bottom: 0.2em;
}
.subtitle {
    text-align: center;
    opacity: 0.7;
    font-size: 1.05em;
    margin-top: 0;
}
.example-row {
    background: var(--block-background-fill);
    border-radius: 8px;
    padding: 12px 16px;
    margin: 6px 0;
    font-size: 0.92em;
}
footer { display: none !important; }
"""

with gr.Blocks(title="Apni Awaaz", css=CSS, theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        "# 🎙️ Apni Awaaz\n"
        "#### Dub English video into the Hindi people actually speak",
        elem_classes="main-title",
    )
    gr.Markdown(
        '_No more "मुझे यह कार्य संपन्न करना आवश्यक है"_ — '
        '_just "ये करना पड़ेगा यार"_',
        elem_classes="subtitle",
    )

    with gr.Row(equal_height=True):
        # ── left column: inputs ──
        with gr.Column(scale=1):
            vid_in = gr.Video(label="Upload an English clip (< 3 min)")
            voice_radio = gr.Radio(
                ["Male", "Female"],
                value="Male",
                label="Hindi voice",
            )
            btn = gr.Button("🎬  Dub it in apni bhasha!", variant="primary", size="lg")

        # ── right column: outputs ──
        with gr.Column(scale=1):
            vid_out = gr.Video(label="Dubbed output")
            log_box = gr.Textbox(
                label="Translation log  (EN → HI)",
                lines=12,
                interactive=False,
                show_copy_button=True,
            )

    # ── "what it does" section ──
    with gr.Accordion("How is this different from normal dubbing?", open=False):
        gr.Markdown(
            "Most Hindi dubs use **शुद्ध हिंदी** — overly formal, Sanskritized language "
            "that nobody actually speaks at home.\n\n"
            "Apni Awaaz translates into **everyday Hindustani** — the natural mix of "
            "Hindi, Urdu, and English that your family actually uses at the dinner table.\n\n"
            "| Official dub | Apni Awaaz |\n"
            "|---|---|\n"
            '| "मुझे इस विषय पर विचार करने दीजिए" | "सोचने दे एक second" |\n'
            '| "यह अत्यंत मूल्यवान है" | "बहुत महँगा है यार" |\n'
            '| "कृपया मुझे अनुमति प्रदान करें" | "please, करने दे ना" |\n'
        )

    btn.click(
        fn=dub_video,
        inputs=[vid_in, voice_radio],
        outputs=[vid_out, log_box],
    )


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