import gradio as gr
from qdrant_client import QdrantClient
from qdrant_client import models
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_url
from dotenv import load_dotenv
import os
from functools import lru_cache

load_dotenv()

URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
sentence_embedding_model = SentenceTransformer("BAAI/bge-large-en")

print(URL)
print(QDRANT_API_KEY)
collection_name = "dataset_cards"
client = QdrantClient(
    url=URL,
    api_key=QDRANT_API_KEY,
)


def format_results(results):
    markdown = ""
    for result in results:
        hub_id = result.payload["id"]
        url = hf_hub_url(hub_id, "README.md", repo_type="dataset")
        header = f"## [{hub_id}]({url})"
        markdown += header + "\n"
        markdown += result.payload["section_text"] + "\n"
    return markdown


@lru_cache()
def search(query: str):
    query_ = sentence_embedding_model.encode(
        f"Represent this sentence for searching relevant passages:{query}"
    )
    results = client.search(
        collection_name="dataset_cards",
        query_vector=query_,
        limit=10,
    )
    return format_results(results)


@lru_cache()
def hub_id_qdrant_id(hub_id):
    matches = client.scroll(
        collection_name="dataset_cards",
        scroll_filter=models.Filter(
            must=[
                models.FieldCondition(key="id", match=models.MatchValue(value=hub_id)),
            ]
        ),
        limit=1,
        with_payload=True,
        with_vectors=False,
    )
    try:
        return matches[0][0].id
    except IndexError as e:
        raise gr.Error(
            f"Hub id {hub_id} not in out database. This could be because it is very new or because it doesn't have much documentation."
        ) from e


@lru_cache()
def recommend(hub_id):
    positive_id = hub_id_qdrant_id(hub_id)
    results = client.recommend(collection_name=collection_name, positive=[positive_id])
    return format_results(results)


def query(search_term, search_type):
    if search_type == "Recommend similar datasets":
        return recommend(search_term)
    else:
        return search(search_term)


with gr.Blocks() as demo:
    gr.Markdown("## 🤗 Semantic Dataset Search")
    with gr.Row():
        gr.Markdown(
            "This Gradio app allows you to search for datasets based on their descriptions. You can either search for similar datasets to a given dataset or search for datasets based on a query."
        )
    with gr.Row():
        search_term = gr.Textbox(value="movie review sentiment",
            label="hub id i.e. IMDB or query i.e. movie review sentiment"
        )
    with gr.Row():
        with gr.Row():
            find_similar_btn = gr.Button("Search")
            search_type = gr.Radio(
                ["Recommend similar datasets", "Semantic Search"],
                label="Search type",
                value="Semantic Search",
                interactive=True,
            )

    results = gr.Markdown()
    find_similar_btn.click(query, [search_term, search_type], results)


demo.launch()