#!/usr/bin/env python3
"""
search.py -- command-line interface for the arxiv-complete retrieval layer.

Examples
--------
    # semantic search (uses the local hashing embedder unless an API key is set)
    python search.py --query "speculative decoding"

    # restrict to recent papers in a category
    python search.py --query "neural machine translation" --top-k 15 --year-min 2018 --category cs.CL

    # only survey / overview / taxonomy papers
    python search.py --query "graph neural networks" --survey-only

    # force the deterministic keyword fallback (no embeddings at all)
    python search.py --query "speculative decoding" --keyword-only

Build the ANN index first (once) with the build script:

    python build_index.py
"""

from __future__ import annotations

import argparse
import os
import sys

# Allow running as a script from the repo root or as a module.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from data import DataStore  # noqa: E402
from embed import HashingEmbedder, get_embedder  # noqa: E402
from query import QueryOptions, Result, query  # noqa: E402


def _print_results(results: list[Result], query_text: str, method_note: str) -> None:
    print(f"\nQuery: {query_text!r}")
    print(f"Index: {method_note}  ({len(results)} result(s))\n")
    if not results:
        print("  (no matches)")
        return
    for i, r in enumerate(results, 1):
        print(f"{i:>2}. [{r.method}] {r.score:.3f}  {r.title}")
        meta = f"   {r.category} · {r.year} · {r.paper_id}"
        print(meta)
        print(f"   {r.url}")
        print(f"   {_snippet_lines(r.snippet)}")
        print()


def _snippet_lines(snippet: str) -> str:
    snippet = snippet.strip()
    if not snippet:
        snippet = "(no abstract available)"
    width = 88
    if len(snippet) <= width:
        return "   " + snippet
    return "   " + snippet[: width - 3] + "…"


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="search.py",
        description="Search the arxiv-complete corpus (semantic + keyword).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--query", "-q", required=True, help="Natural-language query.")
    p.add_argument("--top-k", type=int, default=10, help="Max results to return.")
    p.add_argument("--year-min", type=int, default=None, help="Only papers from this year onward.")
    p.add_argument("--year-max", type=int, default=None, help="Only papers up to this year.")
    p.add_argument("--category", default=None, help="Filter by primary_category, e.g. cs.LG.")
    p.add_argument("--survey-only", action="store_true", help="Only survey/overview/taxonomy papers.")
    p.add_argument(
        "--keyword-only",
        action="store_true",
        help="Skip the ANN index entirely; use the pure-keyword fallback.",
    )
    p.add_argument(
        "--embedder",
        default=None,
        choices=["hashing", "remote"],
        help="Embedding provider. Default: auto (remote if key set, else local hashing).",
    )
    p.add_argument(
        "--index-name",
        default="arxiv",
        help="Name of the persisted FAISS index to use/load.",
    )
    p.add_argument(
        "--parquet",
        default=None,
        help="Local parquet path or HF URL overriding the default metadata source.",
    )
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    store = DataStore(source=args.parquet if args.parquet else os.environ.get(
        "ARXIV_METADATA_PARQUET", DataStore().source
    ))

    opts = QueryOptions(
        top_k=args.top_k,
        year_min=args.year_min,
        year_max=args.year_max,
        primary_category=args.category,
        survey_only=args.survey_only,
    )

    if args.keyword_only:
        results = _keyword_search_only(store, args, opts)
        _print_results(results, args.query, "keyword fallback (no embeddings)")
        return 0

    # Semantic path: try to load an existing index; if absent, build it.
    from index import IndexBuilder, load_index  # imported late to keep CLI fast

    embedder = get_embedder(args.embedder)
    builder = IndexBuilder(args.index_name, embedder)
    index_path = builder.default_path()

    if os.path.exists(index_path):
        try:
            ann, meta = load_index(args.index_name, embedder)
            note = f"ANN index '{args.index_name}' ({meta.count} papers, {meta.embedder})"
        except Exception as exc:
            print(f"[query] failed to load index ({exc}); building fresh.")
            ann = None
    else:
        print(f"[build] no index at {index_path}; building now (this embeds the corpus).")
        added = builder.build(store.scan_filtered(), store)
        meta = builder.save()
        ann, meta = load_index(args.index_name, embedder)
        note = f"freshly built ANN index ({meta.count} papers, {meta.embedder})"

    results = query(args.query, store, ann=ann, embedder=embedder, opts=opts)
    _print_results(results, args.query, note)
    return 0


def _keyword_search_only(store, args, opts) -> list[Result]:
    from query import _keyword_search

    return _keyword_search(args.query, store, opts)


if __name__ == "__main__":
    raise SystemExit(main())