#!/usr/bin/env python3
"""
Probe the journal reputation filter against live arXiv results.

Runs the *full* ArXivSearchEngine pipeline (including the journal filter
as a preview filter) and prints which papers passed/were dropped, along
with the journal_quality score the filter attached.

Usage:
    python scripts/dev/test_journal_filter_arxiv.py "graph neural networks" -n 25
"""

import argparse
import sys
from unittest.mock import Mock

from local_deep_research.utilities.thread_context import search_context
from local_deep_research.web_search_engines.engines.search_engine_arxiv import (
    ArXivSearchEngine,
)
from local_deep_research.web_search_engines.engines.search_engine_openalex import (
    OpenAlexSearchEngine,
)
from local_deep_research.web_search_engines.engines.search_engine_pubmed import (
    PubMedSearchEngine,
)
from local_deep_research.web_search_engines.engines.search_engine_semantic_scholar import (
    SemanticScholarSearchEngine,
)

ENGINES = {
    "arxiv": ArXivSearchEngine,
    "openalex": OpenAlexSearchEngine,
    "pubmed": PubMedSearchEngine,
    "s2": SemanticScholarSearchEngine,
}


DOMAIN_QUERIES = {
    "fusion": "tokamak plasma confinement",
    "llm": "large language model alignment",
    "graph_nn": "graph neural networks",
    "astro": "exoplanet atmosphere spectroscopy",
    "biomed": "CRISPR gene editing therapy",
    "condmat": "high temperature superconductor cuprate",
    "climate": "climate model ocean heat content",
    "quantum": "quantum error correction surface code",
    "math": "riemann hypothesis zeta function",
    "robotics": "reinforcement learning robotic manipulation",
}


def run_one(
    query: str,
    n: int,
    threshold: int,
    snapshot: dict,
    engine_name: str = "arxiv",
) -> dict:
    cls = ENGINES[engine_name]
    kwargs = {"max_results": n, "llm": Mock(), "settings_snapshot": snapshot}
    # PubMed/S2 use the LLM for query optimization — Mock would corrupt
    # the optimized query string. Disable optimization for the probe so
    # we send the raw query verbatim.
    if engine_name in ("pubmed", "s2"):
        kwargs["optimize_queries"] = False
    with search_context({"username": None, "user_password": None}):
        try:
            engine = cls(**kwargs)
            out = engine.run(query)
        except Exception as e:
            return {
                "query": query,
                "n_total": 0,
                "n_no_jref": 0,
                "n_scored": 0,
                "avg_q": 0,
                "high": 0,
                "passthrough": 0,
                "results": [],
                "error": str(e)[:80],
            }

    scored = [r for r in out if "journal_quality" in r]
    no_jref = [r for r in out if not r.get("journal_ref")]
    qualities = [r["journal_quality"] for r in scored]
    return {
        "query": query,
        "n_total": len(out),
        "n_no_jref": len(no_jref),
        "n_scored": len(scored),
        "avg_q": sum(qualities) / len(qualities) if qualities else 0,
        "high": sum(1 for q in qualities if q >= 7),
        "passthrough": sum(1 for q in qualities if q == threshold),
        "results": out,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "query", nargs="?", help="arXiv search query (omit to run dataset)"
    )
    ap.add_argument("-n", type=int, default=20, help="number of papers")
    ap.add_argument("-t", "--threshold", type=int, default=4)
    ap.add_argument(
        "--dataset",
        action="store_true",
        help="run the multi-domain dataset of queries",
    )
    ap.add_argument(
        "--engines",
        default="arxiv",
        help="comma-separated list: arxiv,openalex,pubmed,s2 (or 'all')",
    )
    args = ap.parse_args()
    if args.engines == "all":
        engine_names = list(ENGINES)
    else:
        engine_names = [e.strip() for e in args.engines.split(",")]

    # Minimal settings snapshot — defaults are returned for missing keys.
    # Threshold and arxiv journal-filter enable flag are wired explicitly.
    snapshot = {
        "search.journal_reputation.threshold": args.threshold,
        "search.journal_reputation.exclude_non_published": False,
        "search.journal_reputation.max_context": 3000,
        "search.journal_reputation.reanalysis_period": 365,
        "search.engine.web.arxiv.journal_reputation.enabled": True,
        "search.engine.web.openalex.journal_reputation.enabled": True,
        "search.engine.web.pubmed.journal_reputation.enabled": True,
        "search.engine.web.semantic_scholar.journal_reputation.enabled": True,
    }

    if args.dataset:
        print(
            f"Running dataset of {len(DOMAIN_QUERIES)} domain queries × "
            f"{len(engine_names)} engines (n={args.n}, threshold={args.threshold})\n"
        )
        all_rows = []
        for engine_name in engine_names:
            for domain, q in DOMAIN_QUERIES.items():
                print(
                    f"  • {engine_name:<9} {domain:<10} {q!r} ...", flush=True
                )
                s = run_one(q, args.n, args.threshold, snapshot, engine_name)
                all_rows.append((engine_name, domain, s))

        print(
            f"\n{'engine':<9} {'domain':<10} {'tot':>4} {'noJ':>4} "
            f"{'scor':>4} {'avg':>5} {'≥7':>3} {'pass':>4}  err"
        )
        print("-" * 80)
        for engine_name, domain, s in all_rows:
            err = s.get("error", "")
            print(
                f"{engine_name:<9} {domain:<10} "
                f"{s['n_total']:>4} {s['n_no_jref']:>4} {s['n_scored']:>4} "
                f"{s['avg_q']:>5.1f} {s['high']:>3} {s['passthrough']:>4}  {err}"
            )
        return 0

    if not args.query:
        ap.error("query required (or use --dataset)")

    with search_context({"username": None, "user_password": None}):
        engine = ArXivSearchEngine(
            max_results=args.n, llm=Mock(), settings_snapshot=snapshot
        )
        print(f"Running ArXivSearchEngine for: {args.query!r}\n")
        out = engine.run(args.query)

    print(f"{'#':<3} {'qual':<5} {'journal_ref':<45} title")
    print("-" * 120)
    for i, r in enumerate(out, 1):
        q = r.get("journal_quality", "—")
        jref = (r.get("journal_ref") or "—")[:43]
        title = (r.get("title") or "")[:60]
        print(f"{i:<3} {str(q):<5} {jref:<45} {title}")

    print(
        f"\nReturned {len(out)} results from pipeline "
        f"(threshold={args.threshold})."
    )
    return 0


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