from typing import Any

import gradio as gr
import spaces
from sentence_transformers import CrossEncoder


MODEL_ID = "ramitha2002/genieai-product-reranker"

# ZeroGPU requires CUDA placement at module level.
model = CrossEncoder(
    MODEL_ID,
    device="cuda",
    max_length=384,
)


def format_value(value: Any) -> str:
    if isinstance(value, list):
        return ", ".join(str(item) for item in value)
    return str(value)


def build_product_text(product: dict[str, Any]) -> str:
    fields = [
        ("Title", product.get("title") or product.get("name")),
        ("Description", product.get("description") or product.get("summary")),
        ("Features", product.get("features")),
        ("Brand", product.get("brand")),
        ("Color", product.get("color")),
        ("Category", product.get("category")),
    ]

    return "\n".join(
        f"{label}: {format_value(value)}"
        for label, value in fields
        if value is not None and value != ""
    )


@spaces.GPU(duration=30)
def rerank(
    query: str,
    products: list[dict[str, Any]],
    top_n: int,
) -> dict[str, Any]:
    query = query.strip()

    if not query:
        raise gr.Error("Query is required.")

    if not isinstance(products, list) or not products:
        raise gr.Error("Products must be a non-empty JSON array.")

    if len(products) > 30:
        raise gr.Error("Maximum 30 products per request.")

    pairs = [
        (query, build_product_text(product))
        for product in products
    ]

    scores = model.predict(
        pairs,
        batch_size=min(16, len(pairs)),
        show_progress_bar=False,
    )

    ranked = sorted(
        [
            {
                **product,
                "rerankerScore": float(score),
            }
            for product, score in zip(products, scores)
        ],
        key=lambda product: product["rerankerScore"],
        reverse=True,
    )

    return {
        "results": ranked[:max(1, min(int(top_n), len(ranked)))]
    }


sample_products = [
    {
        "id": "flowers-1",
        "name": "Pink Rose Bouquet",
        "description": "Fresh roses arranged for birthdays",
        "brand": "Bloom House",
        "color": "Pink"
    },
    {
        "id": "mouse-1",
        "name": "Wireless Gaming Mouse",
        "description": "RGB computer mouse",
        "brand": "GamePoint",
        "color": "Black"
    }
]


with gr.Blocks(title="GenieAI Product Reranker") as demo:
    gr.Markdown("# GenieAI Product Reranker")

    query_input = gr.Textbox(
        label="Search query",
        value="birthday flowers for mother",
    )

    products_input = gr.JSON(
        label="Products",
        value=sample_products,
    )

    top_n_input = gr.Slider(
        minimum=1,
        maximum=30,
        value=4,
        step=1,
        label="Number of results",
    )

    rerank_button = gr.Button("Rerank", variant="primary")
    output = gr.JSON(label="Ranked products")

    rerank_button.click(
        fn=rerank,
        inputs=[query_input, products_input, top_n_input],
        outputs=output,
        api_name="rerank",
    )

demo.queue(default_concurrency_limit=2).launch()