app.py

import os, re, math, uuid, time, shutil, logging, tempfile, threading, requests, asyncio, numpy as np
from datetime import datetime, timedelta
from collections import Counter

import gradio as gr
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from keybert import KeyBERT
import edge_tts
from moviepy.editor import (
    VideoFileClip, AudioFileClip, concatenate_videoclips, concatenate_audioclips,
    CompositeAudioClip, AudioClip, TextClip, CompositeVideoClip, VideoClip
)

# ------------------- Configuración & Globals -------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

PEXELS_API_KEY = os.getenv("PEXELS_API_KEY")
if not PEXELS_API_KEY:
    raise RuntimeError("Debes definir PEXELS_API_KEY en 'Settings' -> 'Variables & secrets'")

# --- Modelos inicializados como None para Carga Perezosa (Lazy Loading) ---
tokenizer = None
gpt2_model = None
kw_model = None
# ---

RESULTS_DIR = "video_results"
os.makedirs(RESULTS_DIR, exist_ok=True)
TASKS = {}

# --- Lista de Voces Fija para un Arranque Instantáneo ---
SPANISH_VOICES = [
    "es-ES-ElviraNeural", "es-ES-AlvaroNeural", "es-MX-DaliaNeural", "es-MX-JorgeNeural",
    "es-AR-ElenaNeural", "es-AR-TomasNeural", "es-CO-SalomeNeural", "es-CO-GonzaloNeural",
    "es-US-PalomaNeural", "es-US-AlonsoNeural"
]

# ------------------- Funciones para cargar modelos bajo demanda -------------------
def get_tokenizer():
    global tokenizer
    if tokenizer is None:
        logger.info("Cargando tokenizer por primera vez...")
        tokenizer = GPT2Tokenizer.from_pretrained("datificate/gpt2-small-spanish")
        if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
    return tokenizer

def get_gpt2_model():
    global gpt2_model
    if gpt2_model is None:
        logger.info("Cargando modelo GPT-2 por primera vez...")
        gpt2_model = GPT2LMHeadModel.from_pretrained("datificate/gpt2-small-spanish").eval()
    return gpt2_model

def get_kw_model():
    global kw_model
    if kw_model is None:
        logger.info("Cargando modelo KeyBERT por primera vez...")
        kw_model = KeyBERT("distilbert-base-multilingual-cased")
    return kw_model

# ------------------- Funciones del Pipeline de Vídeo -------------------
def gpt2_script(prompt: str, max_len: int = 160) -> str:
    local_tokenizer = get_tokenizer()
    local_gpt2_model = get_gpt2_model()
    instruction = f"Escribe un guion corto y coherente sobre: {prompt}"
    inputs = local_tokenizer(instruction, return_tensors="pt", truncation=True, max_length=512)
    outputs = local_gpt2_model.generate(
        **inputs, max_length=max_len + inputs["input_ids"].shape[1], do_sample=True,
        top_p=0.9, top_k=40, temperature=0.7, no_repeat_ngram_size=3,
        pad_token_id=local_tokenizer.pad_token_id, eos_token_id=local_tokenizer.eos_token_id,
    )
    text = local_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text.split("sobre:")[-1].strip()[:max_len]

async def edge_tts_synth(text: str, voice: str, path: str):
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(path)

def keywords(text: str) -> list[str]:
    local_kw_model = get_kw_model()
    clean_text = re.sub(r"[^\w\sáéíóúñÁÉÍÓÚÑ]", "", text.lower())
    try:
        kws = local_kw_model.extract_keywords(clean_text, stop_words="spanish", top_n=5)
        return [k.replace(" ", "+") for k, _ in kws if k]
    except Exception as e:
        logger.warning(f"KeyBERT falló, usando método simple. Error: {e}")
        words = [w for w in clean_text.split() if len(w) > 4]
        return [w for w, _ in Counter(words).most_common(5)] or ["naturaleza"]

def pexels_search(query: str, count: int) -> list[dict]:
    res = requests.get(
        "https://api.pexels.com/videos/search",
        headers={"Authorization": PEXELS_API_KEY},
        params={"query": query, "per_page": count, "orientation": "landscape"},
        timeout=20,
    )
    res.raise_for_status()
    return res.json().get("videos", [])

def download_file(url: str, folder: str) -> str | None:
    try:
        name = uuid.uuid4().hex + ".mp4"
        path = os.path.join(folder, name)
        with requests.get(url, stream=True, timeout=60) as r:
            r.raise_for_status()
            with open(path, "wb") as f:
                for chunk in r.iter_content(1024 * 1024): f.write(chunk)
        return path if os.path.exists(path) and os.path.getsize(path) > 1000 else None
    except Exception as e:
        logger.error(f"Fallo al descargar {url}: {e}")
        return None

