#!/usr/bin/env python3
"""
Test script to compare old vs new response formats.
Run this to see the enhanced explanatory responses.
"""

from pathlib import Path
from rag_pipeline import RAGPipeline, DocumentStore

def test_student_question():
    """Test with a student-focused question"""
    print("=" * 80)
    print("Testing Enhanced RAG with Student Question")
    print("=" * 80)
    
    # Initialize
    vector_store_path = Path("vector_store")
    doc_store = DocumentStore(
        persist_dir=vector_store_path,
        embedding_model="sentence-transformers/all-MiniLM-L6-v2"
    )
    
    # Load existing vector store
    src = Path("data")
    pdfs = doc_store.discover_pdfs(src)
    doc_store.build_vector_store(pdfs, force_rebuild=False)
    
    # Initialize RAG
    rag = RAGPipeline(
        doc_store=doc_store,
        model="llama-3.3-70b-versatile",
        temperature=0.1,
        top_k=8
    )
    
    # Test question
    question = "As a student, what do I need to know about the new tax law?"
    
    print(f"\nQuestion: {question}\n")
    print("Generating answer...\n")
    
    answer = rag.query(question, verbose=True)
    
    print("\n" + "=" * 80)
    print("ANSWER:")
    print("=" * 80)
    print(answer)
    print("=" * 80)


def test_multiple_personas():
    """Test different persona questions"""
    questions = [
        ("As a student, what do I need to know about the new tax law?", "Student"),
        ("How does the Development Levy affect my small business?", "Business Owner"),
        ("What PAYE deductions can I claim as an employee?", "Employee"),
        ("What is the corporate income tax rate?", "General")
    ]
    
    print("=" * 80)
    print("Testing Multiple Personas")
    print("=" * 80)
    
    # Initialize once
    vector_store_path = Path("vector_store")
    doc_store = DocumentStore(
        persist_dir=vector_store_path,
        embedding_model="sentence-transformers/all-MiniLM-L6-v2"
    )
    
    src = Path("data")
    pdfs = doc_store.discover_pdfs(src)
    doc_store.build_vector_store(pdfs, force_rebuild=False)
    
    rag = RAGPipeline(
        doc_store=doc_store,
        model="llama-3.3-70b-versatile",
        temperature=0.1,
        top_k=6
    )
    
    for question, persona_type in questions:
        print(f"\n{'=' * 80}")
        print(f"PERSONA: {persona_type}")
        print(f"QUESTION: {question}")
        print("=" * 80)
        
        answer = rag.query(question, verbose=False)
        print(answer)
        print("\n")


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "--all":
        test_multiple_personas()
    else:
        test_student_question()
        print("\n\nTip: Run with --all flag to test multiple personas")