import gradio as gr
import torch
import os
import subprocess
import tempfile
from inference import main as inference_main  # Предполагается, что вы интегрируете функцию main из inference.py
import shutil

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 "Пожалуйста, загрузите видео/изображение и аудио файл, а также выберите чекпойнт."

    # Создание временной директории для обработки
    with tempfile.TemporaryDirectory() as temp_dir:
        # Копирование загруженных файлов во временную директорию
        video_path = os.path.join(temp_dir, video.name)
        audio_path = os.path.join(temp_dir, audio.name)
        with open(video_path, "wb") as f:
            f.write(video.read())
        with open(audio_path, "wb") as f:
            f.write(audio.read())

        # Определение выходного файла
        output_path = os.path.join(temp_dir, "output.mp4")

        # Подготовка аргументов для inference.py
        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:
            # Вызов функции инференса напрямую
            # Если вы интегрируете код из inference.py, вызовите функцию непосредственно
            # Например:
            # inference_main(args)
            # Но для простоты, здесь мы будем использовать subprocess
            cmd = ["python", "inference.py"] + args
            subprocess.run(cmd, check=True)
        except subprocess.CalledProcessError as e:
            return f"Произошла ошибка при обработке: {e}"

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

        # Чтение выходного файла
        with open(output_path, "rb") as f:
            output_video = f.read()

        return output_video

with gr.Blocks() as ui:
    gr.Markdown("## Wav2Lip - Синхронизация губ в видео")
    with gr.Row():
        video = gr.File(label="Видео или Изображение", type="file")
        audio = gr.File(label="Аудио", type="file")
        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
    )

ui.queue(concurrency_count=1).launch(debug=True)