import gradio as gr
import os
import subprocess
from gtts import gTTS
from pydub import AudioSegment
from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
from PIL import Image

def get_thumbnail(image_path, width):
    image = Image.open(image_path)
    image.thumbnail((width, width))
    return image

generator = pipeline('text-generation', model="checkpoints")
tokenizer = GPT2Tokenizer.from_pretrained('checkpoints')
os.environ["TOKENIZERS_PARALLELISM"] = "true"

def generate_output(name, date_of_birth, image):
    prompt = f"Tu carta astral de hoy {name} es:"
    input_tokens = tokenizer.encode(prompt, add_special_tokens=False)
    input_text = tokenizer.decode(input_tokens)
    gpt2_output = generator(input_text, max_length=120, do_sample=True, temperature=0.9)
    generated_text = gpt2_output[0]['generated_text']
    generated_text = generated_text.replace(input_text, "").strip()
    if len(gpt2_output) == 0 or 'generated_text' not in gpt2_output[0]:
        return None, "No se pudo generar el texto."

    try:
        tts = gTTS(generated_text, lang='es')
        temp_audio_path = "temp_audio.mp3"
        tts.save(temp_audio_path)
        audio_path = "audio.wav"
        audio = AudioSegment.from_mp3(temp_audio_path)
        audio.export(audio_path, format="wav")
        print("Archivo de audio generado:", audio_path)
    except Exception as e:
        return None, f"No se pudo generar el audio: {str(e)}"   

    command = f"python3 inference.py --checkpoint_path checkpoints/wav2lip_gan.pth --face {image} --audio audio.wav --outfile video.mp4 --nosmooth"
    process = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if process.returncode != 0:
        error_message = process.stderr.decode("utf-8")
        return None, f"No se pudo generar el video: {error_message}"

    output_video_path = "video.mp4"
    os.remove(temp_audio_path)

    if os.path.isfile(output_video_path):
        return output_video_path, None
    return None, "No se pudo generar el video"

def error_message_fn(error_message):
    if error_message is not None:
        return gr.outputs.Textbox(text=error_message, placeholder="Error")
    else:
        return None

def get_image_choices():
    image_paths = ["1.jpg", "2.jpg", "3.jpg", "4.jpg", "5.jpg", "6.jpg"]
    image_choices = []
    for image_path in image_paths:
        image = Image.open(image_path)
        image.thumbnail((50, 50))
        image_choices.append(image)
    return image_choices

def select_image(image):
    return image

iface = gr.Interface(
    fn=select_image,
    inputs=[
        gr.inputs.Image(label="Selecciona una imagen:", type="pil", choices=get_image_choices(), preview="label"),
        gr.inputs.Textbox(lines=1, label="Nombre", placeholder="Ingresa tu nombre"),
        gr.inputs.Textbox(lines=1, label="Fecha de Nacimiento", placeholder="DD/MM/AAAA")
    ],
    outputs=[
        gr.outputs.Video(label="Respuesta de Andrea (un minuto aproximadamente)").style(width=256),
        gr.outputs.Textbox(label="Mensaje de error", type="text")
    ],
    title="Oráculo de Inteligencia Artificial v2.1",
    layout="vertical",
    description="Por favor, ingresa tu nombre y fecha de nacimiento."
)

iface.launch()