#!/usr/bin/env python3
"""
ArXiv Query Optimization Benchmark
====================================
This script measures whether LLM-based query reformulation improves
search result relevance on arXiv, which uses lexical (keyword) search
rather than semantic search.

HOW IT WORKS:
  1. Takes a set of natural language research questions
  2. Sends each question to arXiv AS-IS and records the results
  3. Reformulates each question into keywords using an LLM
  4. Sends the optimized query to arXiv and records the results
  5. Prints a side-by-side comparison for human evaluation

USAGE:
  # With Ollama running locally (default):
  python benchmark_arxiv_query_optimization.py

  # With a specific Ollama model:
  python benchmark_arxiv_query_optimization.py --model llama3:8b

  # With OpenAI:
  python benchmark_arxiv_query_optimization.py --provider openai --model gpt-4o-mini

  # With Gemini:
  # python benchmark_arxiv_query_optimization.py --provider gemini --model gemini-1.5-flash

  # Dry run (skip LLM optimization, only test raw arXiv queries):
  python benchmark_arxiv_query_optimization.py --dry-run

  # Save results to YAML for future reference:
  python benchmark_arxiv_query_optimization.py --output results.yaml

NOTE: This is a prototype benchmark script. Results require human evaluation
      to determine whether the optimized queries return more relevant papers.
"""

import argparse
import pytest
import yaml
import os
import time
from typing import Any, Optional
from pathlib import Path
import sys
from local_deep_research.web_search_engines.engines.search_engine_arxiv import (
    ArXivSearchEngine,
)

repo_root = Path(__file__).resolve().parents[4]
src_path = repo_root / "src"
if str(src_path) not in sys.path:
    sys.path.append(str(src_path))

# ------------- MOCK SETTINGS ---------
# mock settings to pass to the LDR engine without needing the full server context
MOCK_SETTINGS = {"llm.ollama.url": "http://localhost:11434"}

# ------------- TEST QUERIES -------------

TEST_QUERIES = [
    {
        "question": "Which machine learning benchmarks are the most well-known for weather forecasting?",
        "topic": "ML + climate",
        "expected_keywords": [
            "machine learning",
            "benchmarks",
            "weather forecasting",
        ],
    },
    {
        "question": "How can explainable AI improve diagnostic transparency in healthcare?",
        "topic": "Explainable AI (XAI)",
        "expected_keywords": ["explainable AI", "healthcare", "transparency"],
    },
    {
        "question": "What are the recent breakthroughs in using deep learning for protein folding?",
        "topic": "AI in Biology",
        "expected_keywords": [
            "deep learning",
            "protein folding",
            "structure prediction",
        ],
    },
    {
        "question": "How are contrastive loss functions used to align text and images?",
        "topic": "Methodology / Multimodal AI",
        "expected_keywords": ["contrastive loss", "alignment", "text", "image"],
    },
    {
        "question": "What are the latest advancements in continual learning for spatio-temporal graphs?",
        "topic": "Continual Learning on Graphs",
        "expected_keywords": [
            "continual learning",
            "spatio-temporal",
            "graphs",
        ],
    },
    {
        "question": "Search for papers aboud 3D cellular automata applications as I need one to present one possible application in my seminar.",
        "topic": "Complex Systems / Fire Modeling",
        "expected_keywords": [
            "3D modeling",
            "mathematical systems",
            "fire modeling",
        ],
    },
    {
        "question": "What is the difference between a qubit and a classical computer bit?",
        "topic": "Quantum Computing",
        "expected_keywords": ["qubit", "classical bit", "quantum computing"],
    },
    {
        "question": "Summarize the content of the paper 'Dynamic Graph Echo State Networks' by Tortorella and Micheli.",
        "topic": "Reservoir Computing / Deep Learning for Graphs",
        "expected_keywords": [
            "Dynamic Graph Echo State Networks",
            "Tortorella",
            "Micheli",
        ],
    },
    {
        "question": "What are the ethical implications and privacy concerns of using facial recognition systems in public surveillance?",
        "topic": "Ethics in AI / Computer Vision",
        "expected_keywords": [
            "ethical",
            "privacy",
            "facial recognition",
            "surveillance",
        ],
    },
    {
        "question": "Which are the best lightweight CNN architectures for mobile devices?",
        "topic": "Edge AI / Model Compression",
        "expected_keywords": [
            "lightweight",
            "CNN",
            "mobile",
            "Convolutional Neural Networks",
        ],
    },
]


