"""Reranking via Qwen3-Reranker over Ollama."""
import requests
import os
import re
from .config import CONFIG

RERANK_MODEL = os.getenv("KB_RERANK_MODEL", "dengcao/Qwen3-Reranker-4B:Q5_K_M")


def score_relevance(query: str, document: str) -> float:
    """Get a relevance score in [0, 1] for (query, document) pair."""
    if len(document) > 4000:
        document = document[:4000] + "..."
    
    # Pre-fill <think></think> as empty so the model jumps straight to answer
    prompt = (
        f"<|im_start|>system\n"
        f"You evaluate whether a document is relevant to a query about quantitative trading. "
        f"Be strict: the document must specifically address the query. "
        f"Output only \"yes\" or \"no\".<|im_end|>\n"
        f"<|im_start|>user\n"
        f"<Query>: {query}\n"
        f"<Document>: {document}<|im_end|>\n"
        f"<|im_start|>assistant\n"
        f"<think>\n\n</think>\n\n"
    )
    
    try:
        resp = requests.post(
            f"{CONFIG.ollama_url}/api/generate",
            json={
                "model": RERANK_MODEL,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "num_predict": 5,
                    "temperature": 0.0,
                    "top_p": 1.0,
                },
                "raw": True,
            },
            timeout=60,
        )
        resp.raise_for_status()
        text = resp.json().get("response", "").strip().lower()
        text = re.sub(r"</?think>", "", text).strip()
        resp.raise_for_status()
        text = resp.json().get("response", "").strip()
        text = re.sub(r"</?think>", "", text).strip()
        
        # DEBUG: print every response so we can see what's happening
        print(f"DEBUG rerank response: {text[:80]!r}", flush=True)
        
        text = text.lower()
        
        if text.startswith("yes"):
            return 1.0
        elif text.startswith("no"):
            return 0.0
        else:
            print(f"Rerank ambiguous response: {text[:100]!r}", flush=True)
            return 0.5
    except Exception as e:
        print(f"Rerank error: {e}", flush=True)
        return 0.5


def rerank(query: str, candidates: list, top_k: int = 10) -> list:
    if not candidates:
        return []
    
    scored = []
    for cand in candidates:
        chunk_meta = cand.get("chunk_metadata") or {}
        doc_text = chunk_meta.get("raw_chunk") or cand.get("content", "")
        score = score_relevance(query, doc_text)
        cand_with_score = dict(cand)
        cand_with_score["rerank_score"] = score
        cand_with_score["original_similarity"] = cand.get("similarity")
        scored.append(cand_with_score)
    
    scored.sort(key=lambda x: x["rerank_score"], reverse=True)
    return scored[:top_k]