"""YouRunAI: compare retrieved passages before and after local reranking."""
import json
from pathlib import Path
import gradio as gr
from sentence_transformers import CrossEncoder

model = CrossEncoder("Qwen/Qwen3-Reranker-0.6B", device="cpu", max_length=2048)
out = Path("outputs")
out.mkdir(exist_ok=True)
sample = {"query": "How do I return an item?", "candidates": [
    {"file": "shipping.md", "text": "Cedar Shop ships weekdays. Standard shipping takes three to five business days."},
    {"file": "returns.md", "text": "Cedar Shop accepts unused items within 30 days. Contact support with the order number. Refunds go to the original payment method."},
    {"file": "hours.md", "text": "Cedar Shop support is open Monday to Friday, 9 am to 5 pm."}]}

def rerank(upload):
    try:
        data = json.loads(Path(upload).read_text(encoding="utf-8")) if upload else sample
        query, candidates = data["query"], data["candidates"]
        if not isinstance(query, str) or not query.strip() or not 1 <= len(candidates) <= 20:
            raise ValueError("Use a nonempty query and between 1 and 20 candidates.")
        for row in candidates:
            if not isinstance(row.get("text"), str) or not row["text"].strip() or len(row["text"]) > 8000:
                raise ValueError("Each candidate needs text of 1 to 8,000 characters.")
        scores = model.predict([(query, row["text"]) for row in candidates], batch_size=1)
        ranked = sorted([{**row, "before_rank": i + 1, "rerank_score": float(score)} for i, (row, score) in enumerate(zip(candidates, scores))], key=lambda row: row["rerank_score"], reverse=True)
        path = out / "reranked.json"
        path.write_text(json.dumps({"query": query, "candidates": ranked}, indent=2, ensure_ascii=False), encoding="utf-8")
        before = [[i + 1, r.get("file", "passage"), r["text"]] for i, r in enumerate(candidates)]
        after = [[i + 1, r["before_rank"], r.get("file", "passage"), round(r["rerank_score"], 4), r["text"]] for i, r in enumerate(ranked)]
        return query, before, after, str(path)
    except Exception as error:
        raise gr.Error(f"Cannot rerank this input: {error}") from error

with gr.Blocks(title="Rerank local search results") as app:
    gr.Markdown("# Make search results more relevant\nUpload candidates.json from the private-search guide, or leave empty for the supplied three-passage example. Reranking reorders existing results; it cannot find a missing document.")
    upload = gr.File(label="Optional candidates.json", file_types=[".json"], type="filepath")
    button = gr.Button("Rerank results", variant="primary")
    query = gr.Textbox(label="Query", interactive=False)
    before = gr.Dataframe(headers=["Before rank", "File", "Passage"], interactive=False, wrap=True)
    after = gr.Dataframe(headers=["After rank", "Before rank", "File", "Relevance score", "Passage"], interactive=False, wrap=True)
    download = gr.File(label="Download reranked.json")
    button.click(rerank, inputs=upload, outputs=[query, before, after, download])
app.launch(server_name="127.0.0.1", server_port=7861, share=False)