import os
import uuid
import asyncio
import subprocess
import shutil

import gradio as gr
import edge_tts
from deep_translator import GoogleTranslator
from faster_whisper import WhisperModel
import spaces


MAX_VIDEO_DURATION = 60

model = WhisperModel("small", device="cpu", compute_type="int8")


language_mapping = {
    "English": ("en", "en-US-EricNeural"),
    "Spanish": ("es", "es-ES-AlvaroNeural"),
    "French": ("fr", "fr-FR-HenriNeural"),
    "German": ("de", "de-DE-ConradNeural"),
    "Italian": ("it", "it-IT-DiegoNeural"),
    "Portuguese": ("pt", "pt-PT-DuarteNeural"),
    "Polish": ("pl", "pl-PL-MarekNeural"),
    "Turkish": ("tr", "tr-TR-AhmetNeural"),
    "Russian": ("ru", "ru-RU-DmitryNeural"),
}


def unique(ext):
    return f"{uuid.uuid4().hex}{ext}"


def transcribe(audio):

    segments, _ = model.transcribe(audio)

    text = ""

    for seg in segments:
        text += seg.text + " "

    return text.strip()


async def tts(text, voice, out):

    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(out)


@spaces.GPU
def process(video, target_language):

    video = video.name if hasattr(video, "name") else video

    try:

        uid = uuid.uuid4().hex[:6]

        # безопасное имя файла
        safe_video = f"{uid}_input.mp4"
        shutil.copy(video, safe_video)

        resized = f"{uid}_720.mp4"

        subprocess.run(
            f'ffmpeg -y -i "{safe_video}" -vf scale=-2:720 "{resized}"',
            shell=True,
            check=True
        )

        audio = f"{uid}.wav"

        subprocess.run(
            f'ffmpeg -y -i "{resized}" -vn -ac 1 -ar 16000 "{audio}"',
            shell=True,
            check=True
        )

        text = transcribe(audio)

        lang, voice = language_mapping[target_language]

        translated = GoogleTranslator(
            source="auto",
            target=lang
        ).translate(text)

        speech = f"{uid}_tts.wav"

        asyncio.run(tts(translated, voice, speech))

        out = f"{uid}_dubbed.mp4"

        subprocess.run(
            f'ffmpeg -y -i "{resized}" -i "{speech}" -c:v copy -c:a aac -map 0:v -map 1:a "{out}"',
            shell=True,
            check=True
        )

        # cleanup
        for f in [safe_video, resized, audio, speech]:
            if os.path.exists(f):
                os.remove(f)

        return out, ""

    except Exception as e:

        return None, str(e)


with gr.Blocks() as demo:

    gr.Markdown("# AI Video Dubbing")

    video = gr.File(file_types=["video"])

    lang = gr.Dropdown(
        list(language_mapping.keys()),
        value="Spanish"
    )

    run = gr.Button("Process")

    out = gr.Video()
    status = gr.Textbox()

    run.click(
        process,
        inputs=[video, lang],
        outputs=[out, status]
    )

demo.queue()

demo.launch(server_name="0.0.0.0", server_port=7860)