import gradio as gr
import torch
import yt_dlp
import os
import subprocess
import json
from transformers import AutoTokenizer, AutoModelForCausalLM
import moviepy.editor as mp
import langdetect
import uuid

HF_TOKEN = os.environ.get("HF_TOKEN")
print("Starting the program...")

model_path = "Qwen/Qwen2.5-7B-Instruct"
print(f"Loading model {model_path}...")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True)
model = model.eval()
print("Model successfully loaded.")

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

def cleanup_files(*files):
    for file in files:
        if file and os.path.exists(file):
            os.remove(file)
            print(f"Removed file: {file}")

def download_youtube_audio(url):
    print(f"Downloading audio from YouTube: {url}")
    output_path = generate_unique_filename(".wav")
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
        }],
        'outtmpl': output_path,
        'keepvideo': True,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([url])
    
    if os.path.exists(output_path + ".wav"):
        os.rename(output_path + ".wav", output_path)
    
    if not os.path.exists(output_path):
        raise FileNotFoundError(f"Error: File {output_path} not found after download.")
        
    print(f"Audio download completed. File saved at: {output_path}")
    print(f"File size: {os.path.getsize(output_path)} bytes")
    return output_path

def transcribe_audio(file_path):
    print(f"Starting transcription of file: {file_path}")
    temp_audio = None
    if file_path.endswith(('.mp4', '.avi', '.mov', '.flv')):
        print("Video file detected. Extracting audio...")
        video = mp.VideoFileClip(file_path)
        temp_audio = generate_unique_filename(".wav")
        video.audio.write_audiofile(temp_audio)
        file_path = temp_audio
    
    print(f"Does the file exist? {os.path.exists(file_path)}")
    print(f"File size: {os.path.getsize(file_path) if os.path.exists(file_path) else 'N/A'} bytes")
    
    output_file = generate_unique_filename(".json")
    command = [
        "insanely-fast-whisper",
        "--file-name", file_path,
        "--device-id", "cpu",
        "--model-name", "openai/whisper-large-v3",
        "--task", "transcribe",
        "--timestamp", "chunk",
        "--transcript-path", output_file
    ]
    print(f"Executing command: {' '.join(command)}")
    try:
        subprocess.run(command, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running insanely-fast-whisper: {e.stderr}")
        raise RuntimeError(f"Whisper transcription failed: {e.stderr}")
    
    print(f"Reading transcription file: {output_file}")
    with open(output_file, "r") as f:
        transcription_data = json.load(f)
    
    result_text = transcription_data.get("text", " ".join([chunk["text"] for chunk in transcription_data.get("chunks", [])]))
    print("Transcription completed.")
    
    cleanup_files(output_file)
    if temp_audio:
        cleanup_files(temp_audio)
    
    return result_text

def generate_summary_stream(transcription):
    if not transcription or len(transcription.strip()) < 10:
        return "Transcription is too short to summarize."
        
    print("Starting summary generation...")
    detected_language = langdetect.detect(transcription)
    
    prompt = f"""Summarize the following video transcription in 150-300 words. 
    The summary should be in the same language as the transcription, which is detected as {detected_language}.
    Please ensure that the summary captures the main points and key ideas of the transcription:
    {transcription[:300000]}..."""
    
    response, history = model.chat(tokenizer, prompt, history=[])
    print("Summary generation completed.")
    return response

def process_youtube(url):
    if not url:
        return "Please enter a YouTube URL.", ""
    
    print(f"Processing YouTube URL: {url}")
    audio_file = None
    try:
        audio_file = download_youtube_audio(url)
        transcription = transcribe_audio(audio_file)
        return transcription, ""
    except Exception as e:
        print(f"Error processing YouTube: {e}")
        return f"Error processing YouTube: {str(e)}", ""
    finally:
        cleanup_files(audio_file)
        print(f"Directory content after processing: {os.listdir('.')}")

def process_uploaded_video(video_path):
    # CAMBIO: La comprobación ahora está aquí.
    if video_path is None:
        return "Please upload a video file first.", ""
        
    print(f"Processing uploaded video: {video_path}")
    try:
        transcription = transcribe_audio(video_path)
        return transcription, ""
    except Exception as e:
        print(f"Error processing video: {e}")
        return f"Error processing video: {str(e)}", ""

print("Setting up Gradio interface...")
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎥 Video Transcription and Smart Summary
        Upload a video or provide a YouTube link to get a transcription and AI-generated summary.
        """
    )
    
    with gr.Tabs():
        with gr.TabItem("📤 Video Upload"):
            video_input = gr.Video(label="Drag and drop or click to upload")
            video_button = gr.Button("🚀 Process Video", variant="primary")
        
        with gr.TabItem("🔗 YouTube Link"):
            url_input = gr.Textbox(label="Paste YouTube URL here", placeholder="https://www.youtube.com/watch?v=...")
            url_button = gr.Button("🚀 Process URL", variant="primary")
    
    with gr.Row():
        with gr.Column():
            transcription_output = gr.Textbox(label="📝 Transcription", lines=10)
        with gr.Column():
            summary_output = gr.Textbox(label="📊 Summary", lines=10)
    
    summary_button = gr.Button("📝 Generate Summary", variant="secondary")
    
    # CAMBIO: Se han unificado las llamadas a los botones
    video_button.click(process_uploaded_video, inputs=[video_input], outputs=[transcription_output, summary_output])
    url_button.click(process_youtube, inputs=[url_input], outputs=[transcription_output, summary_output])
    summary_button.click(generate_summary_stream, inputs=[transcription_output], outputs=[summary_output])

print("Launching Gradio interface...")
demo.launch(server_name="0.0.0.0", server_port=7860)