import io
import torch
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import StreamingResponse
from PIL import Image

# Importaciones de IA
from diffusers import StableDiffusionPipeline, DDIMScheduler
from huggingface_hub import hf_hub_download
from ip_adapter.ip_adapter_faceid import IPAdapterFaceID
from insightface.app import FaceAnalysis

app = FastAPI(title="IP-Adapter-FaceID API")

# Variables globales para evitar recargar los modelos en cada petición
app_insight = None
ip_model = None

@app.on_event("startup")
def load_models():
    global app_insight, ip_model
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Iniciando servidor. Hardware detectado: {device.upper()}")
    
    # 1. Inicializar InsightFace para detectar rostros
    print("Cargando InsightFace...")
    app_insight = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
    app_insight.prepare(ctx_id=0, det_size=(640, 640))
    
    # 2. Inicializar Stable Diffusion v1.5
    print("Cargando modelo base de Stable Diffusion...")
    base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
    noise_scheduler = DDIMScheduler(
        num_train_timesteps=1000,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        clip_sample=False,
        set_alpha_to_one=False,
        steps_offset=1,
    )
    
    pipe = StableDiffusionPipeline.from_pretrained(
        base_model_path,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        scheduler=noise_scheduler
    ).to(device)
    
    # 3. Descargar y cargar las capas de IP-Adapter FaceID
    print("Descargando pesos de IP-Adapter-FaceID...")
    ip_ckpt_path = hf_hub_download(
        repo_id="h94/IP-Adapter-FaceID", 
        filename="ip-adapter-faceid_sd15.bin", 
        repo_type="model"
    )
    
    print("Acoplando IP-Adapter al modelo base...")
    ip_model = IPAdapterFaceID(pipe, ip_ckpt_path, device, num_tokens=4)
    print("¡Todos los modelos cargados exitosamente!")

@app.get("/")
def read_root():
    return {"status": "running", "message": "API de IP-Adapter-FaceID lista. Realiza un POST a /predict"}

@app.post("/predict")
async def predict(
    prompt: str = Form(...),
    negative_prompt: str = Form("monochrome, lowres, bad anatomy, worst quality, low quality"),
    num_inference_steps: int = Form(30),
    face_image: UploadFile = File(...)
):
    try:
        # 1. Leer y preparar la imagen
        face_bytes = await face_image.read()
        image = Image.open(io.BytesIO(face_bytes)).convert("RGB")
        cv_img = np.array(image)
        
        # 2. Extraer el rostro (Embedding) usando InsightFace
        faces = app_insight.get(cv_img)
        
        if len(faces) == 0:
            return {"error": "No se detectó ningún rostro en la imagen provista. Por favor intenta con otra foto."}
            
        face_id_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0).unsqueeze(0)
        
        # 3. Generar la imagen
        images = ip_model.generate(
            prompt=prompt,
            negative_prompt=negative_prompt,
            face_id_embed=face_id_embed,
            num_samples=1,
            num_inference_steps=num_inference_steps,
            seed=42 # Cambia o recibe este parámetro si quieres resultados aleatorios
        )
        
        # 4. Formatear la respuesta como PNG
        output_image = images[0]
        img_byte_arr = io.BytesIO()
        output_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        return StreamingResponse(img_byte_arr, media_type="image/png")
        
    except Exception as e:
        return {"error": str(e)}