# ------------- LLM QUERY OPTIMIZER -------------
# This mirrors the pattern used in PubMed's _optimize_query_for_pubmed() adapted for arXiv's search syntax.


ARXIV_OPTIMIZATION_PROMPT = """Convert the natural language question into an arXiv search query. Return ONLY the query, nothing else.

QUESTION: "{query}"

CRITICAL RULES:
1. Return ONLY the query - no explanation, no prefix like "Query:" or "Here is:"
2. Extract 2-3 core technical keywords. Drop filler words like "latest", "recent", "most interesting", "summarize"
3. Boolean operators uppercase only: AND, OR, ANDNOT
4. Use quotes only for fixed 2-3 word technical terms: "deep learning", "string theory", "convolutional neural networks"
5. Never quote long phrases: BAD "ocean plastic pollution", GOOD: ocean AND plastic AND pollution
6. Use (ti:"X" OR abs:"X") only for the 2/3 most important concepts - not for every concept
7. Use au: only for real author names in the question, NEVER for topics. If applicable, use au:"author lastname in the question"

EXAMPLES:
Original: "Show me reinforcement learning applications in robotics or robot control."
✓ GOOD: (ti:"reinforcement learning" OR abs:"reinforcement learning") AND (ti:robotics OR abs:robot) AND (ti:control OR abs:control)

Original: "Find papers by Geoffrey Hinton focusing on computer vision."
✓ GOOD: au:"Geoffrey Hinton" AND (ti:"computer vision" OR abs:"computer vision")

Original: "How does ocean plastic pollution affect marine life?"
✓ GOOD: (ti:ocean OR abs:ocean OR ti:marine OR abs:marine) AND (ti:plastic OR abs:plastic) AND (ti:pollution OR abs:pollution)

Original: "What are the impacts of blockchain technology in finance?"
✗ BAD: Here is the query: "blockchain technology" AND finance
(REASON: Fails Rule 1 by adding conversational text. Fails Rule 3 by missing ti: and abs: fields.)

Original: "I need papers about cybersecurity and deep learning."
✗ BAD: papers about "cybersecurity" and "deep learning"
(REASON: Fails Rule 1. No Boolean logic used.)

Original: "What are the best machine learning solutions for urban traffic congestion?"
✗ BAD: (ti:"machine learning solutions" OR abs:"machine learning solutions") AND (ti:"urban traffic congestion" OR abs:"urban traffic congestion")
(REASON: Fails Rule 2. "machine learning solutions" and "urban traffic congestion" are descriptive phrases. Only the strict 2-word concept "machine learning" can be in quotes. The rest MUST be atomic: (ti:"machine learning" OR abs:"machine learning") AND (ti:traffic OR abs:traffic) AND (ti:congestion OR abs:congestion)

Return ONLY the query string.
"""


