#!/usr/bin/env python3
"""
Test script for the RAG system
"""

import sys
import os

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

def test_imports():
    """Test that all imports work"""
    try:
        from config import BASE_DIR, LANGUAGE_MODEL
        from pdf_processor import PDFProcessor
        from vector_store import VectorStore
        from mistral_integration import MistralRAG
        print("✓ All imports successful")
        return True
    except ImportError as e:
        print(f"✗ Import error: {e}")
        return False

def test_model_loading():
    """Test model loading"""
    try:
        from mistral_integration import MistralRAG
        print("Testing model loading...")
        rag = MistralRAG()
        print("✓ Model loaded successfully")
        return True
    except Exception as e:
        print(f"✗ Model loading failed: {e}")
        return False

def test_vector_store():
    """Test vector store functionality"""
    try:
        from vector_store import VectorStore
        print("Testing vector store...")
        vs = VectorStore()

        # Test with dummy data
        dummy_chunks = [
            {"content": "This is a test document about artificial intelligence.", "metadata": {"source": "test.pdf", "chunk_id": 0}},
            {"content": "Machine learning is a subset of AI.", "metadata": {"source": "test.pdf", "chunk_id": 1}},
        ]

        vs.add_documents(dummy_chunks)
        results = vs.search("artificial intelligence", k=2)

        if len(results) > 0:
            print("✓ Vector store working")
            return True
        else:
            print("✗ Vector store not returning results")
            return False
    except Exception as e:
        print(f"✗ Vector store error: {e}")
        return False

def test_pdf_processor():
    """Test PDF processor"""
    try:
        from pdf_processor import PDFProcessor
        print("Testing PDF processor...")
        processor = PDFProcessor()

        # Just test instantiation and basic methods
        new_files = processor.get_new_pdfs()
        print(f"✓ PDF processor working (found {len(new_files)} new PDFs)")
        return True
    except Exception as e:
        print(f"✗ PDF processor error: {e}")
        return False

def main():
    print("=== RAG System Test ===")

    tests = [
        ("Imports", test_imports),
        ("PDF Processor", test_pdf_processor),
        ("Vector Store", test_vector_store),
        ("Model Loading", test_model_loading),
    ]

    passed = 0
    total = len(tests)

    for test_name, test_func in tests:
        print(f"\nRunning {test_name} test...")
        if test_func():
            passed += 1
        else:
            print(f"{test_name} test failed")

    print(f"\n=== Test Results: {passed}/{total} tests passed ===")

    if passed == total:
        print("🎉 All tests passed! System is ready.")
        return 0
    else:
        print("❌ Some tests failed. Check the errors above.")
        return 1

if __name__ == "__main__":
    sys.exit(main())