import sys
import json
from pathlib import Path

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent.parent / "03_embedding"))

from retriever_multi import get_unified_context_for_llm, export_retrieval_to_json
from embed_query import load_client
import config


def test_retrieval():
    """Test retrieval and output JSON for LLM integration."""
    print("=" * 60)
    print("Multi-Namespace Retrieval Test")
    print("=" * 60)
    
    # Load NVIDIA NIM embedding client
    print("\n Loading NVIDIA NIM embedding model...")
    client = load_client(config.EMBEDDING_MODEL, config.NVIDIA_API_KEY, "NONE")
    print("    Model loaded")
    
    while True:
        # Get query from user
        query = input("\n Enter query (or 'quit' to exit): ").strip()
        
        if query.lower() in ['quit', 'exit', 'q']:
            print("\n Goodbye!")
            break
        
        if not query:
            continue
        
        # Embed query
        print("\n Embedding query...")
        query_embedding = client.embed_query(query)
        
        # Retrieve from all namespaces
        print(" Retrieving from Pinecone...")
        context = get_unified_context_for_llm(
            query_embedding=query_embedding,
            query_text=query,
            doc_top_k=5,
            memory_top_k=3,
            include_related=True,
        )
        
        # Print summary
        print(f"\n Results:")
        print(f"    Document chunks: {context['doc_count']}")
        print(f"    Memory chunks: {context['memory_count']}")
        print(f"    Related chunks: {context['related_count']}")
        
        # Export to JSON
        output_path = Path(__file__).parent / "output" / "retrieval_output.json"
        output = export_retrieval_to_json(query, context, str(output_path))
        
        print(f"\n Saved to: {output_path}")
        
        # Also print the JSON
        print("\n" + "=" * 60)
        print("JSON Output:")
        print("=" * 60)
        print(json.dumps(output, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    test_retrieval()