from fastapi import FastAPI, File, UploadFile, HTTPException 
from fastapi.responses import JSONResponse
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
import uvicorn
import logging
import base64

# Configurar logging
logging.basicConfig(level=logging.INFO)

app = FastAPI()


face_app = FaceAnalysis(name='buffalo_l')
face_app.prepare(ctx_id=0, det_size=(640, 640))
# Carga el modelo desde la carpeta 'models'
swapper = insightface.model_zoo.get_model('models/inswapper_128.onnx', download=False, download_zip=False)

def process_images(image1_bytes: bytes, image2_bytes: bytes):
    # Convertir bytes a imágenes con OpenCV
    nparr1 = np.frombuffer(image1_bytes, np.uint8)
    img1 = cv2.imdecode(nparr1, cv2.IMREAD_COLOR)
    nparr2 = np.frombuffer(image2_bytes, np.uint8)
    img2 = cv2.imdecode(nparr2, cv2.IMREAD_COLOR)
    if img1 is None or img2 is None:
        raise ValueError("Error al decodificar una de las imágenes.")
    
    # Detección de rostros
    faces1 = face_app.get(img1)
    faces2 = face_app.get(img2)
    if not faces1 or not faces2:
        raise ValueError("No se detectó ningún rostro en una o ambas imágenes.")
    
    # Usar el primer rostro detectado en cada imagen
    face1 = faces1[0]
    face2 = faces2[0]
    
    # Realizar el intercambio de rostros
    img1_swap = swapper.get(img1, face1, face2, paste_back=True)
    img2_swap = swapper.get(img2, face2, face1, paste_back=True)
    return img1_swap, img2_swap

def image_to_base64(image) -> str:
    # Codificar la imagen a JPEG y luego a base64
    success, buffer = cv2.imencode(".jpg", image)
    if not success:
        raise ValueError("No se pudo codificar la imagen.")
    return base64.b64encode(buffer).decode("utf-8")

@app.post("/swap_faces")
async def swap_faces(image1: UploadFile = File(...), image2: UploadFile = File(...)):
    """
    Endpoint que recibe dos imágenes por POST y devuelve un JSON con las imágenes
    con los rostros intercambiados codificadas en base64.
    """
    try:
        # Leer los archivos
        image1_bytes = await image1.read()
        image2_bytes = await image2.read()
        # Procesar las imágenes
        img1_swap, img2_swap = process_images(image1_bytes, image2_bytes)
        # Convertir imágenes a base64
        result = {
            "image1": image_to_base64(img1_swap),
            "image2": image_to_base64(img2_swap)
        }
        return JSONResponse(content=result)
    except Exception as e:
        logging.exception("Error al procesar las imágenes")
        raise HTTPException(status_code=500, detail=str(e))

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