"""
Visual Search API - HuggingFace Space

Provides image embedding endpoint using Jina CLIP v2.
Queries Pinecone for similar products.

Deploy to HuggingFace Spaces with ZeroGPU (free).
"""

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

# Pinecone config from HF Secrets
PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
PINECONE_HOST = os.environ.get('PINECONE_HOST')

# 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
        )
        if torch.cuda.is_available():
            model = model.cuda()
        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)  # L2 normalize
        if len(emb) > 512:
            emb = emb[:512]
        return emb.tolist()


def query_pinecone(embedding: list, top_k: int = 12) -> list:
    """Query Pinecone for similar products."""
    if not PINECONE_API_KEY or not PINECONE_HOST:
        return []

    import requests

    resp = requests.post(
        f"https://{PINECONE_HOST}/query",
        headers={
            "Api-Key": PINECONE_API_KEY,
            "Content-Type": "application/json"
        },
        json={
            "vector": embedding,
            "topK": top_k,
            "includeMetadata": True
        },
        timeout=15
    )

    if resp.status_code != 200:
        return []

    matches = resp.json().get('matches', [])
    return [
        {
            'handle': m.get('metadata', {}).get('handle', m.get('id')),
            'title': m.get('metadata', {}).get('title', ''),
            'score': m.get('score', 0),
            'image_url': m.get('metadata', {}).get('image_url', '')
        }
        for m in matches
    ]


def search(image: Image.Image) -> dict:
    """
    Main search function.
    Returns embedding and similar products.
    """
    if image is None:
        return {"error": "No image provided"}

    # Get embedding
    embedding = get_embedding(image)

    # Query Pinecone
    products = query_pinecone(embedding)

    return {
        "embedding": embedding,
        "products": products
    }


def search_simple(image: Image.Image) -> str:
    """Simple search returning product handles."""
    if image is None:
        return "No image"

    embedding = get_embedding(image)
    products = query_pinecone(embedding)

    if not products:
        return "No similar products found"

    return "\n".join([
        f"{i+1}. {p['title']} ({p['handle']}) - {p['score']:.2f}"
        for i, p in enumerate(products)
    ])


# Gradio Interface
with gr.Blocks(title="Visual Search API") as demo:
    gr.Markdown("# Visual Product Search")
    gr.Markdown("Upload an image to find similar products.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            search_btn = gr.Button("Search", variant="primary")

        with gr.Column():
            output = gr.Textbox(label="Results", lines=15)

    search_btn.click(
        fn=search_simple,
        inputs=[image_input],
        outputs=[output]
    )

    gr.Markdown("---")
    gr.Markdown("### API Endpoint")
    gr.Markdown("""
    Use the `/api/predict` endpoint for programmatic access:

    ```python
    from gradio_client import Client

    client = Client("YOUR_SPACE_URL")
    result = client.predict(image_path, api_name="/predict")
    ```
    """)


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