import argparse
from pathlib import Path
import sqlite3

from sentence_transformers import CrossEncoder

from check_handbook import CASES
from handbook_search import DEFAULT_DB, search

MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
REVISION = "233902d25c440f23af6f7d6e94d2946bac0bee0a"


def rerank(model, question, candidates):
    if not candidates:
        return []
    pairs = [(question, passage.text) for passage in candidates]
    scores = model.predict(pairs, show_progress_bar=False)
    return sorted(zip(candidates, scores), key=lambda item: float(item[1]), reverse=True)


def print_results(model, database, question, mode, count):
    candidates = search(database, question, mode=mode, limit=count)
    ranked = rerank(model, question, candidates)
    print("Candidate order:", ", ".join(f"{hit.source} p.{hit.page}" for hit in candidates) or "none")
    print("Reranked:")
    for hit, score in ranked:
        print(f"  {score:.2f} | {hit.source}, page {hit.page}")


def evaluate(model, database, mode, count):
    retrieved = first_before = first_after = 0
    for question, expected in CASES:
        candidates = search(database, question, mode=mode, limit=count)
        reranked = rerank(model, question, candidates)
        before = [p.page if p.source == "library-handbook.pdf" else None for p in candidates]
        after = [p.page if p.source == "library-handbook.pdf" else None for p, _ in reranked]
        retrieved += expected in before
        first_before += bool(before and before[0] == expected)
        first_after += bool(after and after[0] == expected)
        print(f"Page {expected} | before {before} | after {after} | {question}")
    print(f"Candidate coverage: {retrieved}/{len(CASES)}")
    print(f"First result: {first_before}/{len(CASES)} before, {first_after}/{len(CASES)} after")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Rerank retrieved PDF passages on the CPU.")
    parser.add_argument("--db", type=Path, default=DEFAULT_DB)
    parser.add_argument("--mode", choices=("keyword", "vector", "hybrid"), default="hybrid")
    parser.add_argument("--candidates", type=int, choices=range(1, 9), default=3)
    parser.add_argument("--offline", action="store_true")
    parser.add_argument("question", nargs="?")
    args = parser.parse_args()
    try:
        model = CrossEncoder(MODEL, revision=REVISION, device="cpu", max_length=256,
                             local_files_only=args.offline, model_kwargs={"use_safetensors": True})
        if args.question:
            print_results(model, args.db, args.question, args.mode, args.candidates)
        else:
            evaluate(model, args.db, args.mode, args.candidates)
    except (OSError, RuntimeError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None