import os
import glob
import spaces
from natsort import natsorted
import gradio as gr

from inference_util import init_model, infenrece
from attributtes_utils import input_pose, input_emotion, input_blink

model = init_model()

@spaces.GPU
def process(
    video_source, input_vid, upload_vid,
    audio_source, audio_path, upload_audio,
    pose_select, emotion_select, blink_select
):
    pose = input_pose(pose_select)
    emotion = input_emotion(emotion_select)
    blink = input_blink(blink_select)

    # Decide which video to use
    if video_source == "From list":
        face_path = os.path.join("./assets/videos/", input_vid)
    else:
        # upload_vid is a temp file path from gr.Video
        face_path = upload_vid

    # Decide which audio to use
    if audio_source == "From list":
        audio_file = os.path.join("./assets/audios/", audio_path)
    else:
        # upload_audio is a temp file path from gr.Audio(type="filepath")
        audio_file = upload_audio

    print("face_path: ", face_path)
    print("audio_file: ", audio_file)

    result = infenrece(
        model,
        face_path,
        audio_file,
        pose,
        emotion,
        blink,
    )

    print("result: ", result)
    print("finished !")

    return result  # , gr.Group.update(visible=True)


available_videos = natsorted(glob.glob("./assets/videos/*.mp4"))
available_videos = [os.path.basename(x) for x in available_videos]

# prepare audio
for video in available_videos:
    audio = video.replace(".mp4", ".wav")
    if not os.path.exists(os.path.join("./assets/audios/", audio)):
        os.system(
            f"ffmpeg -y -loglevel error -i ./assets/videos/{video} "
            f"-vn -acodec pcm_s16le -ar 16000 -ac 1 ./assets/audios/{audio}"
        )

available_audios = natsorted(glob.glob("./assets/audios/*.wav"))
available_audios = [os.path.basename(x) for x in available_audios]


def toggle_video_source(src):
    return (
        gr.update(visible=(src == "From list")),
        gr.update(visible=(src == "Upload")),
    )


def toggle_audio_source(src):
    return (
        gr.update(visible=(src == "From list")),
        gr.update(visible=(src == "Upload")),
    )


with gr.Blocks() as demo:
    gr.HTML(
        """
            <h1 style="text-align: center; font-size: 40px; font-family: 'Times New Roman', Times, serif;">
                Free-View Expressive Talking Head Video Editing
            </h1>
            <p style="text-align: center; font-size: 20px; font-family: 'Times New Roman', Times, serif;">
                <a style="text-align: center; display:inline-block"
                    href="https://sky24h.github.io/websites/icassp2023_free-view_video-editing">
                    <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/paper-page-sm.svg#center"
                    alt="Project Page">
                </a>
                <a style="text-align: center; display:inline-block" href="https://huggingface.co/spaces/sky24h/Free-View_Expressive_Talking_Head_Video_Editing?duplicate=true">
                    <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm.svg#center" alt="Duplicate Space">
                </a>
            </p>
            <p style="text-align: center; font-size: 16px; font-family: 'Times New Roman', Times, serif;">
                If you wish to use your custom input files, please duplicate this space or clone it to your local environment.</p>
            <p style="text-align: center; font-size: 16px; font-family: 'Times New Roman', Times, serif;">
                Alternatively, you can check our official <a href="https://github.com/sky24h/Free-View_Expressive_Talking_Head_Video_Editing">repository</a> on GitHub.
            </p>
            """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column():
                # Preview area
                video_preview = gr.Video(
                    label="Video Preview",
                    elem_id="video-preview",
                    value="./assets/videos/sample1.mp4",
                )
                audio_preview = gr.Audio(
                    label="Audio Preview",
                    elem_id="audio-preview",
                    value="./assets/audios/sample2.wav",
                )

                # Video source choice
                video_source = gr.Radio(
                    ["From list", "Upload"],
                    label="Video source",
                    value="From list",
                )
                video_input = gr.Dropdown(
                    available_videos, label="Input Video", value="sample1.mp4"
                )
                upload_vid = gr.Video(
                    label="Upload Video",
                    visible=False,
                )

                # Audio source choice
                audio_source = gr.Radio(
                    ["From list", "Upload"],
                    label="Audio source",
                    value="From list",
                )
                audio_input = gr.Dropdown(
                    available_audios, label="Input Audio", value="sample2.wav"
                )
                upload_audio = gr.Audio(
                    type="filepath",
                    label="Upload Audio",
                    visible=False,
                )

                pose_select = gr.Radio(
                    ["front", "left_right_shaking"],
                    label="Pose",
                    value="front",
                )
                emotion_select = gr.Radio(
                    ["neutral", "happy", "angry", "surprised"],
                    label="Emotion",
                    value="neutral",
                )
                blink_select = gr.Radio(
                    ["yes", "no"],
                    label="Blink",
                    value="yes",
                )

            with gr.Column():
                video_out = gr.Video(
                    label="Video Output",
                    elem_id="video-output",
                    height=360,
                )
                submit_btn = gr.Button("Generate video")

        inputs = [
            video_source,
            video_input,
            upload_vid,
            audio_source,
            audio_input,
            upload_audio,
            pose_select,
            emotion_select,
            blink_select,
        ]
        outputs = [video_out]

        video_preview_output = [video_preview]
        audio_preview_output = [audio_preview]

    # Update previews when selecting from list
    video_input.select(
        lambda x: "./assets/videos/" + x,
        video_input,
        video_preview_output,
    )
    audio_input.select(
        lambda x: "./assets/audios/" + x,
        audio_input,
        audio_preview_output,
    )

    # Toggle list vs upload widgets
    video_source.change(
        toggle_video_source,
        inputs=video_source,
        outputs=[video_input, upload_vid],
    )
    audio_source.change(
        toggle_audio_source,
        inputs=audio_source,
        outputs=[audio_input, upload_audio],
    )

    submit_btn.click(process, inputs, outputs)

demo.queue(max_size=10).launch()