"""Test script for domain-aware retrieval."""
import sys
import asyncio
sys.path.insert(0, '/home/neil/Documents/textiq-doc-extraction')

from src.knowledge.domain_aware_retriever import DomainAwareRetriever


async def test_retrieval():
    """Test domain-aware retrieval with sample queries."""
    retriever = DomainAwareRetriever()

    print("="*80)
    print("DOMAIN-AWARE RETRIEVAL TESTS")
    print("="*80)

    # Test 1: Finance query
    print("\n[TEST 1] Finance Query: 'What was our Q4 EBITDA and revenue?'")
    print("-"*80)
    try:
        results = await retriever.retrieve(
            query="What was our Q4 EBITDA and revenue?",
            top_k=3
        )

        print(f"✓ Found {len(results)} results")
        for i, result in enumerate(results[:3]):
            print(f"\n  Result {i+1}:")
            print(f"    Category: {result.get('domain_category', 'N/A')}")
            print(f"    Terms: {result.get('domain_terms', [])[:5]}...")
            print(f"    Score: {result.get('score', 0):.4f}")
            print(f"    Text: {result.get('text_content', '')[:80]}...")
    except Exception as e:
        print(f"✗ Error: {e}")

    # Test 2: HR query
    print("\n\n[TEST 2] HR Query: 'How many employees were hired last quarter?'")
    print("-"*80)
    try:
        results = await retriever.retrieve(
            query="How many employees were hired last quarter?",
            top_k=3
        )

        print(f"✓ Found {len(results)} results")
        for i, result in enumerate(results[:3]):
            print(f"\n  Result {i+1}:")
            print(f"    Category: {result.get('domain_category', 'N/A')}")
            print(f"    Score: {result.get('score', 0):.4f}")
    except Exception as e:
        print(f"✗ Error: {e}")

    # Test 3: Generic query (should still work)
    print("\n\n[TEST 3] Generic Query: 'Tell me about the data'")
    print("-"*80)
    try:
        results = await retriever.retrieve(
            query="Tell me about the data",
            top_k=3
        )

        print(f"✓ Found {len(results)} results")
        print(f"  (Generic query - should fall back to full search)")
    except Exception as e:
        print(f"✗ Error: {e}")

    print("\n" + "="*80)
    print("RETRIEVAL TESTS COMPLETE")
    print("="*80)


if __name__ == "__main__":
    asyncio.run(test_retrieval())