import spaces
import gradio as gr
import numpy as np
import torch
from sentence_transformers import SentenceTransformer


MODEL_ID = "google/embeddinggemma-2"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = (
    torch.bfloat16
    if DEVICE == "cuda" and torch.cuda.is_bf16_supported()
    else torch.float32
)

model = SentenceTransformer(
    MODEL_ID,
    device=DEVICE,
    model_kwargs={"torch_dtype": DTYPE},
    config_kwargs={"vision_config": None, "audio_config": None},
)


TASKS = {
    "Web & document search": (
        "search result",
        "Search for documents relevant to the user's query.",
    ),
    "Question answering": (
        "question answering",
        "Retrieve passages that answer the user's question.",
    ),
    "Code search": (
        "code retrieval",
        "Retrieve code relevant to the user's query.",
    ),
    "Fact checking": (
        "fact checking",
        "Retrieve evidence relevant to the user's claim.",
    ),
}

DEFAULT_DOCUMENTS = """Northern lights | The aurora borealis appears when charged particles from the Sun interact with gases in Earth's atmosphere.
Ocean currents | Large-scale ocean currents redistribute heat around the planet and influence regional climates.
Urban trees | Street trees cool cities by providing shade and releasing water vapor through their leaves.
Sleep and memory | During sleep, the brain strengthens and reorganizes memories formed during the day.
Solar panels | Photovoltaic cells convert sunlight directly into electrical energy using semiconductor materials."""


def parse_documents(raw_documents: str) -> list[tuple[str, str]]:
    """Parse one document per line, with an optional title before a vertical bar."""
    documents = []
    for line in raw_documents.splitlines():
        line = line.strip()
        if not line:
            continue
        if "|" in line:
            title, content = line.split("|", 1)
            title, content = title.strip(), content.strip()
        else:
            title, content = "none", line
        if content:
            documents.append((title or "none", content))
    return documents


@spaces.GPU(duration=90)
def search_documents(
    query: str,
    raw_documents: str,
    task_name: str,
    dimensions: int,
    result_count: int,
) -> tuple[list[list[object]], str]:
    """Embed a query and document set, then return the closest matches and vector details."""
    if not query.strip():
        raise gr.Error("Enter a query to search the collection.")

    documents = parse_documents(raw_documents)
    if not documents:
        raise gr.Error("Add at least one non-empty document, one per line.")

    task_instruction, query_instruction = TASKS[task_name]
    query_text = f"task: {task_instruction} | query: {query.strip()}"
    document_texts = [f"title: {title} | text: {content}" for title, content in documents]

    embeddings = model.encode(
        [query_text, *document_texts],
        convert_to_numpy=True,
        normalize_embeddings=True,
        truncate_dim=int(dimensions),
        show_progress_bar=False,
    )
    query_vector = embeddings[0]
    document_vectors = embeddings[1:]
    scores = document_vectors @ query_vector
    ranked = np.argsort(scores)[::-1][: int(result_count)]

    rows = []
    for rank, index in enumerate(ranked, start=1):
        title, content = documents[int(index)]
        rows.append([rank, title, content, round(float(scores[index]), 4)])

    details = (
        f"**{len(documents)} documents searched** · {dimensions}-dimensional vectors · "
        f"{len(query_vector)} values per embedding\n\n"
        f"Query instruction: `task: {task_instruction} | query: …`\n\n"
        "Document vectors use the `title: … | text: …` format."
    )
    return rows, details


with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")) as demo:
    gr.Markdown(
        """# 🔎 EmbeddingGemma 2 Search Lab
See what an embedding model understands. Add a few documents, ask a question, and inspect how task-aware vectors rank the results.

**740M multimodal model · 100+ languages · 768D native vectors · Apache 2.0**"""
    )

    with gr.Row():
        with gr.Column(scale=5):
            query = gr.Textbox(
                label="Your search query",
                placeholder="What helps cities stay cool during heat waves?",
                lines=2,
            )
            documents = gr.Textbox(
                label="Document collection",
                info="One per line. Optionally add a title like `Title | Document text`.",
                value=DEFAULT_DOCUMENTS,
                lines=9,
            )
        with gr.Column(scale=3):
            task = gr.Dropdown(
                choices=list(TASKS),
                value="Web & document search",
                label="Retrieval task",
            )
            dimensions = gr.Radio(
                choices=[128, 256, 512, 768],
                value=768,
                label="Embedding size",
                info="Shorter vectors use less storage. The model re-normalizes them for cosine search.",
            )
            result_count = gr.Slider(
                minimum=1, maximum=10, value=5, step=1, label="Results to show"
            )
            search = gr.Button("Search collection", variant="primary")

    results = gr.Dataframe(
        headers=["Rank", "Title", "Document", "Cosine score"],
        datatype=["number", "str", "str", "number"],
        label="Nearest documents",
        interactive=False,
        wrap=True,
    )
    details = gr.Markdown()

    gr.Examples(
        examples=[
            ["How does the Sun create colorful lights in the sky?", DEFAULT_DOCUMENTS],
            ["What natural methods can lower temperatures on city streets?", DEFAULT_DOCUMENTS],
            ["How does rest help us remember things?", DEFAULT_DOCUMENTS],
        ],
        inputs=[query, documents],
        label="Try an example",
    )

    search.click(
        fn=search_documents,
        inputs=[query, documents, task, dimensions, result_count],
        outputs=[results, details],
        api_name="search",
    )
    query.submit(
        fn=search_documents,
        inputs=[query, documents, task, dimensions, result_count],
        outputs=[results, details],
        api_name="search_from_query",
    )


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