import gradio as gr
import gradio
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, frame_processor_checkbox, face_swapper_model_dropdown, face_enhancer_model_dropdown, frame_enhancer_model_dropdown, face_debugger_items_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown, face_mask_type_checkbox, *video_args = args
    if not os.path.exists(source):
        return "Source file does not exist"
    if not os.path.exists(target):
        return "Target file does not exist"
    remove_old_directories("./output", num_minutes=60)
    filename = os.path.basename(target)
    os.makedirs(f"./output/{unique_id}",exist_ok=True)
    output = f"./output/{unique_id}/{filename}"
    frame_processor = frame_processor_checkbox
    selected_frame_processors = ' '.join(frame_processor)

    face_analyser_direction = face_analyser_direction_dropdown
    face_recognition = face_analyser_age_dropdown
    face_analyser_gender = face_analyser_gender_dropdown

    cmd = (
        f"python run.py --execution-providers cpu --headless -s {source} -t {target} -o {output} "
        f"--frame-processors {selected_frame_processors} "
        f"--face-analyser-order {face_analyser_direction} "
    )
    if "face_swapper" in selected_frame_processors:
        cmd += f"--face-swapper-model {face_swapper_model_dropdown} "
    if "face_enhancer" in selected_frame_processors:
        cmd += f"--face-enhancer-model {face_enhancer_model_dropdown} "
    if "frame_enhancer" in selected_frame_processors:
        cmd += f"--frame-enhancer-model {frame_enhancer_model_dropdown} "
    if "face_debugger" in selected_frame_processors:
        cmd += f"--face-debugger-items {face_debugger_items_checkbox} "
    if face_recognition != 'none':
        cmd += f"--face-selector-mode {face_recognition} "
    if face_analyser_gender != 'none':
        cmd += f"--face-analyser-gender {face_analyser_gender} "

    if video_args:
        skip_audio, keep_fps, keep_temp = video_args
        if skip_audio:
            cmd += "--skip-audio "
        if keep_fps:
            cmd += "--keep-fps "
        if keep_temp:
            cmd += "--keep-temp "

    try:
        print("Started...", cmd)
        output_text = sp.run(cmd, shell=True, capture_output=True, text=True).stdout
        print(output_text)
        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):
            print("Trying to delete ")
            for filename in os.listdir(output_path):
                file_path = os.path.join(output_path, filename)
                if os.path.isfile(file_path):
                    os.remove(file_path)
                    print(f"Output files in {output_path} are deleted")
                    return "Output files for unique_id deleted"
                else:
                    print(f"Output files in {output_path} does not exist")
                    return "Output directory for (output_path} 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 # Convert to minutes

            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)
                    pass

def get_css() -> str:
    fixes_css = """
:root:root:root button:not([class])
{
    border-radius: 0.375rem;
    float: left;
    overflow: hidden;
    width: 100%;

}
"""
    overrides_css = """
:root:root:root input[type="number"]
{
    max-width: 6rem;
}

:root:root:root [type="checkbox"],
:root:root:root [type="radio"]
{
    border-radius: 50%;
    height: 1.125rem;
    width: 1.125rem;
}

:root:root:root input[type="range"]
{
    height: 0.5rem;
}

:root:root:root input[type="range"]::-moz-range-thumb,
:root:root:root input[type="range"]::-webkit-slider-thumb
{
    background: var(--neutral-300);
    border: unset;
    border-radius: 50%;
    height: 1.125rem;
    width: 1.125rem;
}

:root:root:root input[type="range"]::-webkit-slider-thumb
{
    margin-top: 0.375rem;
}

:root:root:root .grid-wrap.fixed-height
{
    min-height: unset;
}

:root:root:root .grid-container
{
    grid-auto-rows: minmax(5em, 1fr);
    grid-template-columns: repeat(var(--grid-cols), minmax(5em, 1fr));
    grid-template-rows: repeat(var(--grid-rows), minmax(5em, 1fr));
}
"""
    return fixes_css + overrides_css

def get_theme() -> gradio.Theme:
    return gradio.themes.Base(
        primary_hue = gradio.themes.colors.red,
        secondary_hue = gradio.themes.colors.neutral,
        font = gradio.themes.GoogleFont('Open Sans')
    ).set(
        background_fill_primary = '*neutral_100',
        block_background_fill = 'white',
        block_border_width = '0',
        block_label_background_fill = '*primary_100',
        block_label_background_fill_dark = '*primary_600',
        block_label_border_width = 'none',
        block_label_margin = '0.5rem',
        block_label_radius = '*radius_md',
        block_label_text_color = '*primary_500',
        block_label_text_color_dark = 'white',
        block_label_text_weight = '600',
        block_title_background_fill = '*primary_100',
        block_title_background_fill_dark = '*primary_600',
        block_title_padding = '*block_label_padding',
        block_title_radius = '*block_label_radius',
        block_title_text_color = '*primary_500',
        block_title_text_size = '*text_sm',
        block_title_text_weight = '600',
        block_padding = '0.5rem',
        border_color_primary = 'transparent',
        border_color_primary_dark = 'transparent',
        #button_large_padding = '2rem 0.5rem',
        #button_large_text_weight = 'normal',
        #button_primary_background_fill = '*primary_500',
        #button_primary_text_color = 'white',
        #button_secondary_background_fill = 'white',
        #button_secondary_border_color = 'transparent',
        #button_secondary_border_color_dark = 'transparent',
        #button_secondary_border_color_hover = 'transparent',
        #button_secondary_border_color_hover_dark = 'transparent',
        #button_secondary_text_color = '*neutral_800',
        #button_small_padding = '0.75rem',
        checkbox_background_color = '*neutral_200',
        checkbox_background_color_selected = '*primary_600',
        checkbox_background_color_selected_dark = '*primary_700',
        checkbox_border_color_focus = '*primary_500',
        checkbox_border_color_focus_dark = '*primary_600',
        checkbox_border_color_selected = '*primary_600',
        checkbox_border_color_selected_dark = '*primary_700',
        checkbox_label_background_fill = '*neutral_50',
        checkbox_label_background_fill_hover = '*neutral_50',
        checkbox_label_background_fill_selected = '*primary_500',
        checkbox_label_background_fill_selected_dark = '*primary_600',
        checkbox_label_text_color_selected = 'white',
        input_background_fill = '*neutral_50',
        shadow_drop = 'none',
        slider_color = '*primary_500',
        slider_color_dark = '*primary_600'
    )


