"""Compare multiple relevance-filter prompt variants on the same arXiv results.

Usage:
    LDR_BOOTSTRAP_ALLOW_UNENCRYPTED=true \\
    LDR_TEST_OLLAMA_BASE_URL=http://localhost:11434 \\
    LDR_TEST_OLLAMA_MODEL=qwen3.5:9b \\
    pdm run python tests/performance/relevance_filter/eval_prompt.py

Fetches arXiv results ONCE per query and then runs the filter with each
candidate prompt variant. Prints a side-by-side comparison so a human can
judge which variant keeps the right papers and drops the right ones.

This is a dev tool — not a pytest test — because the right answer is a
human judgment, not a hard assertion.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Dict, List

import arxiv
import requests

PROJECT_ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from local_deep_research.security import safe_get  # noqa: E402
from local_deep_research.web_search_engines.relevance_filter import (  # noqa: E402
    filter_previews_for_relevance,
)

OLLAMA_URL = os.environ.get(
    "LDR_TEST_OLLAMA_BASE_URL", "http://localhost:11434"
)
OLLAMA_MODEL = os.environ.get("LDR_TEST_OLLAMA_MODEL", "qwen3.5:9b")
ARXIV_MAX = int(os.environ.get("LDR_TEST_ARXIV_MAX", "40"))


QUERIES = [
    "LLM interpretability latest research",
    "sparse autoencoders for mechanistic interpretability",
]


# The shared scaffolding around the guidance block. The variants below only
# swap the middle paragraph — query block, date, results block, and output
# contract stay identical so the comparison is clean.
def _make_template(guidance: str) -> str:
    return (
        "This is a relevance-filtering step. Kept results move forward and may be used in the final answer; dropped ones are excluded from further processing.\n\n"
        'Query: "{query}"\n'
        "Current date: {current_date}\n\n"
        "Search results:\n"
        "{preview_text}\n\n"
        f"{guidance}\n\n"
        "Output ONLY the 0-based indices of relevant results as a comma-separated list, nothing else.\n"
        "Example: 0, 2, 5"
    )


VARIANTS: Dict[str, str] = {
    # V0 — the prompt currently shipping on main after the selectivity-
    # bias fix. Run against this as the baseline to validate new variants.
    "V0_current": _make_template(
        "Direct topic match matters more than keyword match — results that just mention the query terms usually don't help."
    ),
    # V1 — explicit "prefer smaller" selectivity bias (the version that
    # was on an earlier commit of this PR before we dropped it). Kept as
    # a variant so we can quantify the regression if someone re-adds it.
    "V1_prefer_smaller": _make_template(
        "Direct topic match matters more than keyword match — results that just mention the query terms usually don't help. "
        "Prefer a smaller high-confidence selection over a broader one."
    ),
    # V2 — explicitly allow related/adjacent work. Targets the observed
    # false negatives on mech-interp papers that share only one of two query terms.
    "V2_inclusive_adjacent": _make_template(
        "Keep results whose main subject is the query topic or closely relates to it. "
        "Reject results that share keywords with the query but whose main subject is clearly something else."
    ),
    # V3 — flip the framing to emphasise what to keep rather than what to reject.
    "V3_keep_framing": _make_template(
        "A result is relevant if reading it would help someone answering the query. "
        "Borderline-relevant work is often worth keeping; reject only when the main subject is clearly elsewhere."
    ),
    # V4 — terser, closer to the historical "MUST directly address" line.
    "V4_terse": _make_template(
        "Keep results whose main subject directly addresses the query. "
        "Mere keyword overlap with off-topic work is not relevant."
    ),
}


def fetch_previews(query: str, max_results: int) -> List[Dict]:
    client = arxiv.Client(page_size=max_results)
    search = arxiv.Search(
        query=query,
        max_results=max_results,
        sort_by=arxiv.SortCriterion.Relevance,
    )
    return [
        {
            "title": p.title.strip(),
            "snippet": (p.summary or "").strip(),
            "url": p.entry_id,
        }
        for p in client.results(search)
    ]


def check_ollama(url: str) -> bool:
    try:
        resp = safe_get(
            f"{url.rstrip('/')}/api/tags",
            timeout=3,
            allow_private_ips=True,
        )
        return resp.status_code == 200
    except requests.RequestException:
        return False


def run_one(
    variant_name: str, template: str, llm, previews, query
) -> List[Dict]:
    return filter_previews_for_relevance(
        llm=llm,
        previews=previews,
        query=query,
        max_filtered_results=None,
        engine_name=f"Eval[{variant_name}]",
        batch_size=10,
        max_parallel_batches=4,
        prompt_template=template,
    )


def summarise(
    kept_by_variant: Dict[str, List[Dict]], previews: List[Dict]
) -> None:
    all_titles = {p["title"] for p in previews}
    kept_titles_by_variant = {
        name: {p["title"] for p in kept}
        for name, kept in kept_by_variant.items()
    }

    print("\n=== Per-variant counts ===")
    # Iterate over the collected results rather than VARIANTS so a
    # variant that raised (and was skipped in main) doesn't KeyError.
    for name, titles in kept_titles_by_variant.items():
        print(f"  {name:24s}  kept {len(titles):3d} / {len(previews)}")

    if len(kept_titles_by_variant) < 2:
        return

    consensus = set.intersection(*kept_titles_by_variant.values())
    print(f"\n=== Consensus (kept by ALL variants): {len(consensus)} ===")
    for t in sorted(consensus):
        print(f"  = {t[:110]}")

    rejected_by_all = all_titles - set.union(*kept_titles_by_variant.values())
    print(f"\n=== Rejected by ALL variants: {len(rejected_by_all)} ===")
    for t in sorted(rejected_by_all):
        print(f"  × {t[:110]}")

    print("\n=== Disagreements (kept by some, rejected by others) ===")
    for title in sorted(all_titles - consensus - rejected_by_all):
        kept_in = [
            name
            for name, titles in kept_titles_by_variant.items()
            if title in titles
        ]
        print(
            f"  [{','.join(v.split('_')[0] for v in kept_in):14s}] {title[:100]}"
        )


def main() -> int:
    if not check_ollama(OLLAMA_URL):
        print(f"Ollama endpoint {OLLAMA_URL} not reachable.", file=sys.stderr)
        return 1

    from langchain_ollama import ChatOllama

    llm = ChatOllama(model=OLLAMA_MODEL, base_url=OLLAMA_URL, num_ctx=8192)

    for query in QUERIES:
        print("\n" + "=" * 78)
        print(f"QUERY: {query!r}")
        print(f"Judge: {OLLAMA_MODEL} @ {OLLAMA_URL}")
        print("=" * 78)

        previews = fetch_previews(query, ARXIV_MAX)
        print(f"Fetched {len(previews)} arXiv previews.\n")

        kept_by_variant: Dict[str, List[Dict]] = {}
        for name, template in VARIANTS.items():
            try:
                kept = run_one(name, template, llm, previews, query)
            except Exception as exc:
                print(f"  {name:24s}: ERROR {exc!r}")
                continue
            kept_by_variant[name] = kept
            print(f"  {name:24s}: kept {len(kept)}")

        summarise(kept_by_variant, previews)

    return 0


if __name__ == "__main__":
    sys.exit(main())
