import os
import io
import uuid
import spaces
import torch
import torch.nn.functional as F
import gradio as gr
from PIL import Image
from ultralytics import YOLO
from transformers import AutoImageProcessor, AutoModel
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import cloudinary
import cloudinary.uploader
from typing import List

# ==================== 1. CONFIGURACIÓN DE CREDENCIALES ====================
QDRANT_URL = os.environ.get("QDRANT_URL")
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY")

cloudinary.config(
    cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"),
    api_key=os.environ.get("CLOUDINARY_API_KEY"),
    api_secret=os.environ.get("CLOUDINARY_API_SECRET"),
    secure=True
)

# Conectar a Qdrant y crear la colección si no existe
qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
COLLECTION_NAME = "mascotas_perdidas"
VECTOR_SIZE = 1024  # Dimensión exacta de facebook/convnext-base-224

if not qdrant.collection_exists(COLLECTION_NAME):
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE)
    )

# ==================== 2. CARGAR MODELOS DE IA ====================
print("Cargando YOLOv8 y ConvNeXt...")
detector = YOLO("yolov8n.pt")
model_name = "facebook/convnext-base-224"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.eval()
print("¡Modelos listos!")

# ==================== 3. LÓGICA DE VISIÓN Y EMBEDDINGS (CON ZEROGPU) ====================
def crop_pet(pil_img: Image.Image) -> Image.Image:
    """Detecta y recorta el perro o gato de la imagen."""
    results = detector(pil_img, verbose=False)
    PET_CLASSES = [15, 16]  # 15: gato, 16: perro
    for r in results:
        for box in r.boxes:
            if int(box.cls[0]) in PET_CLASSES:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                w, h = pil_img.size
                x1, y1 = max(0, x1 - 10), max(0, y1 - 10)
                x2, y2 = min(w, x2 + 10), min(h, y2 + 10)
                return pil_img.crop((x1, y1, x2, y2))
    return pil_img

@spaces.GPU
def get_average_embedding(pil_images: List[Image.Image]) -> List[float]:
    """Genera el ADN visual promedio usando ZeroGPU."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    
    embeddings = []
    for img in pil_images:
        cropped = crop_pet(img.convert("RGB"))
        inputs = processor(images=cropped, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
            tensor = outputs.last_hidden_state.mean(dim=[2, 3])
        embeddings.append(F.normalize(tensor, p=2, dim=1))
    
    mean_emb = torch.mean(torch.stack(embeddings), dim=0)
    final_emb = F.normalize(mean_emb, p=2, dim=1)
    return final_emb.squeeze().cpu().tolist()

# ==================== 4. FUNCIONES (INTERFAZ + API REST AUTOMÁTICA) ====================
def registrar_mascota(nombre: str, contacto: str, descripcion: str, archivos):
    if not archivos:
        return "⚠️ Sube entre 1 y 3 fotos del perro."
    
    imgs = [Image.open(f.name if hasattr(f, "name") else f) for f in archivos[:3]]

    # 1. Subir la foto principal a Cloudinary
    buf = io.BytesIO()
    imgs[0].convert("RGB").save(buf, format="JPEG", quality=85)
    buf.seek(0)
    upload_res = cloudinary.uploader.upload(buf, folder="mascotas_perdidas")
    foto_url = upload_res["secure_url"]

    # 2. Calcular embedding promedio (soporta cachorro + adulto)
    vector = get_average_embedding(imgs)

    # 3. Guardar en Qdrant Cloud
    pet_id = str(uuid.uuid4())
    qdrant.upsert(
        collection_name=COLLECTION_NAME,
        points=[
            PointStruct(
                id=pet_id,
                vector=vector,
                payload={
                    "nombre": nombre,
                    "contacto": contacto,
                    "descripcion": descripcion,
                    "foto_url": foto_url
                }
            )
        ]
    )
    return f"✅ ¡Mascota guardada en Qdrant y Cloudinary!\nID: {pet_id}\nFoto URL: {foto_url}"

def buscar_mascota(imagen_pil: Image.Image, umbral: float = 0.75):
    if imagen_pil is None:
        return "⚠️ Sube una foto del perro encontrado.", None

    vector = get_average_embedding([imagen_pil])
    
    resultados = qdrant.query_points(
        collection_name=COLLECTION_NAME,
        query=vector,
        limit=3,
        score_threshold=umbral
    ).points

    if not resultados:
        return "❌ No se encontraron coincidencias por encima del umbral.", None

    texto = "🎉 ¡POSIBLES COINCIDENCIAS ENCONTRADAS!\n\n"
    for idx, hit in enumerate(resultados, 1):
        similitud = round(hit.score * 100, 2)
        texto += f"#{idx} | Coincidencia: {similitud}%\n"
        texto += f"🐶 Nombre: {hit.payload.get('nombre')}\n"
        texto += f"📞 Contacto: {hit.payload.get('contacto')}\n"
        texto += f"📝 Notas: {hit.payload.get('descripcion')}\n\n"

    return texto, resultados[0].payload.get("foto_url")

# ==================== 5. INTERFAZ WEB Y API ====================
with gr.Blocks(title="Buscador IA de Mascotas") as demo:
    gr.Markdown("# 🐾 Sistema Híbrido de Reconocimiento de Mascotas")
    
    with gr.Tab("1. Reportar Mascota Perdida"):
        in_nombre = gr.Textbox(label="Nombre de la mascota")
        in_contacto = gr.Textbox(label="Teléfono / WhatsApp del dueño")
        in_desc = gr.Textbox(label="Señas particulares (collar, zona, etc.)")
        in_fotos = gr.File(file_count="multiple", file_types=["image"], label="Sube 1 a 3 fotos (ej. cachorro y adulto)")
        btn_reg = gr.Button("Guardar Mascota Perdida", variant="primary")
        out_reg = gr.Textbox(label="Estado del registro")
        btn_reg.click(
            fn=registrar_mascota,
            inputs=[in_nombre, in_contacto, in_desc, in_fotos],
            outputs=out_reg,
            api_name="registrar"
        )

    with gr.Tab("2. Buscar Perro Encontrado en la Calle"):
        in_buscar_img = gr.Image(type="pil", label="Sube la foto tomada en la calle")
        in_umbral = gr.Slider(0.60, 0.95, value=0.75, step=0.01, label="Sensibilidad (0.75 recomendado para cambios de edad)")
        btn_bus = gr.Button("Buscar Coincidencias", variant="primary")
        out_bus_txt = gr.Textbox(label="Resultados", lines=8)
        out_bus_img = gr.Image(label="Foto de la mejor coincidencia")
        btn_bus.click(
            fn=buscar_mascota,
            inputs=[in_buscar_img, in_umbral],
            outputs=[out_bus_txt, out_bus_img],
            api_name="buscar"
        )

demo.launch()