with gr.Blocks(theme=get_theme(), css=get_css(), title="DeepFakeAI 2.0.1") as ui:
    with gr.Row():
        gr.HTML('<center><a href="https://codegenius.me">DeepFakeAI 2.0.1</a></center>')

    with gr.Row():
        with gr.Column(scale=3):
            frame_processor_checkbox = gr.CheckboxGroup(
                choices = ['face_swapper', 'face_enhancer', 'frame_enhancer', 'face_debugger'],
                label = 'FRAME PROCESSORS',
                value = ['face_swapper']  # Default value
            )
    with gr.Row():
        with gr.Column(scale=3):
            face_swapper_model_dropdown = gr.Dropdown(
                choices = ['blendswap_256', 'inswapper_128', 'inswapper_128_fp16', 'simswap_256', 'simswap_512_unofficial'],
                label = 'FACE SWAPPER MODEL',
                value = 'inswapper_128',
                interactive=True,
                visible = True #'face_swapper' in frame_processor_checkbox.value
            )
            face_enhancer_model_dropdown = gr.Dropdown(
                choices = ['codeformer', 'gfpgan_1.2', 'gfpgan_1.3', 'gfpgan_1.4', 'gpen_bfr_256', 'gpen_bfr_512', 'restoreformer'],
                label = 'FACE ENHANCER MODEL',
                value = 'gfpgan_1.2',
                visible = True, #'face_enhancer' in frame_processor_checkbox.value
                interactive=True
            )
            frame_enhancer_model_dropdown = gr.Dropdown(
                choices = ['real_esrgan_x2plus', 'real_esrgan_x4plus', 'real_esrnet_x4plus'],
                label = 'FRAME ENHANCER MODEL',
                value = 'real_esrgan_x4plus',
                interactive=True,
                visible = True #'frame-enhancer' in frame_processor_checkbox.value
            )
            face_debugger_items_checkbox = gr.CheckboxGroup(
                choices = ['face-mask', 'bbox', 'kps', 'score'],
                label = 'FACE DEBUGGER ITEMS',
                value = 'face-mask',
                interactive=True,
                visible = True
            )
    with gr.Row():
        with gr.Column(scale=3):
            face_analyser_direction_dropdown = gr.Dropdown(
                label = 'FACE ANALYSER DIRECTION',
                choices = ['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'],
                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'
            )
    with gr.Row():
        with gr.Column(scale=3):
            face_mask_type_checkbox = gr.CheckboxGroup(
                choices = ['box', 'occlusion', 'region'],
                label = 'FACE MASK TYPES',
                value = 'box',
                interactive=True            )
            unique_id = gr.Textbox(value=str(uuid.uuid4()), visible=False)
    with gr.Tab("Image:               "):
        source_image = gr.Image(type="filepath", label="SOURCE IMAGE", sources=["upload", "webcam", "clipboard"])
        target_image = gr.Image(type="filepath", label="TARGET IMAGE", sources=["upload", "webcam", "clipboard"])
        image_button = gr.Button("START")
        clear_button = gr.ClearButton(value="CLEAR")
        image_output = gr.Image(label="OUTPUT")
        clear_button.add(image_output)

        image_button.click(
            run,
            inputs=[source_image, target_image, unique_id, frame_processor_checkbox, face_swapper_model_dropdown, face_enhancer_model_dropdown, frame_enhancer_model_dropdown, face_debugger_items_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown, face_mask_type_checkbox],
            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", sources=["upload", "webcam", "clipboard"])
        target_video = gr.Video(label="TARGET VIDEO", sources=["upload"])
        with gr.Row():
            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")
        clear_video_button = gr.ClearButton(value="CLEAR")
        video_output = gr.Video(label="OUTPUT")
        clear_video_button.add(video_output)
        video_button.click(
            run,
            inputs=[source_image_video, target_video, unique_id, frame_processor_checkbox, face_swapper_model_dropdown, face_enhancer_model_dropdown, frame_enhancer_model_dropdown, face_debugger_items_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown, face_mask_type_checkbox, skip_audio, keep_fps, keep_temp],
            outputs=video_output
        )
        clear_video_button.click(fn=clear_output, inputs=unique_id)

ui.launch(debug=True)