import os
import sys

# ==========================================================
# 🚨 STEP 1: FORCE BYPASS THE DEPRECATED 'HfFolder' HUB BUG
# ==========================================================
class DummyHfFolder:
    @staticmethod
    def get_token(): return None
    @staticmethod
    def save_token(token): pass
    @staticmethod
    def delete_token(): pass

try:
    import huggingface_hub
    sys.modules['huggingface_hub'].HfFolder = DummyHfFolder
except Exception:
    pass

import cv2
import tempfile
import time
import gradio as gr
import insightface
import onnxruntime
import spaces  # <--- Essential for dynamic ZeroGPU routing
from gfpgan import GFPGANer


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

    def setup(self):
        model_dir = "models"
        os.makedirs(model_dir, exist_ok=True)

        # Download targeted checkpoint weights if missing from container storage
        if not os.path.exists(f'{model_dir}/GFPGANv1.4.pth'):
            print("Downloading GFPGAN model...")
            os.system(f'wget https://github.com -O {model_dir}/GFPGANv1.4.pth')
        if not os.path.exists(f'{model_dir}/inswapper_128.onnx'):
            print("Downloading Inswapper model...")
            os.system(f'wget https://huggingface.co -O {model_dir}/inswapper_128.onnx')

        print("Loading models...")
        # ZeroGPU dynamically coordinates inference tasks; initialize on safe CPU runtime execution boundaries
        self.face_swapper = insightface.model_zoo.get_model(
            f"{model_dir}/inswapper_128.onnx",
            providers=['CPUExecutionProvider']
        )

        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 spatial geometry boundary"""
        faces = self.face_analyser.get(img_data)
        if len(faces) == 0:
            return None
        return max(faces, key=lambda x: (x.bbox - x.bbox) * (x.bbox - x.bbox))

    @spaces.GPU # <--- STEP 2: Borrow GPU slices dynamically during face synthesis
    def predict(self, target_image, swap_image):
        """Face swap core execution script layer"""
        try:
            if target_image is None or swap_image is None:
                return None, "❌ Please upload both images before processing."

            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."

            # Map coordinates to the visual array inside the ZeroGPU tracking space
            result = self.face_swapper.get(frame, target_face, source_face, paste_back=True)

            # Reconstruct high fidelity resolution metrics using GFPGAN
            _, _, result = self.face_enhancer.enhance(result, paste_back=True)

            # Export modified matrix configuration data to disk
            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 global instances
predictor = Predictor()

# ==========================================================
# 🚨 STEP 3: CONSTRUCT MODERN VISUAL INTERFACE
# ==========================================================
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)

    # Event binding mapping coordinates
    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__":
    # Standard configuration mapping proxy values cleanly across internal server arrays
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )