import gradio as gr
import subprocess as sp
import os

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

def run(*args):
    source, target, *rest_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"

    filename = os.path.basename(target)
    output = f"./output/{filename}"
    frame_processor = rest_args[0]
    selected_frame_processors = ' '.join(frame_processor)
    
    face_analyser_direction = rest_args[1]
    face_analyser_age = rest_args[2]
    face_analyser_gender = rest_args[3]
    
    cmd = (
        f"python run.py --execution-providers cuda -s {source} -t {target} -o {output} "
        f"--frame-processors {selected_frame_processors} "
        f"--face-analyser-direction {face_analyser_direction} "
        f"--face-analyser-age {face_analyser_age} "
        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)
        sp.run(cmd, shell=True, capture_output=True, text=True).stdout
        return output
    except Exception as e:
        return f"An error occurred: {str(e)}"
def clear_output(output):
    if os.path.exists(output):
        os.remove(output)
        return "Output file deleted"
    else:
        return "Output file does not exist"


def get_theme() -> gr.Theme:
    return gr.themes.Soft(
        primary_hue = gr.themes.colors.red,
        secondary_hue = gr.themes.colors.gray,
        font = gr.themes.GoogleFont('Inter')
    ).set(
        background_fill_primary = '*neutral_50',
        block_label_text_size = '*text_sm',
        block_title_text_size = '*text_sm'
    )

with gr.Blocks(theme=get_theme(), title="DeepFakeAI 1.0.0") as ui:
    with gr.Box():
        gr.HTML('<center><a href="codegenius.me">DeepFakeAI 1.0.0</a></center>')

    with gr.Box():
        with gr.Column(scale=3):
           frame_processor_checkbox = gr.CheckboxGroup(
            choices = ['face_swapper', 'face_enhancer', 'frame_enhancer'],
            label = 'FRAME PROCESSORS',
            value = ['face_swapper']  # Default value
        )
        
    with gr.Box():
        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'],
                value = 'left-right'
            )
            face_analyser_age_dropdown = gr.Dropdown(
                label = 'FACE ANALYSER AGE',
                choices = ['none'] + ['child', 'young', 'adult', 'old'],
                value = 'none'
            )
            face_analyser_gender_dropdown = gr.Dropdown(
                label = 'FACE ANALYSER GENDER',
                choices = ['none'] + ['male', 'female'],
                value = 'none'
            )
    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")
        clear_button = gr.Button("CLEAR")
        image_output = gr.Image(label="OUTPUT")
        
        image_button.click(
            run, 
            inputs=[source_image, target_image, frame_processor_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown],
            outputs=image_output
        )
        clear_button.click(clear_output, inputs=[image_output])

    with gr.Tab("Video:               "):
        source_image_video = gr.Image(type="filepath", label="SOURCE IMAGE")
        target_video = gr.Video(label="TARGET VIDEO")
        with gr.Box():
            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.Button("CLEAR")
        video_output = gr.Video(label="OUTPUT")
        video_button.click(
            run, 
            inputs=[source_image_video, target_video, 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(clear_output, inputs=[video_output])

ui.launch(debug=True)