"""
Visual Search API - HuggingFace Space
Returns embedding vector for external Pinecone queries
"""

import os
import gradio as gr
import torch
import numpy as np
from PIL import Image
import json

# Model (loaded on first use)
model = None


def load_model():
    """Load Jina CLIP v2 model."""
    global model
    if model is None:
        print("Loading Jina CLIP v2...")
        from transformers import AutoModel
        model = AutoModel.from_pretrained(
            "jinaai/jina-clip-v2",
            trust_remote_code=True
        )
        model.eval()
        print("Model loaded!")
    return model


def get_embedding(image: Image.Image) -> list:
    """Generate 512-dim embedding for an image."""
    m = load_model()

    with torch.no_grad():
        emb = m.encode_image(image)
        if hasattr(emb, 'cpu'):
            emb = emb.cpu().numpy()
        emb = emb.flatten()
        emb = emb / np.linalg.norm(emb)
        if len(emb) > 512:
            emb = emb[:512]
        return emb.tolist()


def search(image):
    """Return embedding vector as JSON."""
    if image is None:
        return json.dumps({"error": "No image provided"})

    try:
        print("Generating embedding...")
        embedding = get_embedding(image)
        print(f"Embedding generated: {len(embedding)} dimensions")

        # Return embedding as JSON
        result = {
            "embedding": embedding,
            "dimensions": len(embedding)
        }
        return json.dumps(result, indent=2)

    except Exception as e:
        import traceback
        traceback.print_exc()
        return json.dumps({"error": str(e)})


# Gradio interface - returns embedding as JSON
demo = gr.Interface(
    fn=search,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Textbox(label="Embedding Vector (JSON)", lines=15),
    title="Visual Search - Embedding Generator",
    description="Upload an image to get its 512-dimensional CLIP embedding as JSON."
)

if __name__ == "__main__":
    demo.queue().launch()