"""KaLM-Reranker-V1-Small interactive demo for Hugging Face Spaces."""

from __future__ import annotations

import importlib.util
import math
from functools import lru_cache
from typing import Any

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download


MODEL_ID = "KaLM-Embedding/KaLM-Reranker-V1-Small"
DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
COMPRESSION_FACTORS = (1, 2, 4, 8, 16, 32)
MAX_DOCUMENTS = 50


def load_reranker_class() -> type:
    """Fetch the official lightweight inference wrapper from the model repository."""
    module_path = hf_hub_download(repo_id=MODEL_ID, filename="kalm_reranker.py")
    spec = importlib.util.spec_from_file_location("kalm_reranker", module_path)
    if spec is None or spec.loader is None:
        raise RuntimeError("Unable to load the official KaLM reranker wrapper.")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.KaLMReranker


@lru_cache(maxsize=1)
def get_reranker() -> Any:
    """Load the 1B Small reranker once for this Space process."""
    reranker_class = load_reranker_class()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = "bfloat16" if device == "cuda" else "float32"
    return reranker_class(
        MODEL_ID,
        device=device,
        dtype=dtype,
        batch_size=4,
        query_max_length=512,
        max_length=1024,
        chunk_size=4,
    )


def normalize_factor(value: int | float | str) -> int:
    try:
        factor = int(value)
    except (TypeError, ValueError) as error:
        raise gr.Error("Choose a valid compression factor.") from error
    if factor not in COMPRESSION_FACTORS:
        raise gr.Error("Compression must be one of 1, 2, 4, 8, 16, or 32.")
    return factor


def configure_reranker(compression: int | float | str) -> tuple[Any, int]:
    factor = normalize_factor(compression)
    reranker = get_reranker()
    # The official chunk_size pools every N encoder token states. Larger N means
    # fewer passage representations passed to the decoder, i.e. more compression.
    reranker.chunk_size = factor
    return reranker, factor


def parse_documents(raw_documents: str) -> list[str]:
    documents = [line.strip() for line in (raw_documents or "").splitlines() if line.strip()]
    if not documents:
        raise gr.Error("Please enter at least one document.")
    if len(documents) > MAX_DOCUMENTS:
        raise gr.Error(f"Please limit each run to {MAX_DOCUMENTS} documents.")
    return documents


def validate_query(query: str) -> str:
    query = (query or "").strip()
    if not query:
        raise gr.Error("Please enter a query.")
    return query


def estimated_tokens(reranker: Any, document: str) -> int:
    return len(
        reranker.tokenizer(
            f"<Document>: {document}",
            add_special_tokens=False,
            truncation=True,
            max_length=reranker.max_length,
        )["input_ids"]
    )


@spaces.GPU(duration=90)
def rerank_documents(
    query: str,
    raw_documents: str,
    instruction: str,
    compression: int,
    top_k: int,
) -> tuple[list[list[Any]], dict[str, Any]]:
    query = validate_query(query)
    documents = parse_documents(raw_documents)
    instruction = (instruction or DEFAULT_INSTRUCTION).strip() or DEFAULT_INSTRUCTION
    reranker, factor = configure_reranker(compression)
    rankings = reranker.rank(query, documents, instruction=instruction, top_k=int(top_k))

    rows = [
        [
            rank,
            f"{float(item['score']):.4f}",
            documents[int(item["corpus_id"])],
        ]
        for rank, item in enumerate(rankings, start=1)
    ]
    raw_token_count = sum(estimated_tokens(reranker, document) for document in documents)
    return rows, {
        "model": MODEL_ID,
        "score": "P(yes): probability that the document satisfies the query and instruction",
        "documents_reranked": len(documents),
        "compression_factor": f"{factor}×",
        "estimated_encoder_tokens": raw_token_count,
        "estimated_tokens_after_pooling": math.ceil(raw_token_count / factor),
        "device": "CUDA" if torch.cuda.is_available() else "CPU",
    }


