import gradio as gr
from sentence_transformers import CrossEncoder
import torch

import requests
import ast
import os

# -------------------------------
# MODELS
# -------------------------------
CROSS_ENCODER_RERANK = "cross-encoder/ms-marco-MiniLM-L-12-v2"
JINA_MODEL = "jina-reranker-m0"
JINA_API_KEY = os.getenv("JINA_API_KEY")  # set in HF Space settings
JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"
NV_MODEL = "NV-RerankQA-Mistral-4B-v3"
HF_API_KEY = os.getenv("HF_API_KEY")  # set in HF Space settings

# -------------------------------
# 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 {"Error": 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 = [round(torch.sigmoid(torch.tensor(s)).item(), 4) for s in ce_rerank_scores]
    results["CrossEncoder (MS MARCO)"] = ce_rerank_scores

    # 2. Jina Reranker
    if JINA_API_KEY:
        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 = [0] * len(docs)
            for res in r.json()["results"]:
                jina_scores[res["index"]] = round(res["relevance_score"], 4)
            results["Jina Reranker"] = jina_scores
        except Exception as e:
            results["Jina Reranker"] = [f"Error: {e}"]
    else:
        results["Jina Reranker"] = ["Error: Missing JINA_API_KEY"]

    # 3. NV RerankQA Mistral-4B-v3 (HF Inference API)
    if HF_API_KEY:
        try:
            hf_endpoint = f"https://api-inference.huggingface.co/models/{NV_MODEL}"
            headers = {"Authorization": f"Bearer {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 = [round(res["score"], 4) for res in r.json()]
            results["NV-RerankQA-Mistral-4B-v3"] = nv_scores
        except Exception as e:
            results["NV-RerankQA-Mistral-4B-v3"] = [f"Error: {e}"]
    else:
        results["NV-RerankQA-Mistral-4B-v3"] = ["Error: Missing HF_API_KEY"]

    return results

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 👑 Ranking Battle (Aligned Scores)\nOutputs only **scores aligned to input docs** from 3 models.")

    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.JSON(label="Model Scores")

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

demo.launch()