import hashlib
import os

import gradio as gr
import numpy as np
import pandas as pd
import torch
from sentence_transformers import SentenceTransformer

MODEL_ID = "google/embeddinggemma-2"
ASSETS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")

# cpu-basic: float32. bfloat16 is for GPUs; float16 overflows this model.
torch.set_num_threads(min(2, os.cpu_count() or 2))
model = SentenceTransformer(
    MODEL_ID,
    device="cpu",
    model_kwargs={"dtype": torch.float32},
    config_kwargs={"audio_config": None},
)
model.max_seq_length = 2048

DIMS = ["768", "512", "256", "128"]
SEARCH_TASKS = {
    "Web search": "SearchQuery",
    "Question answering": "QuestionAnswering",
    "Fact checking": "FactChecking",
    "Code search": "CodeRetrieval",
}
SAMPLES = [
    ("car.jpg", "Green car"),
    ("cat.jpg", "Tabby cat"),
    ("dog.jpg", "Spotted dog"),
    ("bee.jpg", "Honey bee"),
    ("beach.jpg", "Sunny beach"),
    ("mountain.jpg", "Snowy mountain"),
    ("bike.jpg", "Bicycle gathering"),
]

DEFAULT_DOCS = """Northern lights | The aurora appears when charged particles from the sun collide with gases in Earth's atmosphere.
Ocean currents | Large-scale currents move heat around the planet and shape regional climates.
Eiffel Tower | The iron tower was built in Paris for the 1889 World's Fair.
Sleep | During sleep the brain strengthens memories formed during the day.
Photosynthesis | Plants turn sunlight, water, and carbon dioxide into glucose and oxygen.
Auroras | Las auroras se producen cuando el viento solar choca con la atmósfera terrestre.
Bread | Il pane si prepara con farina, acqua, lievito e sale.
Markets | The stock market closed higher on Friday after strong earnings reports."""

DEFAULT_SENTENCES = """The cat is sleeping on the sofa.
Le chat dort sur le canapé.
Die Katze schläft auf dem Sofa.
A dog is barking in the garden.
Quarterly revenue grew by 12 percent."""


def _unit(embeddings: torch.Tensor) -> np.ndarray:
    if embeddings.ndim == 1:
        embeddings = embeddings.unsqueeze(0)
    return torch.nn.functional.normalize(embeddings.float(), dim=-1).cpu().numpy()


_IMAGE_CACHE: dict[tuple[str, int], np.ndarray] = {}


def _file_key(image, dim: int) -> tuple[str, int] | None:
    if not isinstance(image, str) or not os.path.isfile(image):
        return None
    digest = hashlib.sha256()
    with open(image, "rb") as handle:
        digest.update(handle.read())
    return digest.hexdigest(), int(dim)


def _embed(inputs, prompt_name: str | None, dim: int) -> np.ndarray:
    embeddings = model.encode(
        inputs,
        prompt_name=prompt_name,
        batch_size=8,
        truncate_dim=int(dim),
        normalize_embeddings=True,
        convert_to_tensor=True,
        show_progress_bar=False,
    )
    return _unit(embeddings)


def _embed_images(images, dim: int) -> np.ndarray:
    """Embed images, reusing vectors for files we have already seen."""
    rows: list[np.ndarray | None] = [None] * len(images)
    missing_images = []
    missing_at = []
    for index, image in enumerate(images):
        key = _file_key(image, dim)
        cached = _IMAGE_CACHE.get(key) if key is not None else None
        if cached is not None:
            rows[index] = cached
        else:
            missing_images.append(image)
            missing_at.append(index)
    if missing_images:
        fresh = _embed([{"image": image} for image in missing_images], None, dim)
        for row, index in enumerate(missing_at):
            rows[index] = fresh[row]
            key = _file_key(images[index], dim)
            if key is not None:
                _IMAGE_CACHE[key] = fresh[row]
    return np.stack(rows)


def _parse_docs(raw: str) -> list[tuple[str, str]]:
    documents = []
    for line in (raw or "").splitlines():
        line = line.strip()
        if not line:
            continue
        if " | " in line:
            title, content = line.split(" | ", 1)
            title, content = title.strip() or "none", content.strip()
        else:
            title, content = "none", line
        if content:
            documents.append((title, content))
    return documents


def _iter_images(gallery):
    if not gallery:
        return
    for item in gallery:
        caption = None
        image = item
        if isinstance(item, (list, tuple)) and item:
            image = item[0]
            caption = item[1] if len(item) > 1 else None
        elif isinstance(item, dict):
            image = item.get("image", item.get("path", item.get("name")))
            caption = item.get("caption")
            if isinstance(image, dict):
                caption = caption or image.get("orig_name")
                image = image.get("path") or image.get("url")
        if image is None:
            continue
        if isinstance(image, str):
            label = caption or os.path.splitext(os.path.basename(image))[0]
        else:
            label = caption or "image"
        yield image, str(label)


