import os
import sys
import subprocess
import gradio as gr
from pydub import AudioSegment
import tempfile
from scipy.io.wavfile import write, read
from TTS.api import TTS
import numpy as np
import torch
import torchaudio
from resemble_enhance.enhancer.inference import denoise

# Установка переменных окружения для принятия лицензионных условий
os.environ["COQUI_TOS_AGREED"] = "1"

# Определение устройства (CUDA или CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Глобальные переменные и настройки
language_options = {
    "English (en)": "en",
    "Spanish (es)": "es",
    "French (fr)": "fr",
    "German (de)": "de",
    "Italian (it)": "it",
    "Portuguese (pt)": "pt",
    "Polish (pl)": "pl",
    "Turkish (tr)": "tr",
    "Russian (ru)": "ru",
    "Dutch (nl)": "nl",
    "Czech (cs)": "cs",
    "Arabic (ar)": "ar",
    "Chinese (zh-cn)": "zh-cn",
    "Japanese (ja)": "ja",
    "Hungarian (hu)": "hu",
    "Korean (ko)": "ko",
    "Hindi (hi)": "hi"
}

other_language = {
    "Vietnamese": "vie",
    "Serbian": "srp",
    "Romanian": "ron",
    "Indonesian": "ind",
    "Philippine": "tgl",
    "Bulgarian": "bul",
}

# Инициализация модели TTS
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2")

# Функции для голосового клонирования
def check_audio_length(audio_path, max_duration=120):
    try:
        audio = AudioSegment.from_file(audio_path)
        duration = audio.duration_seconds
        if duration > max_duration:
            print(f"Audio is too long: {duration} seconds. Max allowed is {max_duration} seconds.")
            return False
        return True
    except Exception as e:
        print(f"Error while checking audio length: {e}")
        return False

def synthesize_and_convert_voice(text, language_iso, voice_audio_path, speed):
    # Синтез речи с помощью TTS
    tts_synthesis = TTS(model_name=f"tts_models/{language_iso}/fairseq/vits")
    wav_data = tts_synthesis.tts(text, speed=speed)

    # Преобразование wav_data из списка в NumPy массив с типом float32
    wav_data_np = np.array(wav_data, dtype=np.float32)

    # Нормализация данных, если необходимо
    max_val = np.max(np.abs(wav_data_np))
    if max_val > 1.0:
        wav_data_np = wav_data_np / max_val

    # Масштабирование до int16 для записи в WAV файл
    wav_data_int16 = np.int16(wav_data_np * 32767)

    # Сохранение синтезированного аудио во временный файл
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_tts_wav_file:
        temp_tts_wav_path = temp_tts_wav_file.name
        write(temp_tts_wav_path, 22050, wav_data_int16)

    # Загрузка синтезированного аудио
    wav_tensor, sample_rate = torchaudio.load(temp_tts_wav_path)

    # Преобразование в моно, если аудио стерео
    if wav_tensor.dim() > 1 and wav_tensor.size(0) > 1:
        wav_tensor = wav_tensor.mean(dim=0, keepdim=True)

    # Применение денойзинга (не перемещаем wav_tensor на устройство)
    denoised_wav_tensor, denoised_sample_rate = denoise(wav_tensor.squeeze(), sample_rate, device)

    # Сохранение денойзенного аудио во временный файл
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_denoised_wav_file:
        temp_denoised_wav_path = temp_denoised_wav_file.name
        torchaudio.save(temp_denoised_wav_path, denoised_wav_tensor.unsqueeze(0).cpu(), denoised_sample_rate)

    # Преобразование голоса с использованием денойзенного аудио
    tts_conversion = TTS(model_name="voice_conversion_models/multilingual/vctk/freevc24", progress_bar=False)

    # Подготовка временного выходного файла
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_output_wav_file:
        temp_output_wav_path = temp_output_wav_file.name

    # Преобразование голоса
    tts_conversion.voice_conversion_to_file(temp_denoised_wav_path, target_wav=voice_audio_path,
                                            file_path=temp_output_wav_path)

    # Чтение преобразованного аудио
    output_sample_rate, output_audio_data = read(temp_output_wav_path)

    # Удаление временных файлов
    os.remove(temp_tts_wav_path)
    os.remove(temp_denoised_wav_path)
    os.remove(temp_output_wav_path)

    return (output_sample_rate, output_audio_data)


def synthesize_speech(text, speaker_wav_path, language_iso, speed):
    # Загрузка аудио говорящего
    speaker_wav_tensor, speaker_sample_rate = torchaudio.load(speaker_wav_path)

    # Преобразование в моно, если аудио стерео
    if speaker_wav_tensor.dim() > 1 and speaker_wav_tensor.size(0) > 1:
        speaker_wav_tensor = speaker_wav_tensor.mean(dim=0, keepdim=True)

    # Применение денойзинга к аудио говорящего
    denoised_speaker_wav_tensor, denoised_speaker_sample_rate = denoise(
        speaker_wav_tensor.squeeze(), speaker_sample_rate, device
    )

    # Сохранение денойзенного аудио говорящего во временный файл
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_denoised_speaker_file:
        temp_denoised_speaker_path = temp_denoised_speaker_file.name
        torchaudio.save(
            temp_denoised_speaker_path,
            denoised_speaker_wav_tensor.unsqueeze(0).cpu(),
            denoised_speaker_sample_rate
        )

    # Генерация речи с помощью TTS и сохранение во временный файл
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_tts_output:
        temp_tts_output_path = temp_tts_output.name
        tts.tts_to_file(
            text=text,
            file_path=temp_tts_output_path,
            speed=speed,
            speaker_wav=temp_denoised_speaker_path,
            language=language_iso
        )

    # Загрузка сгенерированного аудио
    wav_tensor, sample_rate = torchaudio.load(temp_tts_output_path)

    # Преобразование в моно, если аудио стерео
    if wav_tensor.dim() > 1 and wav_tensor.size(0) > 1:
        wav_tensor = wav_tensor.mean(dim=0, keepdim=True)

    # Сохранение сгенерированного аудио во временный файл для voice cloning
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_vc_input_file:
        temp_vc_input_path = temp_vc_input_file.name
        torchaudio.save(temp_vc_input_path, wav_tensor.cpu(), sample_rate)

    # Инициализация модели voice conversion
    tts_conversion = TTS(
        model_name="voice_conversion_models/multilingual/vctk/freevc24",
        progress_bar=False
    )

    # Подготовка временного выходного файла
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_output_wav_file:
        temp_output_wav_path = temp_output_wav_file.name

    # Преобразование голоса
    tts_conversion.voice_conversion_to_file(
        temp_vc_input_path,
        target_wav=temp_denoised_speaker_path,
        file_path=temp_output_wav_path
    )

    # Чтение преобразованного аудио
    output_sample_rate, output_audio_data = read(temp_output_wav_path)

    # Удаление временных файлов
    os.remove(temp_denoised_speaker_path)
    os.remove(temp_tts_output_path)
    os.remove(temp_vc_input_path)
    os.remove(temp_output_wav_path)

    return (output_sample_rate, output_audio_data)

