# -*- coding: UTF-8 -*-
import gradio as gr
import cv2
import numpy as np
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from insightface.app import FaceAnalysis
from insightface.model_zoo.face_swapper import FaceSwapper
from pathlib import Path
import requests
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI
app = FastAPI(title="Face Swap API", description="Gradio-based Face Swap App")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# Health Check Endpoint
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# Download Model Function
MODEL_PATH = Path("models/inswapper_128.onnx")
MODEL_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"

def download_model():
    if not MODEL_PATH.exists():
        logger.info("Downloading model...")
        os.makedirs("models", exist_ok=True)
        response = requests.get(MODEL_URL, stream=True, timeout=30)
        response.raise_for_status()
        with open(MODEL_PATH, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        logger.info("Model downloaded successfully.")

download_model()

# Initialize Face Analyzer & Swapper
face_analyzer = FaceAnalysis(name="buffalo_l")
face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
face_swapper = FaceSwapper(MODEL_PATH)

# Face Swap Function
def swap_faces(source_img, target_img):
    source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)

    source_faces = face_analyzer.get(source_img)
    target_faces = face_analyzer.get(target_img)

    if not source_faces or not target_faces:
        return "Error: No faces detected in one or both images."
    if len(source_faces) > 1 or len(target_faces) > 1:
        return "Error: Multiple faces detected; only one face per image is supported."

    result_img = face_swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)

# Gradio UI
def gradio_face_swap(source_img, target_img):
    try:
        result = swap_faces(source_img, target_img)
        if isinstance(result, str):
            return result  # Error message
        return result
    except Exception as e:
        return f"Error: {str(e)}"

gr_interface = gr.Interface(
    fn=gradio_face_swap,
    inputs=[
        gr.Image(label="Upload Source Face"),
        gr.Image(label="Upload Target Image")
    ],
    outputs=gr.Image(label="Swapped Face"),
    title="AI Face Swap",
    description="Upload two images: one with the source face and another where you want the face swapped.",
    live=True
)

# Launch Gradio
app = gr.mount_gradio_app(app, gr_interface, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)