import os
import subprocess
import requests
import gradio as gr
from moviepy.editor import *
from datetime import datetime
import logging
import re
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import warnings
import time

# Configuración inicial
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Suprimir warnings específicos
warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
warnings.filterwarnings("ignore", category=DeprecationWarning)

PEXELS_API_KEY = os.getenv("PEXELS_API_KEY")

# Lista de voces válidas
VOICES = [
    "es-MX-DaliaNeural", "es-ES-ElviraNeural", "es-AR-ElenaNeural",
    "es-MX-JorgeNeural", "es-ES-AlvaroNeural", "es-AR-TomasNeural",
    "en-US-JennyNeural", "fr-FR-DeniseNeural", "de-DE-KatjaNeural",
    "it-IT-ElsaNeural", "pt-BR-FranciscaNeural", "ja-JP-NanamiNeural"
]

# Cargar modelo GPT-2 con manejo de errores mejorado
def cargar_modelo():
    try:
        tokenizer = GPT2Tokenizer.from_pretrained("datificate/gpt2-small-spanish")
        model = GPT2LMHeadModel.from_pretrained("datificate/gpt2-small-spanish")
        logger.info("Modelo GPT-2 cargado correctamente")
        return tokenizer, model
    except Exception as e:
        logger.error(f"Error cargando modelo: {str(e)}")
        return None, None

tokenizer, model = cargar_modelo()

def generar_guion(tema, custom_script=None):
    """Genera texto basado en el tema sin estructuras predefinidas"""
    if custom_script:
        return custom_script
    
    if model is None or tokenizer is None:
        return f"Texto sobre {tema}. " * 100
    
    try:
        inputs = tokenizer(f"Escribe un texto largo sobre {tema}", return_tensors="pt", truncation=True)
        
        outputs = model.generate(
            inputs.input_ids,
            max_length=1000,
            do_sample=True,
            temperature=0.7,
            top_k=50,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id
        )
        
        return tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        logger.error(f"Error generando texto: {str(e)}")
        return f"Contenido generado sobre {tema}. " * 100

def buscar_videos(tema):
    """Busca videos en Pexels con manejo de errores mejorado"""
    try:
        headers = {"Authorization": PEXELS_API_KEY}
        response = requests.get(
            f"https://api.pexels.com/videos/search?query={tema}&per_page=3",
            headers=headers,
            timeout=10
        )
        return response.json().get("videos", [])[:3]
    except Exception as e:
        logger.error(f"Error buscando videos: {str(e)}")
        return []

def crear_video(prompt, custom_script, voz_seleccionada, musica=None):
    try:
        start_time = time.time()
        
        # 1. Generar guión
        guion = generar_guion(prompt, custom_script)
        logger.info(f"Guion generado ({len(guion)} caracteres)")
        
        # 2. Generar voz
        voz_file = "narracion.mp3"
        subprocess.run([
            'edge-tts',
            '--voice', voz_seleccionada,
            '--text', guion,
            '--write-media', voz_file
        ], check=True)
        
        audio = AudioFileClip(voz_file)
        duracion = audio.duration
        
        # 3. Buscar y descargar videos
        videos = buscar_videos(prompt) or buscar_videos("nature")
        clips = []
        
        for i, video in enumerate(videos):
            try:
                video_file = max(video['video_files'], key=lambda x: x.get('width', 0))
                temp_file = f"temp_vid_{i}.mp4"
                
                with requests.get(video_file['link'], stream=True) as r:
                    r.raise_for_status()
                    with open(temp_file, 'wb') as f:
                        for chunk in r.iter_content(chunk_size=8192):
                            f.write(chunk)
                
                clip = VideoFileClip(temp_file)
                clip = clip.subclip(0, min(10, clip.duration))  # Máximo 10 segundos por clip
                clips.append(clip)
            except Exception as e:
                logger.error(f"Error procesando video {i}: {str(e)}")
        
        # 4. Crear video final
        if not clips:
            final_clip = ColorClip((1280, 720), (0, 0, 0), duration=duracion)
        else:
            final_clip = concatenate_videoclips(clips).set_duration(duracion)
        
        final_clip = final_clip.set_audio(audio)
        
        output_file = f"video_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
        final_clip.write_videofile(
            output_file,
            fps=24,
            codec="libx264",
            audio_codec="aac",
            threads=2,
            preset='fast'
        )
        
        logger.info(f"Video creado en {time.time()-start_time:.2f} segundos")
        return output_file
        
    except Exception as e:
        logger.error(f"Error crítico: {str(e)}")
        return None
    finally:
        # Limpieza
        for f in [voz_file, *[f"temp_vid_{i}.mp4" for i in range(3)]]:
            if os.path.exists(f):
                os.remove(f)

# Interfaz simplificada y funcional
with gr.Blocks(title="Generador de Videos") as app:
    with gr.Row():
        with gr.Column():
            tema = gr.Textbox(label="Tema del video")
            guion = gr.TextArea(label="Guion personalizado (opcional)", lines=5)
            voz = gr.Dropdown(label="Voz", choices=VOICES, value=VOICES[0])
            btn = gr.Button("Generar Video")
        
        with gr.Column():
            salida = gr.Video(label="Resultado")

    btn.click(
        fn=crear_video,
        inputs=[tema, guion, voz],
        outputs=salida
    )

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Desactivado para evitar problemas
    )