import gradio as gr
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import torch
import requests
import ast

# -------------------------------
# MODELS
# -------------------------------
BI_ENCODER = "sentence-transformers/all-MiniLM-L6-v2"
CROSS_ENCODER_RERANK = "cross-encoder/ms-marco-MiniLM-L-12-v2"
CROSS_ENCODER_STS = "cross-encoder/stsb-roberta-large"
CROSS_ENCODER_NLI = "cross-encoder/nli-deberta-v3-base"
JINA_MODEL = "jina-reranker-m0"
JINA_API_KEY = "jina_4075150fa702471c85ddea0a9ad4b306ouE7ymhrCpvxTxX3mScUv5LLDPKQ"
JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"

# -------------------------------
# Load models
# -------------------------------
bi_encoder = SentenceTransformer(BI_ENCODER)
ce_rerank = CrossEncoder(CROSS_ENCODER_RERANK)
ce_sts = CrossEncoder(CROSS_ENCODER_STS)
ce_nli = CrossEncoder(CROSS_ENCODER_NLI, num_labels=3)

# -------------------------------
# Pipeline Function
# -------------------------------
def evaluate_models(query, docs_str):
    try:
        # Parse docs string as Python list
        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. Bi-encoder cosine similarity
    query_emb = bi_encoder.encode(query, convert_to_tensor=True)
    doc_embs = bi_encoder.encode(docs, convert_to_tensor=True)
    cos_scores = util.cos_sim(query_emb, doc_embs)[0].cpu().tolist()
    results["1. Bi-encoder similarity"] = sorted(zip(docs, cos_scores), key=lambda x: x[1], reverse=True)

    # 2. 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["2. CrossEncoder Reranker (MS MARCO)"] = sorted(zip(docs, ce_rerank_scores), key=lambda x: x[1], reverse=True)

    # 3. 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["3. Jina Reranker"] = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
    except Exception as e:
        results["3. Jina Reranker"] = [(f"Error: {e}", 0)]

    # 4. CrossEncoder STS
    ce_sts_scores = ce_sts.predict([(query, d) for d in docs])
    results["4. CrossEncoder STS"] = sorted(zip(docs, ce_sts_scores), key=lambda x: x[1], reverse=True)

    # 5. CrossEncoder NLI
    ce_nli_probs = ce_nli.predict([(query, d) for d in docs], apply_softmax=True)
    ce_nli_scores = [float(p[1] + p[2]) for p in ce_nli_probs]  # neutral + entailment
    results["5. CrossEncoder NLI"] = sorted(zip(docs, ce_nli_scores), key=lambda x: x[1], reverse=True)

    # 6. Bi-encoder raw similarity (duplicate for clarity)
    results["6. Bi-encoder baseline"] = sorted(zip(docs, cos_scores), key=lambda x: x[1], reverse=True)

    # -------------------------------
    # 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("## 🔎 Multi-Model Reranker (HF + Jina)\nPass a **query** and a **list of documents (Python list of strings)**.")

    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()