import os
import sys
import queue
import threading
import torch
import numpy as np
import librosa
import subprocess
import srt
import yt_dlp
import ffmpeg
import whisperx
import gradio as gr
import time
from scipy.io.wavfile import write
from transformers import MarianMTModel, MarianTokenizer, VitsModel, AutoTokenizer
from typing import Dict, Any, List

# =============================================================================
# GLOBAL CONFIGURATION
# =============================================================================

WHISPER_MODEL = "large-v2"
SRC_LANGUAGE  = "en"
OUTPUT_VIDEO  = "video.mp4"
OUTPUT_AUDIO  = "audio.wav"
OUTPUT_SRT    = "subtitles.srt"

TRANSLATED_SRT = "subtitles_fr.srt"
MT_MODEL_NAME  = "Helsinki-NLP/opus-mt-en-fr"
BATCH_SIZE     = 16

TTS_MODEL_NAME      = "facebook/mms-tts-fra"
OUTPUT_MIXED_AUDIO  = "mixed_audio.wav"
FINAL_VIDEO         = "final_dubbed_video.mp4"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# =============================================================================
# PROCESSING LOGIC (STAGES 1, 2, 3)
# =============================================================================

def download_video(url: str, output_path: str):
    ydl_opts = {
        "format": "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best",
        "outtmpl": output_path,
        "quiet": True,
        "overwrites": True
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([url])

def extract_audio(input_video: str, output_audio: str):
    (
        ffmpeg
        .input(input_video)
        .output(output_audio, acodec="pcm_s16le", ac=1, ar="16000")
        .run(overwrite_output=True, capture_stdout=True, capture_stderr=True)
    )

def transcribe_and_align(audio_path: str, device: str, lang: str) -> Dict[str, Any]:
    model = whisperx.load_model(WHISPER_MODEL, device, compute_type="float16" if device == "cuda" else "int8")
    audio = whisperx.load_audio(audio_path)
    result = model.transcribe(audio, language=lang)
    model_a, metadata = whisperx.load_align_model(language_code=lang, device=device)
    return whisperx.align(result["segments"], model_a, metadata, audio, device)

def format_srt_precise(aligned_result: Dict[str, Any], max_words=9, min_words=4) -> str:
    srt_lines = []
    idx = 1
    def fmt_time(s):
        ms = int(round((s % 1) * 1000))
        return f"{int(s//3600):02}:{int((s%3600)//60):02}:{int(s%60):02},{ms:03}"
    for seg in aligned_result["segments"]:
        words = seg.get("words", [])
        if not words: continue
        i = 0
        while i < len(words):
            end_idx = min(i + max_words, len(words))
            split_at = end_idx
            for j in range(end_idx - 1, i + min_words - 1, -1):
                w = words[j]["word"].strip().lower()
                if any(w.endswith(p) for p in ".!?;"):
                    split_at = j + 1
                    break
            chunk = words[i:split_at]
            start_t = next((w["start"] for w in chunk if "start" in w), 0)
            end_t = next((w["end"] for w in reversed(chunk) if "end" in w), start_t + 1)
            text = " ".join(w["word"].strip() for w in chunk)
            srt_lines.append(f"{idx}\n{fmt_time(start_t)} --> {fmt_time(end_t)}\n{text.strip()}\n")
            idx += 1
            i = split_at
    return "".join(srt_lines)

def translate_batch(texts, model, tokenizer, batch_size):
    translated_all = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i: i + batch_size]
        inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=128)
        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        translated_all.extend([t.strip() for t in decoded])
    return translated_all

def synthesize_audio(subs, model, tokenizer, device):
    sampling_rate = model.config.sampling_rate
    chunks = []
    for sub in subs:
        inputs = tokenizer(sub.content.strip(), return_tensors="pt").to(device)
        with torch.no_grad():
            waveform = model(**inputs).waveform
        audio = waveform.squeeze().cpu().numpy().astype(np.float32)
        chunks.append({
            "start": sub.start.total_seconds(),
            "audio": audio,
            "duration": len(audio) / sampling_rate
        })
    return chunks, sampling_rate

def assemble_dub(chunks, sr):
    max_duration = chunks[-1]["start"] + chunks[-1]["duration"] + 120
    dubbed_audio = np.zeros(int(max_duration * sr), dtype=np.float32)
    playhead = 0.0
    for chunk in chunks:
        if chunk["start"] > playhead + 1.5: playhead = chunk["start"]
        elif chunk["start"] > playhead: playhead += 0.1 
        start_idx = int(playhead * sr)
        end_idx = start_idx + len(chunk["audio"])
        dubbed_audio[start_idx:end_idx] += chunk["audio"]
        playhead = end_idx / sr
    return dubbed_audio[:int(playhead * sr)]