def loop_audio(audio_clip: AudioFileClip, duration: float) -> AudioFileClip:
    if audio_clip.duration >= duration: return audio_clip.subclip(0, duration)
    loops = math.ceil(duration / audio_clip.duration)
    return concatenate_audioclips([audio_clip] * loops).subclip(0, duration)

def make_subtitle_clips(script: str, video_w: int, video_h: int, duration: float):
    sentences = [s.strip() for s in re.split(r"[.!?¿¡]", script) if s.strip()]
    if not sentences: return []
    total_words = sum(len(s.split()) for s in sentences) or 1
    time_per_word = duration / total_words
    clips, current_time = [], 0.0
    for sentence in sentences:
        num_words = len(sentence.split())
        sentence_duration = num_words * time_per_word
        if sentence_duration < 0.1: continue
        txt_clip = (
            TextClip(sentence, fontsize=int(video_h * 0.05), color="white",
                     stroke_color="black", stroke_width=1.5, method="caption",
                     size=(int(video_w * 0.9), None), font="Arial-Bold")
            .set_start(current_time)
            .set_duration(sentence_duration)
            .set_position(("center", "bottom"))
        )
        clips.append(txt_clip)
        current_time += sentence_duration
    return clips

def make_grain_clip(size: tuple[int, int], duration: float):
    w, h = size
    def make_frame(t):
        noise = np.random.randint(0, 40, (h, w, 1), dtype=np.uint8)
        return np.repeat(noise, 3, axis=2)
    return VideoClip(make_frame, duration=duration).set_opacity(0.15)

def build_video(script_text: str, generate_script_flag: bool, voice: str, music_path: str | None) -> str:
    tmp_dir = tempfile.mkdtemp()
    try:
        script = gpt2_script(script_text) if generate_script_flag else script_text.strip()
        voice_path = os.path.join(tmp_dir, "voice.mp3")
        asyncio.run(edge_tts_synth(script, voice, voice_path))
        voice_clip = AudioFileClip(voice_path)
        video_duration = voice_clip.duration
        if video_duration < 1: raise ValueError("El audio generado es demasiado corto.")
        video_paths = []
        for kw in keywords(script):
            if len(video_paths) >= 8: break
            for video_data in pexels_search(kw, 2):
                best_file = max(video_data.get("video_files", []), key=lambda f: f.get("width", 0))
                if best_file:
                    path = download_file(best_file.get('link'), tmp_dir)
                    if path: video_paths.append(path)
                if len(video_paths) >= 8: break
        if not video_paths: raise RuntimeError("No se encontraron vídeos en Pexels.")
        segments = []
        for path in video_paths:
            try: segments.append(VideoFileClip(path))
            except Exception as e: logger.warning(f"No se pudo cargar el clip {path}: {e}")
        if not segments: raise RuntimeError("Los clips descargados no son válidos.")
        final_segments = [s.subclip(0, min(8, s.duration)) for s in segments]
        base_video = concatenate_videoclips(final_segments, method="chain")
        if base_video.duration < video_duration:
            num_loops = math.ceil(video_duration / base_video.duration)
            base_video = concatenate_videoclips([base_video] * num_loops, method="chain")
        base_video = base_video.subclip(0, video_duration)
        if music_path:
            music_clip = loop_audio(AudioFileClip(music_path), video_duration).volumex(0.20)
            final_audio = CompositeAudioClip([music_clip, voice_clip])
        else: final_audio = voice_clip
        subtitles = make_subtitle_clips(script, base_video.w, base_video.h, video_duration)
        grain_effect = make_grain_clip(base_video.size, video_duration)
        final_video = CompositeVideoClip([base_video, grain_effect, *subtitles]).set_audio(final_audio)
        output_path = os.path.join(tmp_dir, "final_video.mp4")
        final_video.write_videofile(output_path, fps=24, codec="libx264", audio_codec="aac", threads=2, logger=None)
        return output_path
    finally:
        # Intenta cerrar todos los clips de MoviePy para liberar memoria
        if 'voice_clip' in locals(): voice_clip.close()
        if 'music_clip' in locals(): music_clip.close()
        if 'base_video' in locals(): base_video.close()
        if 'final_video' in locals(): final_video.close()
        if 'segments' in locals():
            for seg in segments: seg.close()

