import gradio as gr
import uuid
import os
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct
from sentence_transformers import SentenceTransformer
from PIL import Image
import numpy as np

# Create uploads folder
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

# Connect to Qdrant
COLLECTION = "lost_and_found"
qclient = QdrantClient(":memory:")  # for demo, replace with your Qdrant server for persistence

# Load CLIP model
model = SentenceTransformer("sentence-transformers/clip-ViT-B-32-multilingual-v1")

# Ensure collection exists
qclient.recreate_collection(
    collection_name=COLLECTION,
    vectors_config={"size": 512, "distance": "Cosine"}
)

# Encode helper
def encode_data(text=None, image=None):
    if image is not None:
        img = Image.open(image).convert("RGB")
        emb = model.encode(img, convert_to_numpy=True, normalize_embeddings=True)
    elif text:
        emb = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
    else:
        raise ValueError("Need text or image")
    return emb.astype(np.float32)

# Add item
def add_item(mode, text, image, name, phone):
    try:
        vector = encode_data(text=text if text else None, image=image if image else None)
        payload = {
            "mode": mode,
            "text": text,
            "has_image": image is not None,
        }

        # Save image file if uploaded
        if image is not None:
            img = Image.open(image).convert("RGB")
            fname = f"{uuid.uuid4().hex}.png"
            fpath = os.path.join(UPLOAD_DIR, fname)
            img.save(fpath)
            payload["image_path"] = fpath

        if mode == "found":
            payload["finder_name"] = name
            payload["finder_phone"] = phone

        qclient.upsert(
            collection_name=COLLECTION,
            points=[PointStruct(id=str(uuid.uuid4()), vector=vector.tolist(), payload=payload)]
        )
        return "✅ Item added successfully!"
    except Exception as e:
        return f"❌ Error: {e}"

# Search items
def search_items(query_image, query_text, limit, min_score):
    try:
        query_vector = encode_data(
            text=query_text if query_text else None,
            image=query_image if query_image else None
        )
        results = qclient.search(
            collection_name=COLLECTION,
            query_vector=query_vector.tolist(),
            limit=limit,
        )
        out_texts, out_imgs = [], []
        for r in results:
            if r.score < min_score:
                continue
            pl = r.payload
            info = f"id:{r.id} | score:{r.score:.4f} | mode:{pl.get('mode','')}"
            if pl.get("text"):
                info += f" | text:{pl['text']}"
            if pl.get("mode") == "found":
                info += f" | found by: {pl.get('finder_name','?')} ({pl.get('finder_phone','?')})"
            out_texts.append(info)

            if pl.get("image_path"):
                out_imgs.append(pl["image_path"])

        return "\n".join(out_texts) if out_texts else "No matches.", out_imgs
    except Exception as e:
        return f"❌ Error: {e}", []

# Clear all images
def clear_all_images():
    try:
        # Clear uploads folder
        for f in os.listdir(UPLOAD_DIR):
            os.remove(os.path.join(UPLOAD_DIR, f))
        # Clear Qdrant
        qclient.delete(
            collection_name=COLLECTION,
            points_selector={"filter": {"must": [{"key": "has_image", "match": {"value": True}}]}}
        )
        return "🗑️ All image items cleared!"
    except Exception as e:
        return f"❌ Error clearing images: {e}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🔎 Lost & Found System")

    with gr.Tab("➕ Add Item"):
        mode = gr.Radio(["lost", "found"], label="Mode")
        text = gr.Textbox(label="Describe the item (optional)")
        img = gr.Image(type="filepath", label="Upload image (optional)")
        name = gr.Textbox(label="Finder's Name (only if found)", placeholder="John Doe")
        phone = gr.Textbox(label="Finder's Phone (only if found)", placeholder="+1234567890")
        add_btn = gr.Button("Add Item")
        add_out = gr.Textbox(label="Add result")
        add_btn.click(add_item, inputs=[mode, text, img, name, phone], outputs=add_out)

    with gr.Tab("🔍 Search"):
        query_text = gr.Textbox(label="Search by text (optional)")
        query_img = gr.Image(type="filepath", label="Search by image (optional)")
        max_results = gr.Slider(1, 10, value=5, step=1, label="Max results")
        score_slider = gr.Slider(0.5, 1.0, value=0.9, step=0.01, label="Min similarity threshold")
        search_btn = gr.Button("Search")
        search_out = gr.Textbox(label="Search results (text)")
        gallery = gr.Gallery(label="Search Results", show_label=True, elem_id="gallery", columns=2, height="auto")
        search_btn.click(search_items, inputs=[query_img, query_text, max_results, score_slider], outputs=[search_out, gallery])

    with gr.Tab("🗑️ Admin"):
        clear_btn = gr.Button("Clear All Images")
        clear_out = gr.Textbox(label="Clear Result")
        clear_btn.click(clear_all_images, outputs=clear_out)

demo.launch()