def apply_ducking_and_merge(video_path, dub_audio, sr, subs):
    orig_y, _ = librosa.load(video_path, sr=sr, mono=True)
    vol_mask = np.ones_like(orig_y)
    for sub in subs:
        s_idx, e_idx = int(sub.start.total_seconds() * sr), int(sub.end.total_seconds() * sr)
        vol_mask[s_idx:min(e_idx, len(vol_mask))] = 0.1 
    orig_ducked = orig_y * vol_mask
    max_len = max(len(orig_ducked), len(dub_audio))
    mixed = np.zeros(max_len, dtype=np.float32)
    mixed[:len(orig_ducked)] += orig_ducked
    mixed[:len(dub_audio)] += dub_audio * 0.9 
    mixed_stereo = np.stack([mixed, mixed], axis=-1)
    write(OUTPUT_MIXED_AUDIO, sr, (mixed_stereo * 32767).astype(np.int16))
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path, "-i", OUTPUT_MIXED_AUDIO,
        "-vf", f"scale=ceil(iw/2)*2:ceil(ih/2)*2,subtitles={TRANSLATED_SRT}", "-c:v", "libx264", "-c:a", "aac",
        "-map", "0:v", "-map", "1:a", "-shortest", FINAL_VIDEO
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

# =============================================================================
# PIPELINE MANAGEMENT
# =============================================================================

class _StreamToQueue:
    def __init__(self, log_queue: queue.Queue):
        self.queue = log_queue
        self._original = sys.stdout
    def write(self, msg):
        if msg.strip(): self.queue.put(msg)
        self._original.write(msg)
    def flush(self):
        self._original.flush()

def _run_full_pipeline(url: str, log_queue: queue.Queue, result: dict):
    try:
        print("[*] STAGE 1: Download & Transcribe")
        download_video(url, OUTPUT_VIDEO)
        extract_audio(OUTPUT_VIDEO, OUTPUT_AUDIO)
        aligned = transcribe_and_align(OUTPUT_AUDIO, DEVICE, SRC_LANGUAGE)
        with open(OUTPUT_SRT, "w", encoding="utf-8") as f:
            f.write(format_srt_precise(aligned))

        print("[*] STAGE 2: Translate Subtitles")
        subs = list(srt.parse(open(OUTPUT_SRT, "r", encoding="utf-8").read()))
        mt_tokenizer = MarianTokenizer.from_pretrained(MT_MODEL_NAME)
        mt_model = MarianMTModel.from_pretrained(MT_MODEL_NAME).to(DEVICE)
        translated_texts = translate_batch([sub.content.strip() for sub in subs], mt_model, mt_tokenizer, BATCH_SIZE)
        translated_subs = [srt.Subtitle(index=s.index, start=s.start, end=s.end, content=t) for s, t in zip(subs, translated_texts)]
        with open(TRANSLATED_SRT, "w", encoding="utf-8") as f:
            f.write(srt.compose(translated_subs))

        print("[*] STAGE 3: TTS & Final Mix")
        tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_NAME)
        tts_model = VitsModel.from_pretrained(TTS_MODEL_NAME).to(DEVICE)
        chunks, sr = synthesize_audio(translated_subs, tts_model, tts_tokenizer, DEVICE)
        final_dub = assemble_dub(chunks, sr)
        apply_ducking_and_merge(OUTPUT_VIDEO, final_dub, sr, translated_subs)

        result["video"] = FINAL_VIDEO
        print("\n✅ Process Complete!")
    except Exception as exc:
        result["error"] = str(exc)
        print(f"\n❌ Error: {exc}")
    finally:
        log_queue.put(None) 

def gradio_run_pipeline(youtube_url: str):
    log_queue, result = queue.Queue(), {"video": None, "error": None}
    old_stdout = sys.stdout
    sys.stdout = _StreamToQueue(log_queue)
    thread = threading.Thread(target=_run_full_pipeline, args=(youtube_url, log_queue, result), daemon=True)
    thread.start()
    log_text = ""
    try:
        while True:
            try:
                msg = log_queue.get(timeout=0.5)
            except queue.Empty:
                yield log_text, None
                continue
            if msg is None: break
            log_text += msg if msg.endswith("\n") else msg + "\n"
            yield log_text, None
    finally:
        sys.stdout = old_stdout
    thread.join()
    yield log_text, result["video"]

# =============================================================================
# GRADIO UI — FINAL DESIGN
# =============================================================================

custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;800&family=Fira+Code&display=swap');
@import url('https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css');

