# benchmarking/inspect_prompts.py

"""Module to inspect LLM prompts generated for search result evaluation.

This script loads a specific query and its search results, applies a
chosen permutation of search engine order, and displays the system and
user prompts that would be generated by the 'evaluate_results.py' script.
It is intended for debugging and fine-tuning prompt generation logic.
"""

import argparse
import json
import os
from itertools import permutations
from typing import Any, Dict, List, Tuple

# Assuming inspect_prompts.py is in the same directory as evaluate_results.py
try:
    from search_engine_ranker import (
        MAX_RESULTS_TO_DISPLAY_PER_ENGINE,
        SEARCH_ENGINE_KEYS,
        SEARCH_ENGINE_NAMES,
        SEARCH_RESULTS_FILE,
        construct_llm_prompt,
        format_search_results_for_prompt,
    )
except ImportError as e:
    print("Error: Could not import necessary components from 'evaluate_results.py'.")
    print(f"Details: {e}")
    print(
        "Please ensure 'evaluate_results.py' is in the same directory "
        "as 'inspect_prompts.py' and that all its global functions and "
        "constants are correctly defined."
    )
    import sys

    sys.exit(1)


def display_prompts_for_query(query_index: int, permutation_index: int) -> None:
    """Loads data, generates, and displays prompts for a specific scenario.

    Args:
        query_index: The index of the query in the search_results.json file.
        permutation_index: The index of the search engine permutation to use (0-5).
    """
    search_results_path = os.path.join(os.path.dirname(__file__), SEARCH_RESULTS_FILE)

    if not os.path.exists(search_results_path):
        print(f"Error: Search results file not found at '{search_results_path}'.")
        print(
            f"Please ensure '{SEARCH_RESULTS_FILE}' exists in the "
            "same directory as this script."
        )
        return

    try:
        with open(search_results_path, encoding="utf-8") as f:
            all_search_data = json.load(f)
        if not isinstance(all_search_data, list):
            print(f"Error: '{search_results_path}' does not contain a JSON list.")
            return
        if not all_search_data:
            print(f"No search data found in '{search_results_path}'.")
            return
    except json.JSONDecodeError:
        print(f"Error: Could not decode JSON from '{search_results_path}'.")
        return
    except Exception as e:
        print(f"Error reading '{search_results_path}': {e}")
        return

    if not 0 <= query_index < len(all_search_data):
        print(
            f"Error: query_index {query_index} is out of bounds. "
            f"File contains {len(all_search_data)} queries "
            f"(indices 0 to {len(all_search_data) - 1})."
        )
        return

    query_data: Dict[str, Any] = all_search_data[query_index]
    query_string: str = query_data.get("query", "Query string not found")

    engine_key_permutations: List[Tuple[str, ...]] = list(
        permutations(SEARCH_ENGINE_KEYS)
    )
    # permutation_index is already validated by argparse choices
    chosen_engine_permutation: Tuple[str, ...] = engine_key_permutations[
        permutation_index
    ]

    print("--- Scenario Details ---")
    print(f"Query Index: {query_index}")
    print(f'Query String: "{query_string}"')
    print(f"Permutation Index: {permutation_index}")
    print(f"Engine Order (Keys): {chosen_engine_permutation}")
    print("-" * 20)

    placeholders = ["Engine A", "Engine B", "Engine C"]
    current_placeholder_to_engine_name_map: Dict[str, str] = {}
    formatted_results_for_llm_prompt: Dict[str, str] = {}

    for i, engine_key in enumerate(chosen_engine_permutation):
        placeholder = placeholders[i]
        # Use SEARCH_ENGINE_NAMES for user-friendly name, fall back to key
        actual_engine_name = SEARCH_ENGINE_NAMES.get(engine_key, engine_key)
        current_placeholder_to_engine_name_map[placeholder] = actual_engine_name

        engine_specific_results = query_data.get(engine_key)
        formatted_engine_results_str = format_search_results_for_prompt(
            engine_specific_results, engine_key, MAX_RESULTS_TO_DISPLAY_PER_ENGINE
        )
        formatted_results_for_llm_prompt[placeholder] = formatted_engine_results_str

    system_prompt_str, user_prompt_str = construct_llm_prompt(
        query_string,
        current_placeholder_to_engine_name_map,
        formatted_results_for_llm_prompt,
    )

    print("\n--- System Prompt ---")
    print(system_prompt_str)
    print("-" * 20)

    print("\n--- User Prompt ---")
    print(user_prompt_str)
    print("-" * 20)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description=(
            "Inspect the LLM prompts generated for a specific query and "
            "engine permutation, using logic from evaluate_results.py."
        )
    )
    parser.add_argument(
        "--query_index",
        type=int,
        default=0,
        help="Index of the query in search_results.json to use (default: 0).",
    )
    parser.add_argument(
        "--perm_index",
        type=int,
        default=0,
        choices=range(len(list(permutations(SEARCH_ENGINE_KEYS)))),
        metavar=f"[0-{len(list(permutations(SEARCH_ENGINE_KEYS))) - 1}]",
        help=(
            "Index of the search engine permutation to use (default: 0). "
            f"There are {len(list(permutations(SEARCH_ENGINE_KEYS)))} permutations "
            f"for {len(SEARCH_ENGINE_KEYS)} engines."
        ),
    )
    args = parser.parse_args()

    display_prompts_for_query(args.query_index, args.perm_index)
