import gradio as gr
import gradio
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)
    
    cmd = f"python run.py --execution-providers cpu -s {source} -t {target} -o {output} --frame-processors {selected_frame_processors}"
    
    if len(rest_args) > 1:
        skip_audio = rest_args[1]
        keep_fps = rest_args[2]
        keep_temp = rest_args[3]
        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 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():
        frame_processor_checkbox = gr.CheckboxGroup(
            choices = ['face_swapper', 'face_enhancer', 'frame_enhancer'],
            label = 'FRAME PROCESSORS',
            value = ['face_swapper']  # Default value
        )

    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()
        
        image_button.click(
            run, 
            inputs=[source_image, target_image, frame_processor_checkbox],
            outputs=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")
        video_output = gr.Video()
        video_button.click(
            run, 
            inputs=[source_image_video, target_video, frame_processor_checkbox, skip_audio, keep_fps, keep_temp],
            outputs=video_output
        )

    # Additional UI components ...
    with gr.Box():
            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'
            )

ui.launch(debug=True)