def get_llm(provider: str, model: str):
    """
    Initialize an LLM for query optimization. Supports Ollama, OpenAI and Gemini for tests.

    Args:
        provider: "ollama", "openai" or "gemini"
        model: Model name (e.g., "llama3:8b" or "gpt-4o-mini" or "gemini-2.5-flash")

    Returns:
        A LangChain LLM instance, or None if initialization fails
    """
    try:
        if provider == "ollama":
            from langchain_ollama import ChatOllama

            llm = ChatOllama(model=model, temperature=0)
            llm.invoke("test")
            print(f"  Using Ollama with model '{model}'...")
            return llm

        if provider == "openai":
            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                print("  ERROR: OPENAI_API_KEY environment variable not set.")
                return None
            from langchain_openai import ChatOpenAI

            print(f"  Using OpenAI with model '{model}'...")
            return ChatOpenAI(model=model, temperature=0, api_key=api_key)

        if provider == "gemini":
            api_key = os.environ.get("GOOGLE_API_KEY")
            if not api_key:
                print("  ERROR: GOOGLE_API_KEY environment variable not set.")
                return None
            from langchain_google_genai import ChatGoogleGenerativeAI

            print(f"  Using Gemini with model '{model}'...")
            return ChatGoogleGenerativeAI(
                model=model, temperature=0, google_api_key=api_key
            )

        print(
            f"  ERROR: Unknown provider '{provider}'. Use 'ollama', 'openai' or 'gemini'."
        )
        return None

    except Exception as e:
        print(f"  WARNING: Could not initialize LLM: {e}")
        return None


def optimize_query_for_arxiv(query: str, llm) -> str:
    """
    Use a LLM to convert a natural language question into an optimized arXiv search query.

    Args:
        query: Natural language research question
        llm: LangChain LLM instance

    Returns:
        Optimized query string (or original if LLM fails)
    """
    try:
        prompt = ARXIV_OPTIMIZATION_PROMPT.format(query=query)
        response = llm.invoke(prompt)

        # Extract the query from the LLM response
        raw = (
            response.content if hasattr(response, "content") else str(response)
        )
        raw = raw.strip()
        lines = [line.strip() for line in raw.split("\n") if line.strip()]
        optimized = lines[0] if lines else query

        # Sanity check: if the output looks like an explanation, fall back
        if (
            optimized.lower().startswith(("i ", "here", "the query", "sure"))
            or len(optimized.split()) > 40
        ):
            print(
                f'    Warning: LLM output looks like an explanation, using original \n     optimized: "{optimized}"'
            )
            return query

        return optimized

    except Exception as e:
        print(f"    Warning: Query optimization failed: {e}")
        return query


def search_arxiv(query: str, max_results: int = 5) -> list[dict[str, Any]]:
    """
    Run a search against arXiv and return the top results.

    Uses the arxiv library directly (already a dependency of LDR) to avoid
    needing the full LDR settings/database context that ArXivSearchEngine
    requires when running inside the server.

    Args:
        query: The search query (raw or optimized)
        max_results: Number of results to fetch

    Returns:
        List of dicts with title, authors, categories, and abstract snippet
    """
    try:
        import arxiv

        client = arxiv.Client(page_size=max_results)
        search = arxiv.Search(
            query=query,
            max_results=max_results,
            sort_by=arxiv.SortCriterion.Relevance,
            sort_order=arxiv.SortOrder.Descending,
        )

        papers = list(client.results(search))

        results = []
        for paper in papers[:max_results]:
            results.append(
                {
                    "title": paper.title,
                    "authors": [a.name for a in paper.authors[:2]],
                    "categories": paper.categories[:3]
                    if paper.categories
                    else [],
                    "abstract_snippet": paper.summary[:200] + "..."
                    if paper.summary
                    else "",
                    "published": paper.published.strftime("%Y-%m")
                    if paper.published
                    else "unknown",
                    "arxiv_id": paper.entry_id.split("/")[-1]
                    if paper.entry_id
                    else "",
                }
            )

        return results

    except Exception as e:
        print(f"    ERROR searching arXiv: {e}")
        return []


def search_arxiv_LDR(
    query: str, max_results: int = 5, engine: ArXivSearchEngine = None
) -> list[dict[str, Any]]:
    """
    Run a search using the real LDR ArXivSearchEngine.
    """
    try:
        ldr_results = engine.run(query=query)

        results = []
        for paper in ldr_results[:max_results]:
            results.append(
                {
                    "title": paper.get("title", "Unknown Title"),
                    "authors": paper.get("authors", []),
                    "categories": paper.get("categories", []),
                    "abstract_snippet": paper.get("snippet", ""),
                    "published": paper.get("published", "unknown")[
                        :7
                    ],  # Keep only YYYY-MM
                    "arxiv_id": paper.get("id", "").split("/")[-1]
                    if paper.get("id")
                    else "",
                }
            )

        return results

    except Exception as e:
        print(f"    ERROR searching arXiv via LDR engine: {e}")
        return []