@spaces.GPU(duration=90)
def score_pair(
    query: str,
    document: str,
    instruction: str,
    compression: int,
) -> tuple[str, dict[str, Any]]:
    query = validate_query(query)
    document = (document or "").strip()
    if not document:
        raise gr.Error("Please enter a document.")
    instruction = (instruction or DEFAULT_INSTRUCTION).strip() or DEFAULT_INSTRUCTION
    reranker, factor = configure_reranker(compression)
    score = float(reranker.predict([(query, document)], instruction=instruction)[0])
    token_count = estimated_tokens(reranker, document)
    return f"## Relevance score: **{score:.4f}**\n\nThis is the model's probability that the answer is **yes**.", {
        "model": MODEL_ID,
        "compression_factor": f"{factor}×",
        "encoder_tokens": token_count,
        "estimated_tokens_after_pooling": math.ceil(token_count / factor),
        "instruction": instruction,
    }


@spaces.GPU(duration=90)
def compare_compression(
    query: str,
    document: str,
    instruction: str,
    factors: list[int] | None,
) -> tuple[list[list[Any]], dict[str, Any]]:
    query = validate_query(query)
    document = (document or "").strip()
    if not document:
        raise gr.Error("Please enter a document.")
    if not factors:
        raise gr.Error("Choose at least one compression factor.")

    instruction = (instruction or DEFAULT_INSTRUCTION).strip() or DEFAULT_INSTRUCTION
    reranker = get_reranker()
    token_count = estimated_tokens(reranker, document)
    selected_factors = sorted({normalize_factor(factor) for factor in factors})
    rows: list[list[Any]] = []
    for factor in selected_factors:
        reranker.chunk_size = factor
        score = float(reranker.predict([(query, document)], instruction=instruction)[0])
        rows.append(
            [
                f"{factor}×",
                token_count,
                math.ceil(token_count / factor),
                f"{score:.4f}",
            ]
        )
    return rows, {
        "how_to_read": "Larger factors pool more encoder token states into each chunk. This reduces the document representation length passed to the decoder.",
        "model": MODEL_ID,
        "instruction": instruction,
    }


EXAMPLE_QUERY = "What is the capital of China?"
EXAMPLE_DOCUMENTS = """The capital of China is Beijing.
Paris is the capital 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 #4f46e5; padding: 0.65rem 0.9rem; background: #eef2ff; border-radius: 0.4rem; }
"""