def worker(task_id: str, mode: str, topic: str, user_script: str, voice: str, music: str | None):
    try:
        text = topic if mode == "Generar Guion con IA" else user_script
        result_tmp_path = build_video(text, mode == "Generar Guion con IA", voice, music)
        final_path = os.path.join(RESULTS_DIR, f"{task_id}.mp4")
        shutil.copy2(result_tmp_path, final_path)
        TASKS[task_id] = {"status": "done", "result": final_path, "timestamp": datetime.utcnow()}
        shutil.rmtree(os.path.dirname(result_tmp_path))
    except Exception as e:
        logger.error(f"Error en la tarea {task_id}: {e}", exc_info=True)
        TASKS[task_id] = {"status": "error", "error": str(e), "timestamp": datetime.utcnow()}

def submit_task(mode, topic, user_script, voice, music):
    content = topic if mode == "Generar Guion con IA" else user_script
    if not content.strip(): return "", "Por favor, ingresa un tema o guion."
    task_id = uuid.uuid4().hex[:8]
    TASKS[task_id] = {"status": "processing", "timestamp": datetime.utcnow()}
    threading.Thread(target=worker, args=(task_id, mode, topic, user_script, voice, music), daemon=True).start()
    return task_id, f"✅ Tarea creada con ID: {task_id}. Comprueba el estado en unos minutos."

def check_task_status(task_id):
    if not task_id or task_id not in TASKS: return None, None, "ID de tarea no válido o no encontrado."
    task_info = TASKS[task_id]
    status = task_info["status"]
    if status == "processing": return None, None, "⏳ La tarea se está procesando..."
    if status == "error": return None, None, f"❌ Error: {task_info['error']}"
    if status == "done": return task_info["result"], task_info["result"], "✅ ¡Vídeo listo!"
    return None, None, "Estado desconocido."

def janitor_thread():
    while True:
        time.sleep(3600)
        now = datetime.utcnow()
        for task_id, info in list(TASKS.items()):
            if now - info["timestamp"] > timedelta(hours=24):
                if info.get("result") and os.path.exists(info["result"]):
                    try:
                        os.remove(info["result"])
                        logger.info(f"Limpiado vídeo antiguo: {info['result']}")
                    except Exception as e:
                        logger.error(f"Error al limpiar {info['result']}: {e}")
                del TASKS[task_id]

threading.Thread(target=janitor_thread, daemon=True).start()

with gr.Blocks(title="Generador de Vídeos IA", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 Generador de Vídeos con IA")
    gr.Markdown("Crea vídeos a partir de texto, con voz, música y efectos visuales.")
    with gr.Tabs():
        with gr.TabItem("1. Crear Vídeo"):
            with gr.Row():
                with gr.Column(scale=2):
                    mode_radio = gr.Radio(["Generar Guion con IA", "Usar Mi Guion"], value="Generar Guion con IA", label="Elige el método")
                    topic_textbox = gr.Textbox(label="Tema para la IA", placeholder="Ej: La historia de la Vía Láctea")
                    script_textbox = gr.Textbox(label="Tu Guion Completo", lines=5, visible=False, placeholder="Pega aquí tu guion...")
                    voice_dropdown = gr.Dropdown(SPANISH_VOICES, value=SPANISH_VOICES[0], label="Elige una voz")
                    music_upload = gr.Audio(type="filepath", label="Música de fondo (opcional)")
                    submit_button = gr.Button("✨ Generar Vídeo", variant="primary")
                with gr.Column(scale=1):
                    task_id_output = gr.Textbox(label="ID de tu Tarea (Guárdalo)", interactive=False)
                    status_output = gr.Textbox(label="Estado", interactive=False)
                    gr.Markdown("---")
                    gr.Markdown("### ¿Cómo funciona?\n1. Elige un método y rellena el texto.\n2. **Copia el ID** que aparecerá.\n3. Ve a la pestaña **'2. Revisar Estado'**.")
        with gr.TabItem("2. Revisar Estado"):
            gr.Markdown("### Consulta el estado de tu vídeo")
            with gr.Row():
                task_id_input = gr.Textbox(label="Pega aquí el ID de tu tarea", scale=3)
                check_button = gr.Button("🔍 Verificar", scale=1)
            status_check_output = gr.Textbox(label="Estado Actual", interactive=False)
            video_output = gr.Video(label="Resultado del Vídeo")
            download_file_output = gr.File(label="Descargar Fichero")
    def toggle_textboxes(mode):
        is_ai_mode = mode == "Generar Guion con IA"
        return gr.update(visible=is_ai_mode), gr.update(visible=not is_ai_mode)
    mode_radio.change(toggle_textboxes, inputs=mode_radio, outputs=[topic_textbox, script_textbox])
    submit_button.click(submit_task, inputs=[mode_radio, topic_textbox, script_textbox, voice_dropdown, music_upload], outputs=[task_id_output, status_output])
    check_button.click(check_task_status, inputs=task_id_input, outputs=[video_output, download_file_output, status_check_output])

if __name__ == "__main__":
    demo.launch()