import os
import io
import json
import faiss
import torch
import pandas as pd
import numpy as np
from PIL import Image
import requests
import gradio as gr
from sentence_transformers import SentenceTransformer

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
CSV_PATH = os.path.join(DATA_DIR, "products.csv")
EMB_PATH = os.path.join(DATA_DIR, "embeddings.npy")
INDEX_PATH = os.path.join(DATA_DIR, "index.faiss")
MAP_PATH = os.path.join(DATA_DIR, "row_mapping.json")

MODEL_NAME = "clip-ViT-B-32"

# Lazy globals
_df = None
_index = None
_row_map = None
_model = None

def load_resources():
    global _df, _index, _row_map, _model
    if _df is None:
        _df = pd.read_csv(CSV_PATH)
    if _index is None:
        if not os.path.exists(INDEX_PATH) or not os.path.exists(EMB_PATH) or not os.path.exists(MAP_PATH):
            raise gr.Error("Index not found. Please run build_index.py first (or deploy to Spaces and let it run once).")
        _index = faiss.read_index(INDEX_PATH)
        with open(MAP_PATH, "r", encoding="utf-8") as f:
            _row_map = json.load(f)
    if _model is None:
        _model = SentenceTransformer(MODEL_NAME)

def fetch_image_from_url(url: str) -> Image.Image:
    r = requests.get(url, timeout=15)
    r.raise_for_status()
    return Image.open(io.BytesIO(r.content)).convert("RGB")

def embed_image(pil_img: Image.Image) -> np.ndarray:
    with torch.no_grad():
        emb = _model.encode([pil_img], convert_to_tensor=True)
        emb = torch.nn.functional.normalize(emb, p=2, dim=1)
        return emb.cpu().numpy().astype(np.float32)

def search(query_img, top_k=12, min_score=0.2, categories=None):
    try:
        load_resources()
    except Exception as e:
        raise gr.Error(str(e))

    if query_img is None:
        raise gr.Error("Please upload an image or provide an image URL.")

    q_emb = embed_image(query_img)  # shape (1, d)
    scores, idxs = _index.search(q_emb, top_k + 50)  # oversample, we will filter by category/score
    scores = scores.flatten().tolist()
    idxs = idxs.flatten().tolist()

    # Build results list
    rows = []
    for score, local_idx in zip(scores, idxs):
        row_id = _row_map[local_idx]
        row = _df.iloc[row_id]
        rows.append({
            "id": int(row["id"]),
            "name": row["name"],
            "category": row["category"],
            "price": float(row["price"]),
            "image_url": row["image_url"],
            "similarity": float(score),
        })

    # Filter by category (if any) and score
    if categories:
        rows = [r for r in rows if r["category"] in categories]
    rows = [r for r in rows if r["similarity"] >= min_score]

    # Take top_k after filtering
    rows = rows[:top_k]

    # Prepare gallery & table
    gallery = [[r["image_url"], f'{r["name"]} | {r["category"]} | ₹{r["price"]} | sim={r["similarity"]:.3f}'] for r in rows]
    table = pd.DataFrame(rows, columns=["id","name","category","price","similarity","image_url"])
    return gallery, table

def infer(upload_image, image_url, top_k, min_score, categories):
    # Prefer uploaded image; if absent, try URL
    img = None
    if upload_image is not None:
        img = upload_image
    elif image_url and image_url.strip():
        try:
            img = fetch_image_from_url(image_url.strip())
        except Exception as e:
            raise gr.Error(f"Could not fetch image from URL: {e}")
    else:
        raise gr.Error("Please upload an image or enter a valid image URL.")

    return search(img, top_k=top_k, min_score=min_score, categories=categories)

with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("# 🔎 Visual Product Matcher\nUpload a product photo (or paste an image URL) to find visually similar items.")
    with gr.Row():
        with gr.Column():
            upload = gr.Image(type="pil", label="Upload Image")
            url = gr.Textbox(label="...or Image URL")
            topk = gr.Slider(1, 30, value=12, step=1, label="Results to show")
            mins = gr.Slider(0.0, 1.0, value=0.2, step=0.01, label="Minimum similarity")
            cats = gr.CheckboxGroup(choices=sorted(list(set(pd.read_csv(CSV_PATH)["category"]))), label="Filter by category (optional)")
            btn = gr.Button("Search", variant="primary")
        with gr.Column():
            gr.Markdown("### Query preview")
            # The query preview is just the uploaded image; URL-only flow won't preview until fetched internally.
            preview = gr.Image(label="Uploaded Image", interactive=False)
            gallery = gr.Gallery(label="Similar Products").style(grid=[3], height="auto")
            table = gr.Dataframe(headers=["id","name","category","price","similarity","image_url"],
                                 label="Results (details)", wrap=True)
    upload.change(lambda i: i, inputs=upload, outputs=preview)
    btn.click(infer, inputs=[upload, url, topk, mins, cats], outputs=[gallery, table])

if __name__ == "__main__":
    # For local run: `python build_index.py` first, then `python app.py`
    demo.launch()