import os
import tempfile
import gradio as gr
from transformers import (
    pipeline,
    WhisperProcessor,
    WhisperForConditionalGeneration,
    MBartForConditionalGeneration,
    MBartTokenizer,
    pipeline as hf_pipeline,
)
from huggingface_hub import hf_hub_download
import torch
from pathlib import Path

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

# Whisper (ASR)
whisper_model = pipeline("automatic-speech-recognition", model="openai/whisper-base")

# NLLB Translation
TRANSLATION_MODEL_NAME = "facebook/nllb-200-distilled-600M"
tokenizer = MBartTokenizer.from_pretrained(TRANSLATION_MODEL_NAME)
translation_model = MBartForConditionalGeneration.from_pretrained(
    TRANSLATION_MODEL_NAME
).to(DEVICE)

# Summarization
summarizer = hf_pipeline(
    "summarization", model="facebook/bart-large-cnn", device=0 if DEVICE == "cuda" else -1
)

# Language codes for NLLB
LANG_MAP = {
    "English": "eng_Latn",
    "Urdu": "urd_Arab",
    "Spanish": "spa_Latn",
    "French": "fra_Latn",
    "German": "deu_Latn",
    "Arabic": "arb_Arab",
    "Chinese": "zho_Hans",
    "Hindi": "hin_Deva",
}

# Helper Functions
def transcribe_audio(video_path: str, language: str = None) -> str:
    result = whisper_model(video_path)
    return result["text"].strip()

def translate_text(text: str, target_lang: str) -> str:
    target = LANG_MAP.get(target_lang, "eng_Latn")
    tokenizer.src_lang = "eng_Latn"
    encoded = tokenizer(text, return_tensors="pt").to(DEVICE)
    generated = translation_model.generate(
        **encoded,
        forced_bos_token_id=tokenizer.lang_code_to_id[target],
        max_length=512
    )
    return tokenizer.decode(generated[0], skip_special_tokens=True)

def summarize_text(text: str) -> str:
    max_chunk = 1000
    chunks = [text[i:i+max_chunk] for i in range(0, len(text), max_chunk)]
    summaries = []
    for chunk in chunks:
        if len(chunk.strip()) < 30:
            continue
        summary = summarizer(chunk, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
        summaries.append(summary)
    return " ".join(summaries)

def generate_subtitles(video_path: str, format="srt") -> str:
    result = whisper_model(video_path)
    text = result["text"].strip()

    # Simulate a single subtitle block for fallback (no segments in HF pipeline)
    lines = []
    if format == "srt":
        lines.append("1")
        lines.append("00:00:00,000 --> 00:00:10,000")
        lines.append(text)
        lines.append("")
    else:
        lines.append("WEBVTT\n")
        lines.append("00:00:00.000 --> 00:00:10.000")
        lines.append(text)
        lines.append("")

    return "\n".join(lines)

# Subtitle formatting helpers (future use)
def _to_srt_time(seconds: float) -> str:
    millis = int((seconds % 1) * 1000)
    secs = int(seconds) % 60
    mins = int(seconds // 60) % 60
    hrs = int(seconds // 3600)
    return f"{hrs:02}:{mins:02}:{secs:02},{millis:03}"

def _to_vtt_time(seconds: float) -> str:
    millis = int((seconds % 1) * 1000)
    secs = int(seconds) % 60
    mins = int(seconds // 60) % 60
    hrs = int(seconds // 3600)
    return f"{hrs:02}:{mins:02}:{secs:02}.{millis:03}"

# Main processing function
def process_video(video_file, target_lang):
    if video_file is None or not video_file.endswith(".mp4"):
        return ["Invalid video format. Please upload an MP4 file."] * 5

    # Step 1: Transcribe
    original_transcript = transcribe_audio(video_file)

    # Step 2: Translate
    translated_text = translate_text(original_transcript, target_lang)

    # Step 3: Summarize
    summary = summarize_text(original_transcript)

    # Step 4: Subtitles (SRT & VTT)
    srt_subs = generate_subtitles(video_file, format="srt")
    vtt_subs = generate_subtitles(video_file, format="vtt")

    base = Path(video_file).stem
    with tempfile.NamedTemporaryFile(delete=False, suffix=".srt") as srt_file:
        srt_file.write(srt_subs.encode("utf-8"))
        srt_path = srt_file.name
    with tempfile.NamedTemporaryFile(delete=False, suffix=".vtt") as vtt_file:
        vtt_file.write(vtt_subs.encode("utf-8"))
        vtt_path = vtt_file.name

    return (
        original_transcript,
        translated_text,
        summary,
        srt_path,
        vtt_path,
    )

# Gradio Interface
languages = list(LANG_MAP.keys())

with gr.Blocks(title="VidScribe AI") as demo:
    gr.Markdown("# 🎬 VidScribe AI")
    gr.Markdown("Upload a short MP4 video to auto-transcribe, translate, summarize, and subtitle it!")

    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video (MP4)")
            target_lang = gr.Dropdown(choices=languages, value="Urdu", label="Translate to")
            run_btn = gr.Button("Process", variant="primary")

        with gr.Column():
            original_out = gr.Textbox(label="Original Transcript", lines=8, interactive=False)
            translated_out = gr.Textbox(label="Translated Transcript", lines=8, interactive=False)
            summary_out = gr.Textbox(label="Summary", lines=4, interactive=False)
            srt_file = gr.File(label="Download .srt subtitles")
            vtt_file = gr.File(label="Download .vtt subtitles")

    run_btn.click(
        fn=process_video,
        inputs=[video_input, target_lang],
        outputs=[original_out, translated_out, summary_out, srt_file, vtt_file],
    )

demo.queue().launch()