"""Interactive multilingual embedding demo for Hugging Face Spaces."""

from __future__ import annotations

from functools import lru_cache
from typing import Any

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


MODEL_ID = "KaLM-Embedding/KaLM-embedding-multilingual-mini-instruct-v2.5"
MAX_TEXTS = 100
DEFAULT_RETRIEVAL_INSTRUCTION = "Given a query, retrieve documents that answer the query"
DEFAULT_SIMILARITY_INSTRUCTION = "Retrieve semantically similar text."


@lru_cache(maxsize=1)
def get_model() -> SentenceTransformer:
    """Load the model once per Space process."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_kwargs: dict[str, Any] = {}
    if device == "cuda":
        model_kwargs["torch_dtype"] = torch.float16

    model = SentenceTransformer(
        MODEL_ID,
        trust_remote_code=True,
        device=device,
        model_kwargs=model_kwargs,
    )
    # A responsive default for a public CPU Space; the model supports longer input.
    model.max_seq_length = 512
    return model


def clean_texts(raw_text: str, label: str) -> list[str]:
    texts = [line.strip() for line in (raw_text or "").splitlines() if line.strip()]
    if not texts:
        raise gr.Error(f"Please enter at least one {label}.")
    if len(texts) > MAX_TEXTS:
        raise gr.Error(f"Please limit {label} to {MAX_TEXTS} lines per run.")
    return texts


def encode_documents(model: SentenceTransformer, documents: list[str]) -> np.ndarray:
    return model.encode(
        documents,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )


def encode_query(model: SentenceTransformer, query: str, instruction: str) -> np.ndarray:
    instruction = instruction.strip()
    if instruction:
        prompt = f"Instruct: {instruction}\nQuery:"
    else:
        # Sentence Transformers 3.x does not expose encode_query(), so apply
        # the model card's default retrieval instruction explicitly.
        prompt = f"Instruct: {DEFAULT_RETRIEVAL_INSTRUCTION}\nQuery:"
    return model.encode(
        [query],
        prompt=prompt,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )[0]


@spaces.GPU(duration=60)
def search(
    query: str, raw_documents: str, top_k: int, instruction: str
) -> tuple[list[list[Any]], dict[str, Any]]:
    query = (query or "").strip()
    if not query:
        raise gr.Error("Please enter a query.")

    documents = clean_texts(raw_documents, "document")
    model = get_model()
    query_embedding = encode_query(model, query, instruction)
    document_embeddings = encode_documents(model, documents)
    scores = document_embeddings @ query_embedding
    best_indices = np.argsort(-scores)[: min(int(top_k), len(documents))]

    rows = [
        [rank, f"{float(scores[index]):.4f}", documents[int(index)]]
        for rank, index in enumerate(best_indices, start=1)
    ]
    metadata = {
        "model": MODEL_ID,
        "embedding_dimension": int(query_embedding.shape[0]),
        "documents_searched": len(documents),
        "device": "CUDA" if torch.cuda.is_available() else "CPU",
        "similarity": "cosine (normalized embeddings)",
    }
    return rows, metadata


@spaces.GPU(duration=60)
def compare(left: str, right: str, instruction: str) -> tuple[str, dict[str, Any]]:
    left = (left or "").strip()
    right = (right or "").strip()
    if not left or not right:
        raise gr.Error("Please enter both texts to compare.")

    model = get_model()
    if instruction.strip():
        prompt = f"Instruct: {instruction.strip()}\nQuery:"
        embeddings = model.encode(
            [left, right],
            prompt=prompt,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )
    else:
        embeddings = model.encode(
            [left, right],
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )
    score = float(embeddings[0] @ embeddings[1])
    return f"## Cosine similarity: **{score:.4f}**", {
        "model": MODEL_ID,
        "embedding_dimension": int(embeddings.shape[1]),
        "interpretation": "Higher scores indicate greater semantic similarity.",
    }


@spaces.GPU(duration=60)
def inspect_embeddings(raw_text: str) -> tuple[list[list[str]], dict[str, Any]]:
    texts = clean_texts(raw_text, "text")
    model = get_model()
    embeddings = model.encode(
        texts,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )
    rows = [
        [index + 1, text, ", ".join(f"{value:.4f}" for value in vector[:12])]
        for index, (text, vector) in enumerate(zip(texts, embeddings))
    ]
    return rows, {
        "model": MODEL_ID,
        "embedding_dimension": int(embeddings.shape[1]),
        "vectors_created": len(texts),
        "normalization": "L2 normalized",
        "preview": "The table shows the first 12 dimensions of each vector.",
    }


EXAMPLE_QUERY = "What is the capital of China?"
EXAMPLE_DOCUMENTS = """Beijing is the capital city of China.
Paris is the capital and most populous city of France.
Gravity attracts bodies with mass toward one another.
中国的首都是北京。"""


CSS = """
.gradio-container { max-width: 1120px !important; }
#hero { text-align: center; margin: 0.5rem 0 1.5rem; }
#hero h1 { margin-bottom: 0.35rem; }
.notice { border-left: 4px solid #6366f1; padding: 0.6rem 0.9rem; background: #eef2ff; border-radius: 0.4rem; }
"""


with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="KaLM Embedding Demo") as demo:
    gr.Markdown(
        """
        <div id="hero">
          <h1>KaLM Embedding —Versatile Text Embedding</h1>
          <p>Explore semantic retrieval and text similarity with <code>KaLM-embedding-multilingual-mini-instruct-v2.5</code>.</p>
        </div>
        """
    )
    gr.HTML("<div class='notice'>The model is loaded on first use. Keep inputs concise for the most responsive public demo.</div>")

    with gr.Tabs():
        with gr.Tab("Semantic Retrieval", id="search"):
            with gr.Row():
                with gr.Column(scale=2):
                    query_input = gr.Textbox(
                        label="Query / 查询",
                        placeholder="Ask in English, Chinese, or another supported language…",
                        lines=2,
                    )
                with gr.Column(scale=1):
                    top_k = gr.Slider(1, 10, value=3, step=1, label="Results to show")
            documents_input = gr.Textbox(
                label="Documents / 文档",
                placeholder="One document per line. They will be ranked by meaning, not keywords.",
                lines=10,
            )
            with gr.Accordion("Advanced retrieval instruction", open=False):
                search_instruction = gr.Textbox(
                    label="Task instruction",
                    value=DEFAULT_RETRIEVAL_INSTRUCTION,
                    lines=2,
                    info="Describe what makes a document relevant. Leave blank to use the model default.",
                )
            search_button = gr.Button("Search documents", variant="primary")
            search_results = gr.Dataframe(
                headers=["Rank", "Cosine score", "Document"],
                datatype=["number", "str", "str"],
                interactive=False,
                wrap=True,
                label="Ranked results",
            )
            search_metadata = gr.JSON(label="Run details")
            search_button.click(
                search,
                inputs=[query_input, documents_input, top_k, search_instruction],
                outputs=[search_results, search_metadata],
            )
            gr.Examples(
                examples=[[EXAMPLE_QUERY, EXAMPLE_DOCUMENTS, 3, DEFAULT_RETRIEVAL_INSTRUCTION]],
                inputs=[query_input, documents_input, top_k, search_instruction],
                label="Try an example",
            )

        with gr.Tab("Semantic Textual Similarity", id="similarity"):
            with gr.Row():
                similarity_left = gr.Textbox(label="Text A", lines=6, value="北京是中国的首都。")
                similarity_right = gr.Textbox(label="Text B", lines=6, value="The capital of China is Beijing.")
            with gr.Accordion("Optional task instruction", open=False):
                similarity_instruction = gr.Textbox(
                    label="Instruction",
                    value=DEFAULT_SIMILARITY_INSTRUCTION,
                    lines=2,
                )
            compare_button = gr.Button("Compare meaning", variant="primary")
            similarity_score = gr.Markdown()
            similarity_metadata = gr.JSON(label="Run details")
            compare_button.click(
                compare,
                inputs=[similarity_left, similarity_right, similarity_instruction],
                outputs=[similarity_score, similarity_metadata],
            )

        with gr.Tab("Embedding inspector", id="inspector"):
            inspect_input = gr.Textbox(
                label="Texts to embed",
                placeholder="One text per line. The table will show a compact vector preview.",
                lines=8,
                value="Semantic search finds relevant information.\n语义检索可以找到与问题相关的信息。",
            )
            inspect_button = gr.Button("Create embeddings", variant="primary")
            inspect_results = gr.Dataframe(
                headers=["#", "Text", "First 12 vector values"],
                datatype=["number", "str", "str"],
                interactive=False,
                wrap=True,
                label="Embedding preview",
            )
            inspect_metadata = gr.JSON(label="Run details")
            inspect_button.click(
                inspect_embeddings,
                inputs=inspect_input,
                outputs=[inspect_results, inspect_metadata],
            )

    gr.Markdown(
        """
        ---
        **About KaLM-Embedding-V2** — KaLM-Embedding-V2, a versatile and compact embedding model that achieves impressive performance in general-purpose text embedding tasks through systematic incentivization of advanced embedding techniques.
        [Model card](https://huggingface.co/KaLM-Embedding/KaLM-embedding-multilingual-mini-instruct-v2.5) ·
        [Homepage](https://kalm-embedding.github.io/)
        """
    )
    with gr.Accordion("Citation", open=True):
        gr.Markdown(
            """
            ```bibtex
            @inproceedings{
                zhao2026kalmembeddingv,
                title={Ka{LM}-Embedding-V2: Superior Training Techniques and Data Inspire A Versatile Embedding Model},
                author={Xinping Zhao and Xinshuo Hu and Zifei Shan and Shouzheng Huang and Yao Zhou and Xin Zhang and Zetian Sun and zhenyu liu and Dongfang Li and Xinyuan Wei and Youcheng Pan and Yang Xiang and Meishan Zhang and Haofen Wang and Jun Yu and Baotian Hu and Min Zhang},
                booktitle={The Fourteenth International Conference on Learning Representations},
                year={2026},
                url={https://openreview.net/forum?id=Y7qzhvWhcz}
            }
            ```
            """
        )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=20).launch()