def print_comparison(
    query_info: dict,
    raw_query: str,
    optimized_query: str,
    raw_results: list,
    optimized_results: list,
) -> dict:
    """
    Print a readable side-by-side comparison of raw vs optimized results.

    Returns a dict with the comparison data for YAML export.
    """
    sep = "─" * 70
    print(f"\n{sep}")
    print(f"TOPIC: {query_info['topic']}")
    print(f"QUESTION: {query_info['question']}")
    print(f"{sep}")

    print(f'\n  RAW QUERY:       "{raw_query}"')
    print(f'  OPTIMIZED QUERY: "{optimized_query}"')

    # Print raw results
    print(f"\n  ── Raw Results ({len(raw_results)} papers) ──")
    if raw_results:
        for i, r in enumerate(raw_results, 1):
            cats = ", ".join(r["categories"])
            print(f"  {i}. [{r['published']}] {r['title']}")
            print(f"     Categories: {cats}")
    else:
        print("  (no results)")

    # Print optimized results
    print(f"\n  ── Optimized Results ({len(optimized_results)} papers) ──")
    if optimized_results:
        for i, r in enumerate(optimized_results, 1):
            cats = ", ".join(r["categories"])
            print(f"  {i}. [{r['published']}] {r['title']}")
            print(f"     Categories: {cats}")
    else:
        print("  (no results)")

    # Count overlap — same paper appearing in both result sets
    raw_titles = {r["title"] for r in raw_results}
    opt_titles = {r["title"] for r in optimized_results}
    overlap = raw_titles & opt_titles
    new_in_optimized = opt_titles - raw_titles

    print("\n  ── Summary ──")
    print(f"  Papers in common:         {len(overlap)}")
    print(f"  New papers in optimized:  {len(new_in_optimized)}")
    if new_in_optimized:
        print("  New titles:")
        for t in list(new_in_optimized)[:3]:
            print(f"    + {t[:80]}")

    return {
        "topic": query_info["topic"],
        "question": query_info["question"],
        "raw_query": raw_query,
        "optimized_query": optimized_query,
        "raw_result_count": len(raw_results),
        "optimized_result_count": len(optimized_results),
        "overlap_count": len(overlap),
        "new_in_optimized": len(new_in_optimized),
        "raw_results": raw_results,
        "optimized_results": optimized_results,
    }


