import gradio as gr
import subprocess as sp
import os
import uuid
import time
import shutil

os.makedirs("./output", exist_ok=True)


def run(*args):
    source, target, unique_id, *rest_args = args

    if not source or not os.path.exists(source):
        return "Source file does not exist"

    if not target or not os.path.exists(target):
        return "Target file does not exist"

    remove_old_directories("./output", num_minutes=60)

    filename = os.path.basename(target)
    output_dir = f"./output/{unique_id}"

    os.makedirs(output_dir, exist_ok=True)

    output = f"{output_dir}/{filename}"

    frame_processor = rest_args[0]
    selected_frame_processors = " ".join(frame_processor)

    face_analyser_direction = rest_args[1]
    face_recognition = rest_args[2]
    face_analyser_gender = rest_args[3]

    cmd = (
        f'python run.py '
        f'--execution-providers cpu '
        f'-s "{source}" '
        f'-t "{target}" '
        f'-o "{output}" '
        f'--frame-processors {selected_frame_processors} '
        f'--face-analyser-direction {face_analyser_direction} '
    )

    if face_recognition != "none":
        cmd += f"--face-recognition {face_recognition} "

    if face_analyser_gender != "none":
        cmd += f"--face-analyser-gender {face_analyser_gender} "

    if len(rest_args) > 4:
        skip_audio = rest_args[4]
        keep_fps = rest_args[5]
        keep_temp = rest_args[6]

        if skip_audio:
            cmd += "--skip-audio "

        if keep_fps:
            cmd += "--keep-fps "

        if keep_temp:
            cmd += "--keep-temp "

    try:
        print("Started...", cmd)

        result = sp.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True
        )

        print(result.stdout)

        if result.stderr:
            print(result.stderr)

        return output

    except Exception as e:
        return f"An error occurred: {str(e)}"


def clear_output(unique_id):
    try:
        output_path = f"./output/{unique_id}"

        if os.path.exists(output_path):
            shutil.rmtree(output_path)
            print(f"Deleted {output_path}")
            return "Output deleted"

        return "Output directory does not exist"

    except Exception as e:
        return f"An error occurred: {str(e)}"


def remove_old_directories(directory, num_minutes=60):
    now = time.time()

    for r, d, f in os.walk(directory):
        for dir_name in d:
            dir_path = os.path.join(r, dir_name)

            timestamp = os.path.getmtime(dir_path)
            age_minutes = (now - timestamp) / 60

            if age_minutes >= num_minutes:
                try:
                    print("Removing", dir_path)
                    shutil.rmtree(dir_path)
                    print("Directory removed:", dir_path)

                except Exception as e:
                    print(e)


def get_theme():
    return gr.themes.Soft(
        primary_hue=gr.themes.colors.red,
        secondary_hue=gr.themes.colors.gray,
        font=gr.themes.GoogleFont("Inter")
    )


with gr.Blocks(
    theme=get_theme(),
    title="DeepFakeAI 1.0.0"
) as ui:

    gr.HTML(
        '<center><h1>DeepFakeAI 1.0.1</h1></center>'
    )

    frame_processor_checkbox = gr.CheckboxGroup(
        choices=[
            "face_swapper",
            "face_enhancer",
            "frame_enhancer"
        ],
        label="FRAME PROCESSORS",
        value=["face_swapper"]
    )

    face_analyser_direction_dropdown = gr.Dropdown(
        label="FACE ANALYSER DIRECTION",
        choices=[
            "left-right",
            "right-left",
            "top-bottom",
            "bottom-top",
            "small-large",
            "large-small"
        ],
        value="left-right"
    )

    face_analyser_age_dropdown = gr.Dropdown(
        label="FACE RECOGNITION",
        choices=["none", "reference", "many"],
        value="reference"
    )

    face_analyser_gender_dropdown = gr.Dropdown(
        label="FACE ANALYSER GENDER",
        choices=["none", "male", "female"],
        value="none"
    )

    unique_id = gr.Textbox(
        value=str(uuid.uuid4()),
        visible=False
    )

    with gr.Tab("Image"):

        source_image = gr.Image(
            type="filepath",
            label="SOURCE IMAGE"
        )

        target_image = gr.Image(
            type="filepath",
            label="TARGET IMAGE"
        )

        image_button = gr.Button("START")

        image_output = gr.Image(
            label="OUTPUT"
        )

        clear_button = gr.Button("CLEAR")

        image_button.click(
            fn=run,
            inputs=[
                source_image,
                target_image,
                unique_id,
                frame_processor_checkbox,
                face_analyser_direction_dropdown,
                face_analyser_age_dropdown,
                face_analyser_gender_dropdown
            ],
            outputs=image_output
        )

        clear_button.click(
            fn=clear_output,
            inputs=unique_id
        )

    with gr.Tab("Video"):

        source_image_video = gr.Image(
            type="filepath",
            label="SOURCE IMAGE"
        )

        target_video = gr.Video(
            label="TARGET VIDEO"
        )

        with gr.Group():

            skip_audio = gr.Checkbox(
                label="SKIP AUDIO"
            )

            keep_fps = gr.Checkbox(
                label="KEEP FPS"
            )

            keep_temp = gr.Checkbox(
                label="KEEP TEMP"
            )

        video_button = gr.Button("START")

        video_output = gr.Video(
            label="OUTPUT"
        )

        clear_video_button = gr.Button("CLEAR")

        video_button.click(
            fn=run,
            inputs=[
                source_image_video,
                target_video,
                unique_id,
                frame_processor_checkbox,
                face_analyser_direction_dropdown,
                face_analyser_age_dropdown,
                face_analyser_gender_dropdown,
                skip_audio,
                keep_fps,
                keep_temp
            ],
            outputs=video_output
        )

        clear_video_button.click(
            fn=clear_output,
            inputs=unique_id
        )

ui.queue()

ui.launch(
    server_name="0.0.0.0",
    server_port=7860,
    debug=True
)