def _sample_gallery():
    return [(os.path.join(ASSETS, name), caption) for name, caption in SAMPLES]


print("Caching sample image embeddings...", flush=True)
_embed_images([os.path.join(ASSETS, name) for name, _caption in SAMPLES], 768)
print(f"Cached {len(_IMAGE_CACHE)} sample image embeddings", flush=True)


def search(query: str, docs_text: str, task: str, dim: int):
    """Rank documents against a text query with EmbeddingGemma 2.

    The query receives the selected task prefix. Each document is formatted as
    `title: … | text: …`. Vectors are Matryoshka-truncated and L2-normalized,
    then scored with cosine similarity.
    """
    query = (query or "").strip()
    documents = _parse_docs(docs_text)
    if not query:
        raise gr.Error("Enter a query.")
    if not documents:
        raise gr.Error("Add at least one document, one per line.")
    if len(documents) > 32:
        raise gr.Error("Use at most 32 documents.")

    prompt_name = SEARCH_TASKS[task]
    query_vec = _embed([query], prompt_name, dim)[0]
    doc_inputs = [f"title: {title} | text: {content}" for title, content in documents]
    doc_vecs = _embed(doc_inputs, None, dim)
    scores = doc_vecs @ query_vec
    order = np.argsort(-scores)

    rows = []
    for rank, index in enumerate(order, start=1):
        title, content = documents[int(index)]
        rows.append([rank, title, content, round(float(scores[index]), 4)])
    note = (
        f"{len(documents)} documents · {int(dim)}-d · "
        f"query prefix `{model.prompts[prompt_name].strip()}`"
    )
    return rows, note


def search_images(query: str, gallery, dim: int):
    """Rank images against a text query in EmbeddingGemma 2's shared vector space.

    The query uses the web-search prefix. Images are embedded with no text prefix.
    Both sides are truncated to the same dimension and L2-normalized before cosine scoring.
    """
    query = (query or "").strip()
    images = list(_iter_images(gallery))
    if not query:
        raise gr.Error("Enter a query.")
    if not images:
        raise gr.Error("Add at least one image.")
    if len(images) > 12:
        raise gr.Error("Use at most 12 images.")

    query_vec = _embed([query], "SearchQuery", dim)[0]
    image_vecs = _embed_images([image for image, _label in images], dim)
    scores = image_vecs @ query_vec
    order = np.argsort(-scores)

    ranked_gallery = []
    rows = []
    for rank, index in enumerate(order, start=1):
        image, label = images[int(index)]
        score = float(scores[index])
        ranked_gallery.append((image, f"{score:.3f} · {label}"))
        rows.append([rank, label, round(score, 4)])
    note = f"{len(images)} images · {int(dim)}-d · query prefix `{model.prompts['SearchQuery'].strip()}`"
    return ranked_gallery, rows, note


def compare(text: str, dim: int):
    """Score pairwise cosine similarity for one sentence per line.

    Every sentence uses the sentence-similarity task prefix. Vectors are
    Matryoshka-truncated and L2-normalized.
    """
    sentences = [line.strip() for line in (text or "").splitlines() if line.strip()]
    if len(sentences) < 2:
        raise gr.Error("Enter at least two sentences, one per line.")
    if len(sentences) > 16:
        raise gr.Error("Use at most 16 sentences.")

    vectors = _embed(sentences, "SentenceSimilarity", dim)
    similarity = np.round((vectors @ vectors.T).astype(np.float64), 3)
    headers = [f"#{i + 1}" for i in range(len(sentences))]
    frame = pd.DataFrame(similarity, columns=headers)
    frame.insert(0, "Sentence", [f"#{i + 1} {sentence}" for i, sentence in enumerate(sentences)])
    note = f"{len(sentences)} sentences · {int(dim)}-d · `{model.prompts['SentenceSimilarity'].strip()}`"
    return frame, note


CSS = """
.gradio-container { max-width: 1100px !important; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="EmbeddingGemma 2") as demo:
    gr.Markdown(
        """# EmbeddingGemma 2