body { 
    background: radial-gradient(circle at center, #1a1a24 0%, #0b0e14 100%); 
    font-family: 'Inter', sans-serif; 
    color: #e2e8f0; 
}

.center-header { text-align: center; padding: 3rem 0; margin-bottom: 1rem; }
.main-title { 
    font-size: 3.5rem !important; 
    font-weight: 800; 
    background: linear-gradient(90deg, #ffffff 0%, #ff4d4d 100%);
    -webkit-background-clip: text; 
    -webkit-text-fill-color: transparent;
    margin-bottom: 0.5rem;
    filter: drop-shadow(0 0 10px rgba(255, 77, 77, 0.3));
}
.sub-title { font-size: 1.25rem !important; color: #a0aec0; }
.dev-name { color: #ff4d4d; font-weight: 800; text-shadow: 0 0 8px rgba(255, 77, 77, 0.4); }

/* Custom Placeholder Sizing & Glassmorphism */
#url-input { 
    flex-grow: 3.8 !important; 
}
#url-input textarea { 
    height: 55px !important; 
    font-size: 16px !important;
    background: rgba(255, 255, 255, 0.05) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 12px !important;
    backdrop-filter: blur(10px);
    color: white !important;
    transition: all 0.3s ease;
}
#url-input textarea:focus {
    border-color: #ff4d4d !important;
    box-shadow: 0 0 15px rgba(255, 77, 77, 0.3) !important;
}

/* Custom Button Sizing */
#run-btn { 
    flex-grow: 0.9 !important; 
    background: linear-gradient(135deg, #ff0000 0%, #990000 100%) !important;
    color: white !important; 
    font-weight: 700 !important; 
    border-radius: 12px !important; 
    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
    height: 55px !important;
    border: none !important;
    box-shadow: 0 4px 15px rgba(255, 0, 0, 0.4) !important;
    font-size: 16px !important;
}
#run-btn:hover { 
    transform: translateY(-3px); 
    box-shadow: 0 8px 25px rgba(255, 0, 0, 0.6) !important; 
}
#run-btn:active { 
    transform: translateY(1px); 
}

#log-box textarea { 
    background: rgba(10, 15, 25, 0.8) !important; 
    color: #00ffcc !important; 
    font-family: 'Fira Code', monospace !important; 
    border-radius: 15px !important;
    border: 1px solid rgba(255, 255, 255, 0.05) !important;
    box-shadow: inset 0 0 20px rgba(0,0,0,0.5);
    padding: 1rem !important;
}
#video-out { 
    border: 1px solid rgba(255, 77, 77, 0.2) !important; 
    border-radius: 15px !important; 
    background: rgba(0, 0, 0, 0.4) !important;
    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.5) !important;
    max-height: 450px !important;
    object-fit: contain !important;
    overflow: hidden !important;
}
#video-out video {
    max-height: 450px !important;
    object-fit: contain !important;
}

.section-title {
    font-size: 1.4rem;
    font-weight: 600;
    color: #e2e8f0;
    margin-bottom: 1rem;
    display: flex;
    align-items: center;
    gap: 0.5rem;
}
.section-title i {
    color: #ff4d4d;
    filter: drop-shadow(0 0 5px rgba(255, 77, 77, 0.5));
}
"""

with gr.Blocks(title="Video Dubbing AI", css=custom_css, theme=gr.themes.Default(neutral_hue="slate")) as demo:
    # Titles with bold colored names
    with gr.Column(elem_classes="center-header"):
        gr.HTML("<h1 class='main-title'><i class='fa-solid fa-language fa-sm'></i> AutoDub AI</h1>")
        gr.HTML('<p class="sub-title">Developed by <span class="dev-name">Meghazi Oussama</span> & <span class="dev-name">Benyettou Selsabil</span></p>')

    # Input Row with sizing adjustments
    with gr.Row():
        url_input = gr.Textbox(
            placeholder="Paste YouTube URL here...", 
            container=False, 
            elem_id="url-input"
        )
        run_btn = gr.Button("START DUBBING", variant="primary", elem_id="run-btn")

    # Reveal Box
    with gr.Column(visible=False) as result_box:
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML('<div class="section-title"><i class="fa-solid fa-microchip"></i> Neural Processing Log</div>')
                log_out = gr.Textbox(label="", lines=12, interactive=False, elem_id="log-box")
            with gr.Column(scale=1):
                gr.HTML('<div class="section-title"><i class="fa-solid fa-film"></i> Final Dubbed Video</div>')
                video_out = gr.Video(label="", elem_id="video-out")

    def start_process(url):
        if not url: return gr.update(visible=False), "", None
        yield gr.update(visible=True), "Initializing pipeline...", None
        for log, vid in gradio_run_pipeline(url):
            yield gr.update(visible=True), log, vid

    run_btn.click(fn=start_process, inputs=url_input, outputs=[result_box, log_out, video_out])

if __name__ == "__main__":
    demo.queue().launch(share=True)