import os
import asyncio
import logging
import tempfile
import requests
from datetime import datetime
from moviepy.editor import VideoFileClip, concatenate_videoclips, AudioFileClip, CompositeAudioClip, afx
import edge_tts
import gradio as gr

from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Inicializa tokenizer y modelo (puedes cambiar modelo si quieres)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2").eval()

def generate_script(prompt, max_length=300):
    logger.info("Generando guion...")
    inputs = tokenizer(prompt, return_tensors="pt", truncation=False)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            do_sample=True,
            top_p=0.95,
            top_k=60,
            temperature=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    logger.info(f"Guion generado, longitud: {len(text)} caracteres")
    return text

async def text_to_speech(text, voice="es-ES-ElviraNeural", output_path="voz.mp3"):
    logger.info("Generando audio TTS...")
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(output_path)
    logger.info(f"Audio guardado en {output_path}")

def download_video_sample(url):
    logger.info(f"Descargando video de ejemplo: {url}")
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    response = requests.get(url, stream=True)
    for chunk in response.iter_content(chunk_size=1024*1024):
        tmp.write(chunk)
    tmp.close()
    return tmp.name

def loop_audio_to_length(audio_clip, target_duration):
    if audio_clip.duration >= target_duration:
        return audio_clip.subclip(0, target_duration)
    loops = int(target_duration // audio_clip.duration) + 1
    audios = [audio_clip] * loops
    concatenated = concatenate_videoclips(audios, method="compose")
    return concatenated.subclip(0, target_duration)

def crear_video(prompt, musica_url=None):
    # 1. Generar guion
    guion = generate_script(prompt, max_length=300)

    # 2. TTS
    voz_archivo = "voz.mp3"
    asyncio.run(text_to_speech(guion, output_path=voz_archivo))

    # 3. Descargar videos de ejemplo (puedes reemplazar por tu búsqueda real)
    # Aquí pongo 3 clips de ejemplo (deberías poner tus URLs)
    video_urls = [
        "https://sample-videos.com/video123/mp4/240/big_buck_bunny_240p_1mb.mp4",
        "https://sample-videos.com/video123/mp4/240/big_buck_bunny_240p_1mb.mp4",
        "https://sample-videos.com/video123/mp4/240/big_buck_bunny_240p_1mb.mp4"
    ]
    clips = []
    for url in video_urls[:3]:
        video_path = download_video_sample(url)
        clip = VideoFileClip(video_path).subclip(0, 10)  # máximo 10 segundos
        clips.append(clip)

    # 4. Concatenar videos
    video_final = concatenate_videoclips(clips, method="compose")

    # 5. Cargar audio TTS
    audio_tts = AudioFileClip(voz_archivo)

    # 6. Música de fondo en loop si está definida
    if musica_url:
        musica_path = download_video_sample(musica_url)
        musica_audio = AudioFileClip(musica_path)
        # Loop música a duración voz
        musica_loop = loop_audio_to_length(musica_audio, audio_tts.duration)
        # Mezclar audio TTS y música
        mezcla = CompositeAudioClip([musica_loop.volumex(0.3), audio_tts.volumex(1.0)])
    else:
        mezcla = audio_tts

    # 7. Asignar audio al video
    video_final = video_final.set_audio(mezcla).subclip(0, audio_tts.duration)

    # 8. Guardar video final
    output_path = f"video_output_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
    video_final.write_videofile(output_path, fps=24, threads=2, logger=None)

    # 9. Limpiar archivos temporales
    os.remove(voz_archivo)
    for clip in clips:
        clip.close()

    return output_path

def run_app(prompt, musica_url):
    logger.info(f"Entrada recibida: {prompt}")
    try:
        video_path = crear_video(prompt, musica_url if musica_url.strip() else None)
        logger.info(f"Video generado en: {video_path}")
        return video_path
    except Exception as e:
        logger.error(f"Error durante la generación: {e}")
        return None

with gr.Blocks() as app:
    gr.Markdown("### Generador simple de video con texto, voz y música en loop")
    with gr.Row():
        prompt_input = gr.Textbox(label="Introduce el tema para generar el guion", lines=2)
        musica_input = gr.Textbox(label="URL de música (opcional) para usar de fondo")
    boton = gr.Button("Generar video")
    salida = gr.Video(label="Video generado")

    boton.click(run_app, inputs=[prompt_input, musica_input], outputs=salida)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)