#!/usr/bin/env python3
"""
Test the complete RAG pipeline
"""

import sys
import os

# Add the current directory to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from main import RAGSystem

def test_rag_pipeline():
    """Test the complete RAG pipeline"""
    print("=== Testing Complete RAG Pipeline ===")

    # Initialize system
    rag_system = RAGSystem()

    # Setup system (process documents)
    print("Setting up system and processing documents...")
    doc_count = rag_system.setup_system()
    print(f"System ready with {doc_count} documents")

    if doc_count == 0:
        print("No documents found. Please add files to the pdf_directory.")
        return False

    # Test questions
    test_questions = [
        "What is artificial intelligence?",
        "What are the types of AI?",
        "What are some applications of AI in healthcare?",
        "What are the ethical challenges of AI?"
    ]

    print("\n=== Testing Q&A ===")

    for question in test_questions:
        print(f"\nQuestion: {question}")
        try:
            result = rag_system.search_and_answer(question)
            if isinstance(result, str):
                print(f"Answer: {result}")
            else:
                print(f"Answer: {result['answer']}")
                if 'sources' in result and result['sources']:
                    print(f"Sources: {len(result['sources'])} relevant chunks found")
        except Exception as e:
            print(f"Error: {e}")

    print("\n=== RAG Pipeline Test Complete ===")
    return True

if __name__ == "__main__":
    success = test_rag_pipeline()
    sys.exit(0 if success else 1)