Search text and images in one vector space with [google/embeddinggemma-2](https://huggingface.co/google/embeddinggemma-2). Task prefixes steer text. Images share the same 768-d space, and vectors can be shortened to 512, 256, or 128 dimensions.
"""
    )

    with gr.Tabs():
        with gr.Tab("Documents"):
            with gr.Row():
                with gr.Column(scale=5):
                    query = gr.Textbox(
                        label="Query",
                        value="What causes the northern lights?",
                        lines=2,
                    )
                    documents = gr.Textbox(
                        label="Documents",
                        info="One per line. Optional title: `Title | text`.",
                        value=DEFAULT_DOCS,
                        lines=12,
                    )
                    with gr.Accordion("Settings", open=False):
                        task = gr.Dropdown(
                            choices=list(SEARCH_TASKS),
                            value="Web search",
                            label="Task",
                        )
                        dim = gr.Radio(choices=DIMS, value="768", label="Dimensions")
                    search_btn = gr.Button("Search", variant="primary")
                with gr.Column(scale=6):
                    doc_note = gr.Markdown()
                    doc_results = gr.Dataframe(
                        headers=["Rank", "Title", "Document", "Cosine"],
                        label="Closest documents",
                        interactive=False,
                        wrap=True,
                    )
            doc_examples = [
                ["What causes the northern lights?", DEFAULT_DOCS, "Web search", "768"],
                ["When was the Eiffel Tower built?", DEFAULT_DOCS, "Question answering", "768"],
                ["Solar wind creates auroras.", DEFAULT_DOCS, "Fact checking", "256"],
            ]
            gr.Examples(
                examples=doc_examples,
                inputs=[query, documents, task, dim],
                outputs=[doc_results, doc_note],
                fn=search,
                cache_examples=True,
                cache_mode="lazy",
                label="Examples",
            )
            search_btn.click(search, [query, documents, task, dim], [doc_results, doc_note], api_name="search")
            query.submit(search, [query, documents, task, dim], [doc_results, doc_note])

        with gr.Tab("Images"):
            with gr.Row():
                with gr.Column(scale=5):
                    image_query = gr.Textbox(
                        label="Query",
                        value="a green car parked by a yellow wall",
                        lines=2,
                    )
                    gallery = gr.Gallery(
                        value=_sample_gallery(),
                        label="Image collection",
                        type="filepath",
                        columns=4,
                        height=360,
                        interactive=True,
                    )
                    image_dim = gr.Radio(choices=DIMS, value="768", label="Dimensions")
                    image_btn = gr.Button("Rank images", variant="primary")
                with gr.Column(scale=6):
                    image_note = gr.Markdown()
                    ranked = gr.Gallery(label="Ranked images", columns=4, height=360, interactive=False)
                    image_table = gr.Dataframe(
                        headers=["Rank", "Image", "Cosine"],
                        label="Scores",
                        interactive=False,
                    )
            gr.Examples(
                examples=[
                    ["a green car parked by a yellow wall"],
                    ["a tabby cat"],
                    ["a honey bee"],
                    ["a sandy beach with umbrellas"],
                    ["people with bicycles"],
                ],
                inputs=[image_query],
                label="Try a query",
            )
            image_btn.click(
                search_images,
                [image_query, gallery, image_dim],
                [ranked, image_table, image_note],
                api_name="search_images",
            )
            image_query.submit(
                search_images,
                [image_query, gallery, image_dim],
                [ranked, image_table, image_note],
            )

        with gr.Tab("Similarity"):
            with gr.Row():
                with gr.Column():
                    sentences = gr.Textbox(
                        label="Sentences",
                        info="One per line. The same prefix is applied to every line.",
                        value=DEFAULT_SENTENCES,
                        lines=10,
                    )
                    sim_dim = gr.Radio(choices=DIMS, value="768", label="Dimensions")
                    sim_btn = gr.Button("Compare", variant="primary")
                with gr.Column():
                    sim_note = gr.Markdown()
                    sim_table = gr.Dataframe(label="Cosine similarity", interactive=False, wrap=True)
            gr.Examples(
                examples=[
                    [DEFAULT_SENTENCES, "768"],
                    [
                        "I love this phone, the battery lasts forever.\nThe screen cracked after one day.\nCe téléphone est génial.\nPantalla rota, muy decepcionado.",
                        "512",
                    ],
                ],
                inputs=[sentences, sim_dim],
                outputs=[sim_table, sim_note],
                fn=compare,
                cache_examples=True,
                cache_mode="lazy",
                label="Examples",
            )
            sim_btn.click(compare, [sentences, sim_dim], [sim_table, sim_note], api_name="compare")

    gr.Markdown(
        "Loads the text and vision encoders from "
        f"[{MODEL_ID}](https://huggingface.co/{MODEL_ID}) in float32. "
        "Images are embedded without a task prefix. Audio and video are supported by the model and omitted here."
    )


if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
        css=CSS,
    )