"""Test search functionality"""

from collections.abc import Callable
import os

# ------------------------------------------------------------------------------
# Search functionality Testing Utils

from oa.vector_stores import Query, MaxNumResults, ResultT, SearchResults


def top_results_contain(results: SearchResults, expected: SearchResults) -> bool:
    """
    Check that the top results contain the expected elements.
    That is, the first len(expected) elements of results match the expected set,
    and if there are less results than expected, the only elements in results are
    contained in expected.
    """
    if len(results) < len(expected):
        return set(results) <= set(expected)
    return set(results[: len(expected)]) == set(expected)


def general_test_for_search_function(
    query,
    top_results_expected_to_contain: SearchResults,
    *,
    search_func: Callable[[Query], SearchResults],
    n_top_results=None,
):
    """
    General test function for search functionality.

    Args:
        query: Query string
        top_results_expected_to_contain: Set of expected document keys
        search_func: Search function to test (keyword-only)
        n_top_results: Number of top results to check. If None, defaults to min(len(results), len(top_results_expected_to_contain)) (keyword-only)

    Example use:

    >>> def search_docs_containing(query):
    ...     docs = {'doc1': 'apple pie recipe', 'doc2': 'car maintenance guide', 'doc3': 'apple varieties'}
    ...     return (key for key, text in docs.items() if query in text)
    >>> general_test_for_search_function(
    ...     query='apple',
    ...     top_results_expected_to_contain={'doc1', 'doc3'},
    ...     search_func=search_docs_containing
    ... )
    """
    # Execute search and collect results
    # TODO: Protect from cases where search_func(query) could be a long generator? Example, a max_results limit?
    results = list(search_func(query))

    # Determine the actual number of top results to check
    if n_top_results is None:
        effective_n_top_results = min(
            len(results), len(top_results_expected_to_contain)
        )
    else:
        effective_n_top_results = n_top_results

    # Get the slice of results to check
    top_results_to_check = results[:effective_n_top_results]

    # Generate helpful error message
    error_context = []
    error_context.append(f"Query: '{query}'")
    error_context.append(f"Expected docs: {top_results_expected_to_contain}")
    error_context.append(f"Actual results: {results}")
    error_context.append(
        f"Checking top {effective_n_top_results} results: {top_results_to_check}"
    )

    error_message = "\n".join(error_context)

    # Perform the assertion
    assert top_results_contain(
        top_results_to_check, top_results_expected_to_contain
    ), error_message


# ─── Test Documents ────────────────────────────────────────────────────────────
docs = {
    "python": "Python is a high‑level programming language emphasizing readability and rapid development.",
    "java": "Java is a class‑based, object‑oriented language designed for portability across platforms.",
    "numpy": "NumPy provides support for large, multi‑dimensional arrays and matrices, along with a collection of mathematical functions.",
    "pandas": "Pandas is a Python library offering data structures and operations for manipulating numerical tables and time series.",
    "apple": "Apple is a fruit that grows on trees and comes in varieties such as Granny Smith, Fuji, and Gala.",
    "banana": "Banana is a tropical fruit with a soft, sweet interior and a peel that changes from green to yellow when ripe.",
    "microsoft": "Microsoft develops software products including the Windows operating system, Office suite, and cloud services.",
}

# ─── Semantic Search Examples ─────────────────────────────────────────────────


def check_search_func(
    search_func: Callable[[Query], SearchResults],
):
    """
    Test the search function with multiple queries using the general test framework.
    """
    # Test case 1: programming language search
    general_test_for_search_function(
        query="object‑oriented programming",
        top_results_expected_to_contain={"java", "python", "numpy"},
        search_func=search_func,
    )

    # Test case 2: fruit category search
    general_test_for_search_function(
        query="tropical fruit",
        top_results_expected_to_contain={"banana", "apple"},
        search_func=search_func,
    )


# ─── Retrieval‑Augmented Generation Example ────────────────────────────────────


def check_find_docs_to_answer_question(
    find_docs_to_answer_question: Callable[[Query], SearchResults],
):
    """
    Test the function that finds documents relevant to a question.
    """
    general_test_for_search_function(
        query="Which documents describe a fruit that is sweet and easy to eat?",
        top_results_expected_to_contain={"apple", "banana"},
        search_func=find_docs_to_answer_question,
    )


# ─── test these test functions with a docs_to_search_func factory function ──────


def check_search_func_factory(
    search_func_factory: Callable[[dict], Callable[[Query], SearchResults]],
):
    """
    Test the search function factory with a set of documents.
    """
    search_func = search_func_factory(docs)

    # Run the search function tests
    check_search_func(search_func)
    check_find_docs_to_answer_question(search_func)


# ------------------------------------------------------------------------------
# Tests

from oa.vector_stores import (
    OaStores,
    OaVectorStoreFiles,
    docs_to_vector_store,
    mk_search_func_for_oa_vector_store,
    docs_to_search_func_factory_via_vector_store,
)


def test_vector_store_search_functionality():
    """Test the vector store search functionality."""

    # Skip test if no OpenAI API key
    if not os.getenv("OPENAI_API_KEY") and not os.getenv("OPEN_AI_KEY"):
        print("Skipping vector store tests - no OpenAI API key found")
        return

    try:
        print("Testing docs_to_vector_store and mk_search_func_for_oa_vector_store...")

        # Test the individual functions first
        print("1. Testing docs_to_vector_store...")
        vs_id, file_mapping = docs_to_vector_store(docs, "test_search_vs")
        print(f"   Created vector store: {vs_id}")
        print(f"   File mapping: {len(file_mapping)} files")

        print("2. Testing mk_search_func_for_oa_vector_store...")
        search_func = mk_search_func_for_oa_vector_store(vs_id, file_mapping)
        print("   Search function created successfully")

        # Test a simple search
        print("3. Testing search functionality...")
        results = search_func("programming")
        print(f"   Search results: {results}")

        # Test with the factory function using our existing test framework
        print("4. Testing with check_search_func_factory...")
        print("   Note: This will create an actual vector store and use OpenAI API")

        # Uncomment the line below to run the full test (uses API calls)
        # check_search_func_factory(docs_to_search_func_factory_via_vector_store)

        print("✓ Vector store search functions created successfully")
        print("✓ Basic functionality verified")
        print("Note: Full search testing requires API calls and is commented out")

    except Exception as e:
        print(f"✗ Vector store test failed: {e}")
        print("This might be due to API key issues or OpenAI service availability")
        import traceback

        traceback.print_exc()


# Function to run the full test with API calls (uncomment to use)
def test_vector_store_search_with_api():
    """Run the full vector store search test with actual API calls."""
    print("Running full vector store search test with API calls...")
    check_search_func_factory(docs_to_search_func_factory_via_vector_store)
    print("✓ Full vector store search test completed successfully")


# ------------------------------------------------------------------------------
# Test runner

if __name__ == "__main__":
    print("=== Running Vector Store Search Tests ===")
    test_vector_store_search_functionality()

    print("\n=== Additional Test Functions Available ===")
    print("To run full API tests, call:")
    print("  test_vector_store_search_with_api()")
    print("  check_search_func_factory(docs_to_search_func_factory_via_vector_store)")
