"""
Test script for the Jina Reranker integration.
This script tests the reranker functionality by comparing results with and without reranking.
"""

import json
import time
import os
from pathlib import Path
from typing import Dict, List, Any, Optional

# Import just what we need for the simple test
from ranking.jina_reranker import JinaReranker, get_jina_reranker

def test_simple_reranker():
    """Test the Jina Reranker with a simple query and documents"""
    # Initialize the reranker directly without parameters (it will read from config)
    try:
        reranker = get_jina_reranker()
        print("Successfully initialized Jina Reranker")
    except Exception as e:
        print(f"Error initializing Jina Reranker: {str(e)}")
        return
    
    # Simple query and documents
    query = "What is quantum computing?"
    documents = [
        "Quantum computing is a type of computation that harnesses quantum mechanics.",
        "Classical computers use bits, while quantum computers use qubits.",
        "Machine learning is a subset of artificial intelligence.",
        "Quantum computers can solve certain problems faster than classical computers."
    ]
    
    print(f"Testing reranker with query: {query}")
    print(f"Documents: {documents}")
    
    # Rerank the documents
    try:
        reranked = reranker.rerank(query, documents)
        print(f"Reranked results: {json.dumps(reranked, indent=2)}")
        
        # Save the results to a file for analysis
        results_dir = Path("results")
        results_dir.mkdir(exist_ok=True)
        results_file = results_dir / f"reranked_results_{int(time.time())}.json"
        
        with open(results_file, "w") as f:
            json.dump(reranked, f, indent=2)
        
        print(f"Results saved to {results_file}")
        return True
    except Exception as e:
        print(f"Error reranking: {str(e)}")
        return False

if __name__ == "__main__":
    # Just run the simple test
    success = test_simple_reranker()
    
    if success:
        print("Jina Reranker test completed successfully!")
    else:
        print("Jina Reranker test failed.")