@pytest.mark.slow
def run_benchmark(
    provider: str = "ollama",
    model: str = "llama3.2",
    max_results: int = 5,
    dry_run: bool = False,
    output_path: Optional[str] = None,
    delay: float = 3.0,
) -> list[dict]:
    """
    Run the full benchmark across all test queries.

    Args:
        provider: LLM provider ("ollama", "openai", or "gemini")
        model: LLM model name
        max_results: Number of arXiv results per query
        dry_run: If True, skip LLM optimization (only test raw queries)
        output_path: Optional path to save YAML results
        delay: Seconds to wait between arXiv requests (avoid rate limiting)

    Returns:
        List of comparison result dicts
    """
    print("=" * 70)
    print("ArXiv Query Optimization Benchmark")
    print("=" * 70)
    print(f"Provider: {provider} | Model: {model}")
    print(f"Results per query: {max_results} | Dry run: {dry_run}")
    print(f"Queries to test: {len(TEST_QUERIES)}")

    # Initialize LLM (unless dry run)
    llm = None
    if not dry_run:
        print("\nInitializing LLM...")
        llm = get_llm(provider, model)
        if llm is None:
            print("  Falling back to dry-run mode (no LLM available).")
            dry_run = True
        engine = ArXivSearchEngine(
            max_results=max_results, llm=llm, settings_snapshot=MOCK_SETTINGS
        )

    all_results = []

    for i, query_info in enumerate(TEST_QUERIES):
        if provider in ["openai", "gemini"] and i % 2 == 0 and i > 0:
            print(
                f"\nWaiting {30} seconds before next query to avoid rate limits..."
            )
            time.sleep(30)

        question = query_info["question"]
        print(
            f"\n[{i + 1}/{len(TEST_QUERIES)}] Processing: {query_info['topic']}..."
        )

        # Step 1: Optimize the query (or use raw if dry run)
        if dry_run or llm is None:
            optimized_query = question
            print("  Dry run — skipping optimization")
        else:
            print("  Optimizing query with LLM...")
            optimized_query = optimize_query_for_arxiv(question, llm)
            print(f'  Optimized: "{optimized_query}"')

        # Step 2: Search arXiv with the raw question
        print("  Searching arXiv with raw query...")
        if not dry_run:
            raw_results = search_arxiv_LDR(question, max_results, engine=engine)
        else:
            raw_results = search_arxiv(question, max_results)
        time.sleep(delay)

        # Step 3: Search arXiv with the optimized query (skip if same)
        if optimized_query != question:
            print("  Searching arXiv with optimized query...")
            optimized_results = search_arxiv_LDR(
                optimized_query, max_results, engine=engine
            )
            time.sleep(delay)
        else:
            optimized_results = raw_results
            print("  Skipping optimized search (query unchanged)")

        # Step 4: Print and store comparison
        comparison = print_comparison(
            query_info,
            question,
            optimized_query,
            raw_results,
            optimized_results,
        )
        all_results.append(comparison)

    # Final summary
    print(f"\n{'=' * 70}")
    print("BENCHMARK COMPLETE — SUMMARY")
    print(f"{'=' * 70}")
    total_new = sum(r["new_in_optimized"] for r in all_results)
    total_overlap = sum(r["overlap_count"] for r in all_results)
    print(f"Total queries tested:           {len(all_results)}")
    print(f"Total new papers via optimized: {total_new}")
    print(f"Total papers in common:         {total_overlap}")
    print("\nNOTE: Human evaluation is required to assess result relevance.")
    print("      Review the paper titles above and judge if the optimized")
    print("      queries returned more on-topic results.")

    # Save to yaml if requested
    if output_path:
        output = {
            "benchmark_config": {
                "provider": provider,
                "model": model,
                "max_results": max_results,
                "dry_run": dry_run,
            },
            "results": all_results,
        }
        with open(output_path, "w") as f:
            yaml.dump(output, f, sort_keys=False, default_flow_style=False)
        print(f"\nResults saved to: {output_path}")

    return all_results


def parse_args():
    parser = argparse.ArgumentParser(
        description="Benchmark arXiv search quality with and without LLM query optimization",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--provider",
        default=os.environ.get("LDR_PROVIDER", "ollama"),
        choices=["ollama", "openai", "gemini"],
        help="LLM provider for query optimization (default: ollama)",
    )
    parser.add_argument(
        "--model",
        default=os.environ.get("LDR_MODEL", "llama3.2"),
        help="LLM model name (default: llama3.2 for Ollama, gpt-4o-mini for OpenAI, gemini-1.5-pro for Gemini)",
    )
    parser.add_argument(
        "--max-results",
        type=int,
        default=5,
        help="Number of arXiv results per query (default: 5)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Skip LLM optimization, only test raw arXiv queries",
    )
    parser.add_argument(
        "--output",
        metavar="FILE",
        help="Save results to a YAML file",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=5.0,
        help="Seconds to wait between arXiv requests to avoid rate limiting (default: 5)",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    run_benchmark(
        provider=args.provider,
        model=args.model,
        max_results=args.max_results,
        dry_run=args.dry_run,
        output_path=args.output,
        delay=args.delay,
    )
