"""Compare multiple local LLMs as relevance-filter judges on the same arXiv results.

Companion to ``eval_relevance_prompt.py`` — that script varies the prompt
with one model held constant; this one varies the model with the prompt
held constant (the prompt currently on the PR).

Usage:
    LDR_BOOTSTRAP_ALLOW_UNENCRYPTED=true \\
    LDR_TEST_OLLAMA_BASE_URL=http://localhost:11434 \\
    pdm run python tests/performance/relevance_filter/eval_models.py

Override the model list via env var:
    LDR_TEST_OLLAMA_MODELS="qwen3:4b,qwen3.5:9b,gemma3:12b"

Fetches arXiv results once per query and runs the filter with each model.
Prints per-model keep counts and the consensus / disagreement breakdown so
a human can judge whether the prompt generalises across families and sizes.
"""

from __future__ import annotations

import os
import sys
import time
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"
)
ARXIV_MAX = int(os.environ.get("LDR_TEST_ARXIV_MAX", "40"))

# Default model set — picked for family / size diversity so we can see
# whether prompt quality transfers. Override via LDR_TEST_OLLAMA_MODELS.
DEFAULT_MODELS = [
    "qwen3:4b",  # tiny — test the floor
    "qwen3.5:9b",  # current default
    "gemma3:12b",  # Gemma family
    "ministral-3:14b",  # Mistral family
    "gpt-oss:20b",  # OpenAI-origin
    "qwen3.5:27b",  # larger qwen — does scaling help
]

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


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 installed_models(url: str) -> set:
    try:
        resp = safe_get(
            f"{url.rstrip('/')}/api/tags",
            timeout=3,
            allow_private_ips=True,
        )
        return {m["name"] for m in resp.json().get("models", [])}
    except requests.RequestException:
        return set()


def run_model(model: 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[{model}]",
        batch_size=10,
        max_parallel_batches=4,
    )


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

    print("\n=== Per-model keep counts ===")
    for name in kept_by_model:
        n = len(kept_titles[name])
        print(
            f"  {name:24s}  kept {n:3d} / {len(previews)}   ({timings[name]:.1f}s)"
        )

    if len(kept_by_model) < 2:
        return

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

    rejected_by_all = all_titles - set.union(*kept_titles.values())
    print(f"\n=== Rejected by ALL models: {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.items() if title in titles
        ]
        tag = ",".join(m.split(":")[0][:6] for m in kept_in)
        print(f"  [{tag:40s}] {title[:100]}")


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

    override = os.environ.get("LDR_TEST_OLLAMA_MODELS")
    if override:
        models = [m.strip() for m in override.split(",") if m.strip()]
    else:
        models = list(DEFAULT_MODELS)

    installed = installed_models(OLLAMA_URL)
    missing = [m for m in models if m not in installed]
    if missing:
        print(f"Skipping models not installed on {OLLAMA_URL}: {missing}")
        models = [m for m in models if m in installed]
    if not models:
        print("No models to test.", file=sys.stderr)
        return 1

    from langchain_ollama import ChatOllama

    print(f"Testing {len(models)} models: {models}")

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

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

        kept_by_model: Dict[str, List[Dict]] = {}
        timings: Dict[str, float] = {}

        for model in models:
            llm = ChatOllama(model=model, base_url=OLLAMA_URL, num_ctx=8192)
            t0 = time.monotonic()
            try:
                kept = run_model(model, llm, previews, query)
            except Exception as exc:
                print(f"  {model:24s}: ERROR {exc!r}")
                continue
            elapsed = time.monotonic() - t0
            kept_by_model[model] = kept
            timings[model] = elapsed
            print(f"  {model:24s}: kept {len(kept):3d} in {elapsed:.1f}s")

        summarise(kept_by_model, previews, timings)

    return 0


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