import os
import uuid
import gradio as gr
import qdrant_client
from qdrant_client import models
from sentence_transformers import SentenceTransformer
from PIL import Image

# ===============================
# Setup
# ===============================
UPLOAD_DIR = "uploaded_images"
os.makedirs(UPLOAD_DIR, exist_ok=True)

COLLECTION = "lost_and_found"

# Qdrant client (in-memory for Hugging Face)
qclient = qdrant_client.QdrantClient(":memory:")
encoder = SentenceTransformer("clip-ViT-B-32")

# Create collection if not exists
if not qclient.collection_exists(COLLECTION):
    qclient.create_collection(
        collection_name=COLLECTION,
        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
    )


# ===============================
# Encode Function (Text or Image)
# ===============================
def encode_data(text=None, image=None):
    if isinstance(image, Image.Image):  # Image is already PIL
        return encoder.encode(image.convert("RGB"))
    elif isinstance(image, str):  # Path to image
        return encoder.encode(Image.open(image).convert("RGB"))
    elif text:
        return encoder.encode([text])[0]
    else:
        return None


# ===============================
# Add Item
# ===============================
def add_item(text, image, uploader_name, uploader_phone):
    try:
        img_path = None
        vector = None

        if isinstance(image, Image.Image):  # PIL image
            img_id = str(uuid.uuid4())
            img_path = os.path.join(UPLOAD_DIR, f"{img_id}.png")
            image.save(img_path)
            vector = encode_data(image=image)

        elif text:
            vector = encode_data(text=text)

        if vector is None:
            return "❌ Please provide an image or text."

        qclient.upsert(
            collection_name=COLLECTION,
            points=[
                models.PointStruct(
                    id=str(uuid.uuid4()),
                    vector=vector.tolist(),
                    payload={
                        "text": text or "",
                        "uploader_name": uploader_name or "N/A",
                        "uploader_phone": uploader_phone or "N/A",
                        "image_path": img_path,
                    },
                )
            ],
        )
        return "✅ Item added to database!"
    except Exception as e:
        return f"❌ Error: {e}"


# ===============================
# Search Function
# ===============================
def search_items(text, image, max_results, min_score):
    try:
        vector = None
        if isinstance(image, Image.Image):  # Search with PIL
            vector = encode_data(image=image)
        elif text:
            vector = encode_data(text=text)

        if vector is None:
            return "❌ Please provide an image or text.", []

        results = qclient.search(
            collection_name=COLLECTION,
            query_vector=vector.tolist(),
            limit=max_results,
            score_threshold=min_score,
        )

        if not results:
            return "No matches found.", []

        # Format results
        result_texts, result_imgs = [], []
        for r in results:
            payload = r.payload
            result_texts.append(
                f"id:{r.id} | score:{r.score:.3f} | "
                f"text:{payload.get('text','')} | "
                f"finder:{payload.get('uploader_name','N/A')} "
                f"({payload.get('uploader_phone','N/A')})"
            )
            if payload.get("image_path") and os.path.exists(payload["image_path"]):
                result_imgs.append(payload["image_path"])

        return "\n".join(result_texts), result_imgs
    except Exception as e:
        return f"❌ Error: {e}", []


# ===============================
# Delete All
# ===============================
def clear_database():
    qclient.delete_collection(COLLECTION)
    qclient.create_collection(
        collection_name=COLLECTION,
        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
    )
    return "🗑️ Database cleared!"


# ===============================
# Gradio UI
# ===============================
with gr.Blocks() as demo:
    gr.Markdown("## 🗝️ Lost & Found - Database")

    # --- Add Item Tab ---
    with gr.Tab("➕ Add Item"):
        with gr.Row():
            text_input = gr.Textbox(label="Description (optional)")
            img_input = gr.Image(type="pil", label="Upload Image")
        with gr.Row():
            uploader_name = gr.Textbox(label="Finder Name")
            uploader_phone = gr.Textbox(label="Finder Phone")
        add_btn = gr.Button("Add to Database")
        add_output = gr.Textbox(label="Status")

        add_btn.click(
            add_item,
            inputs=[text_input, img_input, uploader_name, uploader_phone],
            outputs=add_output,
        )

    # --- Search Tab ---
    with gr.Tab("🔍 Search"):
        with gr.Row():
            search_text = gr.Textbox(label="Search by text (optional)")
            search_img = gr.Image(type="pil", label="Search by image (optional)")
        with gr.Row():
            max_results = gr.Slider(1, 10, value=5, step=1, label="Max results")
            min_score = gr.Slider(0.5, 1.0, value=0.8, step=0.01, label="Min similarity threshold")
        search_btn = gr.Button("Search")
        search_text_out = gr.Textbox(label="Search results (text)")
        search_gallery = gr.Gallery(label="Search Results", columns=2, height="auto")

        search_btn.click(
            search_items,
            inputs=[search_text, search_img, max_results, min_score],
            outputs=[search_text_out, search_gallery],
        )

    # --- Admin Tab ---
    with gr.Tab("🗑️ Admin"):
        clear_btn = gr.Button("Clear Database")
        clear_out = gr.Textbox(label="Status")
        clear_btn.click(clear_database, outputs=clear_out)

demo.launch()