import os
import cv2
import tempfile
import time
import gradio as gr
import insightface
import onnxruntime
from gfpgan import GFPGANer


class Predictor:
    def __init__(self):
        self.setup()

    def setup(self):
        model_dir = "models"

        print("Loading models...")
        self.face_swapper = insightface.model_zoo.get_model(
            f"{model_dir}/inswapper_128.onnx",
            providers=onnxruntime.get_available_providers()
        )

        self.face_enhancer = GFPGANer(
            model_path=f"{model_dir}/GFPGANv1.4.pth",
            upscale=1,
        )

        self.face_analyser = insightface.app.FaceAnalysis(name="buffalo_l")
        self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))

    def get_face(self, img_data):
        """Return the largest detected face"""
        faces = self.face_analyser.get(img_data)
        if len(faces) == 0:
            return None
        return max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))

    def predict(self, target_image, swap_image):
        """Face swap function"""
        try:
            frame = cv2.imread(target_image)
            source = cv2.imread(swap_image)

            if frame is None or source is None:
                return None, "❌ Could not read one of the images."

            target_face = self.get_face(frame)
            source_face = self.get_face(source)

            if target_face is None or source_face is None:
                return None, "❌ Could not detect a clear face in one or both images."

            # Perform face swap
            result = self.face_swapper.get(frame, target_face, source_face, paste_back=True)

            # Enhance with GFPGAN
            _, _, result = self.face_enhancer.enhance(result, paste_back=True)

            # Save result
            out_path = tempfile.mkdtemp() + f"/swapped_{int(time.time())}.jpg"
            cv2.imwrite(out_path, result)

            return out_path, "✅ Face swap completed successfully!"

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


# Load the predictor once at startup
predictor = Predictor()

# Clean UI without examples
with gr.Blocks(title="Swap Face Model", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 😻 Face Swap Model")
    gr.Markdown("Upload a **target image** and a **face to swap in**. Works best with clear, well-lit frontal faces.")

    with gr.Row():
        with gr.Column():
            target_input = gr.Image(
                type="filepath", 
                label="Target Image (the photo you want to change)", 
                height=450
            )
        with gr.Column():
            swap_input = gr.Image(
                type="filepath", 
                label="Swap Face Image (the face you want to use)", 
                height=450
            )

    with gr.Row():
        swap_btn = gr.Button("🔄 Swap Faces", variant="primary", size="large")

    with gr.Row():
        output_image = gr.Image(
            type="filepath", 
            label="Result", 
            height=500
        )

    status = gr.Textbox(label="Status", interactive=False)

    # Button action
    swap_btn.click(
        fn=predictor.predict,
        inputs=[target_input, swap_input],
        outputs=[output_image, status]
    )

    gr.Markdown("**Tips:**\n• Use high-resolution, well-lit photos\n• Frontal or near-frontal faces work best\n• Results improve with similar lighting and angles.")

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )