import gradio as gr
import torch
import os
import subprocess
import json
try:
    from moviepy import VideoFileClip
except ImportError:
    from moviepy.editor import VideoFileClip
import langdetect
import uuid
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM

HF_TOKEN = os.environ.get("HF_TOKEN")

MODEL_PATH = "Qwen/Qwen3.5-2B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
).eval()


def generate_unique_filename(extension):
    return f"/tmp/{uuid.uuid4()}{extension}"


def cleanup_files(*files):
    for f in files:
        if f and os.path.exists(f):
            os.remove(f)


@spaces.GPU(duration=120)
def transcribe_audio(file_path):
    temp_audio = None
    if file_path.lower().endswith((".mp4", ".avi", ".mov", ".flv", ".mkv")):
        video = VideoFileClip(file_path)
        temp_audio = generate_unique_filename(".wav")
        video.audio.write_audiofile(temp_audio, logger=None)
        video.close()
        file_path = temp_audio

    output_file = generate_unique_filename(".json")
    try:
        subprocess.run(
            [
                "insanely-fast-whisper",
                "--file-name", file_path,
                "--device-id", "0",
                "--model-name", "openai/whisper-large-v3",
                "--task", "transcribe",
                "--timestamp", "chunk",
                "--transcript-path", output_file,
                "--batch-size", "24",
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        with open(output_file) as f:
            data = json.load(f)
        result = data.get("text") or " ".join(c["text"] for c in data.get("chunks", []))
    finally:
        cleanup_files(output_file)
        if temp_audio:
            cleanup_files(temp_audio)

    return result.strip()


@spaces.GPU(duration=60)
def generate_summary(transcription):
    if not transcription or transcription.startswith("Processing error"):
        return "No valid transcription to summarize."

    detected_language = langdetect.detect(transcription)

    messages = [
        {
            "role": "system",
            "content": (
                f"You are a concise video summarizer. Always respond in the same language as the input. "
                f"Detected language: {detected_language}."
            ),
        },
        {
            "role": "user",
            "content": (
                f"Summarize the following video transcription in 150-300 words, "
                f"capturing the main points and key ideas:\n\n{transcription[:30000]}"
            ),
        },
    ]

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][inputs.input_ids.shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def process_uploaded_video(video_path):
    if not video_path:
        return "No video uploaded.", ""
    try:
        transcription = transcribe_audio(video_path)
        return transcription, ""
    except Exception as e:
        return f"Processing error: {e}", ""


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎥 Video Transcription & Smart Summary
        Upload a video for transcription and AI-powered summary.
        > ⚡ Powered by Whisper Large v3 + Qwen3-2B
        """
    )

    video_input = gr.Video(label="Upload Video")
    video_button = gr.Button("🚀 Transcribe Video", variant="primary")

    with gr.Row():
        transcription_output = gr.Textbox(label="📝 Transcription", lines=12)
        summary_output = gr.Textbox(label="📊 Summary", lines=12)

    summary_button = gr.Button("✨ Generate Summary", variant="secondary")

    video_button.click(
        process_uploaded_video,
        inputs=[video_input],
        outputs=[transcription_output, summary_output],
    )
    summary_button.click(
        generate_summary,
        inputs=[transcription_output],
        outputs=[summary_output],
    )

demo.launch(theme=gr.themes.Soft())