with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="KaLM Reranker Demo") as demo:
    gr.Markdown(
        """
        <div id="hero">
          <h1>KaLM Reranker — Compressed Document Reranking</h1>
          <p>Rerank documents efficiently with <code>KaLM-Embedding/KaLM-Reranker-V1-Small</code>.</p>
        </div>
        """
    )
    gr.HTML(
        "<div class='notice'><strong>Compression factor</strong> controls the official encoder chunk pooling parameter. 1× preserves all encoder states; larger factors reduce decoder-side passage length and can improve throughput.</div>"
    )

    with gr.Tabs():
        with gr.Tab("Semantic Reranking", id="rerank"):
            with gr.Row():
                query_input = gr.Textbox(label="Query / 查询", lines=2, scale=2)
                top_k = gr.Slider(1, 10, value=3, step=1, label="Results to show", scale=1)
            documents_input = gr.Textbox(
                label="Documents / 文档",
                placeholder="One document per line. The model will rank them by relevance.",
                lines=10,
            )
            with gr.Accordion("Task and compression settings", open=True):
                rerank_instruction = gr.Textbox(
                    label="Task instruction",
                    value=DEFAULT_INSTRUCTION,
                    lines=2,
                )
                rerank_compression = gr.Dropdown(
                    choices=list(COMPRESSION_FACTORS),
                    value=4,
                    label="Compression factor (encoder chunk pooling)",
                    info="1× = no sequence-length reduction; 32× = strongest pooling.",
                )
            rerank_button = gr.Button("Rerank documents", variant="primary")
            rerank_results = gr.Dataframe(
                headers=["Rank", "P(yes)", "Document"],
                datatype=["number", "str", "str"],
                interactive=False,
                wrap=True,
                label="Reranked results",
            )
            rerank_metadata = gr.JSON(label="Run details")
            rerank_button.click(
                rerank_documents,
                inputs=[query_input, documents_input, rerank_instruction, rerank_compression, top_k],
                outputs=[rerank_results, rerank_metadata],
            )
            gr.Examples(
                examples=[[EXAMPLE_QUERY, EXAMPLE_DOCUMENTS, DEFAULT_INSTRUCTION, 4, 3]],
                inputs=[query_input, documents_input, rerank_instruction, rerank_compression, top_k],
                label="Try an example",
            )

        with gr.Tab("Pair Relevance", id="pair-score"):
            with gr.Row():
                pair_query = gr.Textbox(label="Query", lines=5, value=EXAMPLE_QUERY)
                pair_document = gr.Textbox(label="Document", lines=5, value="The capital of China is Beijing.")
            pair_instruction = gr.Textbox(label="Task instruction", value=DEFAULT_INSTRUCTION, lines=2)
            pair_compression = gr.Dropdown(
                choices=list(COMPRESSION_FACTORS),
                value=4,
                label="Compression factor",
            )
            pair_button = gr.Button("Score relevance", variant="primary")
            pair_score = gr.Markdown()
            pair_metadata = gr.JSON(label="Run details")
            pair_button.click(
                score_pair,
                inputs=[pair_query, pair_document, pair_instruction, pair_compression],
                outputs=[pair_score, pair_metadata],
            )

        with gr.Tab("Compression Explorer", id="compression"):
            gr.Markdown(
                "Compare the same query-document pair across pooling factors. Larger factors shorten the encoded document representation before cross-attention."
            )
            compression_query = gr.Textbox(label="Query", lines=3, value=EXAMPLE_QUERY)
            compression_document = gr.Textbox(
                label="Document",
                lines=5,
                value=(
                    "Beijing, also known as Peking, is the capital of China and one of the most populous "
                    "cities in the world. It is the country’s political, educational, and cultural center, "
                    "housing the headquarters of most of China’s largest state-owned companies. It is a "
                    "significant hub for the national highway, expressway, railway, and high-speed rail networks."
                ),
            )
            compression_instruction = gr.Textbox(label="Task instruction", value=DEFAULT_INSTRUCTION, lines=2)
            comparison_factors = gr.CheckboxGroup(
                choices=list(COMPRESSION_FACTORS),
                value=[1, 4, 16, 32],
                label="Compression factors to compare",
            )
            comparison_button = gr.Button("Compare compression", variant="primary")
            comparison_results = gr.Dataframe(
                headers=["Compression", "Encoder tokens", "After pooling", "P(yes)"],
                datatype=["str", "number", "number", "str"],
                interactive=False,
                label="Compression comparison",
            )
            comparison_metadata = gr.JSON(label="How compression works")
            comparison_button.click(
                compare_compression,
                inputs=[compression_query, compression_document, compression_instruction, comparison_factors],
                outputs=[comparison_results, comparison_metadata],
            )

    gr.Markdown(
        """
        ---
        **About KaLM-Reranker-V1** — A fast but not late-interaction reranker that pre-encodes passages and uses cross-attention to model fine-grained query-document relevance.
        [Model card](https://huggingface.co/KaLM-Embedding/KaLM-Reranker-V1-Small) ·
        [Homepage](https://kalm-embedding.github.io/)
        """
    )
    with gr.Accordion("Citation", open=True):
        gr.Markdown(
            """
            ```bibtex
            @misc{zhao2026kalmrerankerv1,
                  title={KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking}, 
                  author={Xinping Zhao and Jiaxin Xu and Ziqi Dai and Xin Zhang and Shouzheng Huang and Danyu Tang and Xinshuo Hu and Meishan Zhang and Baotian Hu and Min Zhang},
                  year={2026},
                  eprint={2606.22807},
                  archivePrefix={arXiv},
                  primaryClass={cs.CL},
                  url={https://arxiv.org/abs/2606.22807}, 
            }
            ```
            """
        )


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