def get_language_code(selected_language):
    if selected_language in language_options:
        return language_options[selected_language]
    elif selected_language in other_language:
        return other_language[selected_language]
    else:
        return None

def process_speech(text, speaker_wav_path, selected_language, speed):
    language_code = get_language_code(selected_language)

    if language_code is None:
        raise ValueError("Выбранный язык не поддерживается.")

    if speaker_wav_path is None:
        error_message = "Пожалуйста, загрузите аудио файл говорящего."
        error = gr.Error(error_message, duration=5)
        raise error

    # Проверка длины аудио
    audio = AudioSegment.from_file(speaker_wav_path)
    duration = audio.duration_seconds
    if duration > 120:
        error_message = "Длина аудио превышает допустимый лимит в 2 минуты."
        error = gr.Error(error_message, duration=5)
        raise error

    if selected_language in other_language:
        output_audio_data = synthesize_and_convert_voice(text, language_code, speaker_wav_path, speed)
    else:
        output_audio_data = synthesize_speech(text, speaker_wav_path, language_code, speed)

    return output_audio_data 

def restart_program():
    python = sys.executable
    os.execl(python, python, *sys.argv)

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

    print(f"Текущая рабочая директория: {os.getcwd()}")
    print(f"Содержимое текущей директории: {os.listdir('.')}")
    print(f"Проверка наличия 'inference.py': {os.path.exists('inference.py')}")

    video_path = video  # Путь к видео или изображению
    audio_path = audio  # Путь к аудио

    print(f"Путь к видео: {video_path}")
    print(f"Путь к аудио: {audio_path}")

    output_dir = "outputs"
    os.makedirs(output_dir, exist_ok=True)

    output_path = os.path.join(output_dir, "output.mp4")
    print(f"Путь к выходному файлу: {output_path}")

    args = [
        "--checkpoint_path", f"checkpoints/{checkpoint}.pth",
        "--segmentation_path", "checkpoints/face_segmentation.pth",
        "--no_seg",
        # "--no_sr",
        "--face", video_path,
        "--audio", audio_path,
        "--outfile", output_path,
        "--resize_factor", str(resize_factor),
        "--face_det_batch_size", "4",
        "--wav2lip_batch_size", "64",
        "--fps", "30",
        "--pads", str(pad_top), str(pad_bottom), str(pad_left), str(pad_right)
    ]

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

    if save_as_video:
        args.append("--save_as_video")

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

    if not os.path.exists(output_path):
        print("Выходной файл не существует.")
        return "Не удалось создать выходное видео."

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

# Создание Gradio интерфейса с вкладками
with gr.Blocks() as app:
    gr.Markdown("# Voice Clone Union")

    with gr.Tabs():
        with gr.TabItem("Voice Clone"):
            # Интерфейс для голосового клонирования
            text_input = gr.Textbox(label="Введите текст для генерации", placeholder="Введите ваш текст здесь...")
            speaker_wav_input = gr.Audio(label="Загрузите аудио файла говорящего (WAV формат)", type="filepath")

            all_languages = list(language_options.keys()) + list(other_language.keys())
            language_input = gr.Dropdown(
                label="Язык",
                choices=all_languages,
                value="English (en)"
            )

            speed_input = gr.Slider(
                label="Скорость синтеза",
                minimum=0.1,
                maximum=10,
                step=0.1,
                value=1.0,
                info="Выберите скорость"
            )

            output_audio = gr.Audio(label="Сгенерированное аудио", type="filepath")

            with gr.Row():
                synthesize_button = gr.Button("Сгенерировать")
                gr.HTML("<div style='width:300px;'></div>")
                reload_button = gr.Button("Перезапустить")

            synthesize_button.click(
                fn=process_speech,
                inputs=[text_input, speaker_wav_input, language_input, speed_input],
                outputs=output_audio
            )

            reload_button.click(fn=restart_program, inputs=None, outputs=None)

        with gr.TabItem("Lipsync"):
            # Интерфейс для липсинка
            gr.Markdown("## Lipsync")
            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", visible=False)
                    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="Отступ снизу")
                    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="Отступ справа")
                    save_as_video = gr.Checkbox(label="Сохранять как видео", value=True)
                    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, save_as_video],
                outputs=result,
            )

    def launch_gradio():
        app.launch(
            # Вы можете добавить параметры запуска здесь, если необходимо
        )

if __name__ == "__main__":
    launch_gradio()