import gradio as gr
import roop.globals
import roop.core
import os
import shutil

# ---------------------------
# GLOBAL CONFIG (RUNS ONCE)
# ---------------------------
roop.globals.headless = True
roop.globals.execution_providers = ['CPUExecutionProvider']
roop.globals.execution_threads = 4  # Faster CPU usage
roop.globals.frame_processors = ['face_swapper']
roop.globals.keep_fps = True
roop.globals.keep_frames = False
roop.globals.skip_audio = True
roop.globals.many_faces = False
roop.globals.reference_face_position = 0
roop.globals.reference_frame_number = 0
roop.globals.similar_face_distance = 0.85
roop.globals.temp_frame_format = 'png'
roop.globals.temp_frame_quality = 95
roop.globals.output_video_encoder = 'libx264'
roop.globals.output_video_quality = 35
roop.globals.max_memory = 8  # limit memory for stability

print("🔥 Roop initialized (CPU optimized)")

def get_valid_image_path(file_obj, suffix=""):
    if file_obj is None:
        return None
    
    # Check if API sent string, dict, or file object
    if isinstance(file_obj, str):
        path = file_obj
    elif isinstance(file_obj, dict) and 'path' in file_obj:
        path = file_obj['path']
    else:
        path = getattr(file_obj, 'name', str(file_obj))
        
    # Agari extension missing hai API ki wajah se, usko .png dedo! (Bug Fixed Here)
    if not any(path.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp']):
        new_path = path + suffix + ".png"
        shutil.copy(path, new_path)
        return new_path
    return path

# ---------------------------
# MAIN FUNCTION
# ---------------------------
def swap_face(source_file, target_file, enable_enhancer):
    if source_file is None or target_file is None:
        return None

    # Ab humein farq nahi padta kon kaisa image bhej raha hai!
    source_path = get_valid_image_path(source_file, "_src")
    target_path = get_valid_image_path(target_file, "_tgt")

    output_path = os.path.join(os.getcwd(), "output.png")

    # Clean old output
    if os.path.exists(output_path):
        os.remove(output_path)

    print(f"📥 Source: {source_path}")
    print(f"🎯 Target: {target_path}")
    print(f"⚙️ Enhancer: {enable_enhancer}")

    # Set paths
    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    roop.globals.output_path = output_path

    # Processors config
    processors = ['face_swapper']
    if enable_enhancer:
        processors.append('face_enhancer')
    roop.globals.frame_processors = processors

    try:
        roop.core.start()

        if os.path.exists(output_path):
            print("✅ Done!")
            return output_path
        else:
            print("❌ Output not created")
            return None

    except Exception as e:
        print(f"❌ Error: {e}")
        return None


# ---------------------------
# UI
# ---------------------------
with gr.Blocks(title="AI Guruji Face Swap") as demo:

    gr.HTML("""
    <div style="background-color:#0b0f19;padding:20px;text-align:center;border-radius:10px;margin-bottom:20px;">
        <h1 style="color:white;">🔥 AI Guruji Face Swap</h1>
        <p style="color:white;">⚡ Fast CPU Optimized Version</p>
    </div>
    """)

    gr.Markdown("### Upload source face and target image")
    gr.Markdown("⚠️ CPU Mode → Processing may take 10–30 seconds")

    with gr.Row():
        with gr.Column():
            source_input = gr.Image(type="filepath", label="Source Face")
            target_input = gr.Image(type="filepath", label="Target Image")
            enable_enhancer = gr.Checkbox(label="Enable Face Enhancer (Slow)", value=False)
            submit_btn = gr.Button("🚀 Swap Face", variant="primary")

        with gr.Column():
            result_image = gr.Image(label="Result", type="filepath")

    submit_btn.click(
        fn=swap_face,
        inputs=[source_input, target_input, enable_enhancer],
        outputs=[result_image]
    )

# ---------------------------
# LAUNCH
# ---------------------------
if __name__ == "__main__":
    demo.queue(max_size=10)  # Prevent overload
    demo.launch()