import gradio as gr
from sentence_transformers import CrossEncoder
import torch
import requests
import ast

# -------------------------------
# MODELS
# -------------------------------
CROSS_ENCODER_RERANK = "cross-encoder/ms-marco-MiniLM-L-12-v2"
JINA_MODEL = "jina-reranker-m0"
JINA_API_KEY = "jina_4075150fa702471c85ddea0a9ad4b306ouE7ymhrCpvxTxX3mScUv5LLDPKQ"
JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"
NV_MODEL = "NV-RerankQA-Mistral-4B-v3"   # Hugging Face hosted

# -------------------------------
# Load models
# -------------------------------
ce_rerank = CrossEncoder(CROSS_ENCODER_RERANK)

# -------------------------------
# Pipeline Function
# -------------------------------
def evaluate_models(query, docs_str):
    try:
        docs = ast.literal_eval(docs_str)
        assert isinstance(docs, list), "Input must be a Python list of strings"
    except Exception as e:
        return f"⚠️ Error parsing documents list: {e}"

    results = {}

    # 1. CrossEncoder reranker (MS MARCO)
    ce_rerank_scores = ce_rerank.predict([(query, d) for d in docs])
    ce_rerank_scores = [torch.sigmoid(torch.tensor(s)).item() for s in ce_rerank_scores]
    results["CrossEncoder (MS MARCO)"] = sorted(zip(docs, ce_rerank_scores), key=lambda x: x[1], reverse=True)

    # 2. Jina Reranker
    headers = {"Authorization": f"Bearer {JINA_API_KEY}", "Content-Type": "application/json"}
    payload = {"model": JINA_MODEL, "query": query, "documents": docs}
    try:
        r = requests.post(JINA_ENDPOINT, headers=headers, json=payload, timeout=30)
        r.raise_for_status()
        jina_scores = [res["relevance_score"] for res in r.json()["results"]]
        results["Jina Reranker"] = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
    except Exception as e:
        results["Jina Reranker"] = [(f"Error: {e}", 0)]

    # 3. NV RerankQA Mistral-4B-v3 (HF Inference API)
    try:
        hf_endpoint = f"https://api-inference.huggingface.co/models/{NV_MODEL}"
        headers = {"Authorization": f"Bearer YOUR_HF_API_KEY"}
        payload = {"inputs": {"query": query, "documents": docs}}
        r = requests.post(hf_endpoint, headers=headers, json=payload, timeout=60)
        r.raise_for_status()
        nv_scores = [res["score"] for res in r.json()]
        results["NV-RerankQA-Mistral-4B-v3"] = sorted(zip(docs, nv_scores), key=lambda x: x[1], reverse=True)
    except Exception as e:
        results["NV-RerankQA-Mistral-4B-v3"] = [(f"Error: {e}", 0)]

    # -------------------------------
    # Format output
    # -------------------------------
    out = ""
    for model_name, ranked in results.items():
        out += f"\n### {model_name}\n"
        for doc, score in ranked:
            out += f"- ({round(score,4)}) {doc}\n"
    return out

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 👑 Ranking Battle (3 Models)\nCompare **NV-RerankQA-Mistral-4B-v3**, **Jina**, and **CrossEncoder**.")

    query = gr.Textbox(label="Query", lines=2, placeholder="Enter your search query...")
    docs = gr.Textbox(
        label="Documents (Python list)", 
        lines=6, 
        placeholder='Example: ["Doc one text", "Doc two text", "Doc three text"]'
    )
    out = gr.Textbox(label="Ranked Results", lines=20)

    btn = gr.Button("Evaluate 🚀")
    btn.click(evaluate_models, inputs=[query, docs], outputs=out)

demo.launch()