import gradio as gr
import subprocess
import os
import tempfile

def generate(video, audio, checkpoint, no_smooth, resize_factor, pad_top, pad_bottom, pad_left, pad_right):
    if video is None or audio is None or checkpoint is None:
        return "Пожалуйста, загрузите видео/изображение и аудио файл, а также выберите чекпойнт."

    # Поскольку Gradio возвращает пути к файлам, используем их напрямую
    video_path = video  # Уже строка с путем к видео
    audio_path = audio  # Уже строка с путем к аудио

    # Логирование путей для отладки
    print(f"Video path: {video_path}, Audio path: {audio_path}")

    # Создание временной директории для сохранения выходного видео
    with tempfile.TemporaryDirectory() as temp_dir:
        # Определение выходного файла
        output_path = os.path.join(temp_dir, "output.mp4")
        print(f"Output path: {output_path}")

        # Подготовка аргументов для инференса
        args = [
            "--checkpoint_path", f"checkpoints/{checkpoint}.pth",
            "--segmentation_path", "checkpoints/face_segmentation.pth",
            "--enhance_face", "gfpgan",
            "--face", video_path,
            "--audio", audio_path,
            "--outfile", output_path,
            "--resize_factor", str(resize_factor),
            "--pads", str(pad_top), str(pad_bottom), str(pad_left), str(pad_right)
        ]

        if no_smooth:
            args.append("--nosmooth")

        try:
            # Вызов команды инференса
            cmd = ["python", "inference.py"] + args
            print(f"Running inference with command: {' '.join(cmd)}")
            subprocess.run(cmd, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Command failed with error: {e}")
            return f"Произошла ошибка при обработке: {e}"

        # Проверка наличия выходного файла
        if not os.path.exists(output_path):
            print("Output file does not exist.")
            return "Не удалось создать выходное видео."

        # Возвращаем путь к выходному файлу
        print(f"Output file created at: {output_path}")
        return output_path  # Gradio автоматически обработает путь и выведет видео

with gr.Blocks() as ui:
    gr.Markdown("## Wav2Lip - Синхронизация губ в видео")
    with gr.Row():
        video = gr.File(label="Видео или Изображение", type="filepath")
        audio = gr.File(label="Аудио", type="filepath")
        with gr.Column():
            checkpoint = gr.Radio(["wav2lip", "wav2lip_gan"], label="Чекпойнт", value="wav2lip_gan")
            no_smooth = gr.Checkbox(label="Без сглаживания", value=False)
            resize_factor = gr.Slider(minimum=1, maximum=4, step=1, label="Фактор изменения размера", value=1)
    with gr.Row():
        with gr.Column():
            pad_top = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ сверху")
            pad_bottom = gr.Slider(minimum=0, maximum=50, step=1, value=10, label="Отступ снизу (рекомендуется 20 для включения подбородка)")
            pad_left = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ слева")
            pad_right = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ справа")
            generate_btn = gr.Button("Сгенерировать")
        with gr.Column():
            result = gr.Video(label="Результат")

    generate_btn.click(
        generate, 
        inputs=[video, audio, checkpoint, no_smooth, resize_factor, pad_top, pad_bottom, pad_left, pad_right], 
        outputs=result,
        concurrency_limit=1  # Устанавливаем лимит на количество одновременно выполняемых операций
    )

ui.launch(debug=True)