#!/usr/bin/env python3
"""
Test script for Mini RAG system
Run this to verify all components work before deployment
"""

import os
import sys
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def test_imports():
    """Test that all required modules can be imported"""
    print("Testing imports...")
    try:
        from chunker import chunk_text
        from llm import LLMProvider
        from pinecone_client import PineconeClient
        from rag_core import RAGCore
        from ingest import load_documents
        print("✅ All imports successful")
        return True
    except ImportError as e:
        print(f"❌ Import failed: {e}")
        return False

def test_chunking():
    """Test text chunking functionality"""
    print("\nTesting chunking...")
    try:
        from chunker import chunk_text
        
        test_text = "This is a test document. " * 50  # Create long text
        chunks = chunk_text(test_text, chunk_size=100, chunk_overlap=20)
        
        if len(chunks) > 1:
            print(f"✅ Chunking works: {len(chunks)} chunks created")
            return True
        else:
            print("❌ Chunking failed: expected multiple chunks")
            return False
    except Exception as e:
        print(f"❌ Chunking test failed: {e}")
        return False

def test_environment():
    """Test environment variable configuration"""
    print("\nTesting environment variables...")
    
    required_vars = ['PINECONE_API_KEY', 'OPENAI_API_KEY']
    optional_vars = ['GROQ_API_KEY', 'COHERE_API_KEY']
    
    missing_required = []
    for var in required_vars:
        if not os.getenv(var):
            missing_required.append(var)
    
    if missing_required:
        print(f"❌ Missing required environment variables: {missing_required}")
        print("Please set these in your .env file")
        return False
    
    print("✅ Required environment variables set")
    
    # Check optional variables
    for var in optional_vars:
        if os.getenv(var):
            print(f"✅ {var} is set")
        else:
            print(f"⚠️  {var} not set (optional)")
    
    return True

def test_document_loading():
    """Test document loading functionality"""
    print("\nTesting document loading...")
    try:
        from ingest import load_documents
        
        # Check if data directory exists
        data_dir = "./data"
        if not os.path.exists(data_dir):
            print(f"⚠️  Data directory {data_dir} not found")
            return False
        
        docs = load_documents(data_dir)
        if docs:
            print(f"✅ Document loading works: {len(docs)} documents found")
            for doc in docs:
                print(f"   - {doc['path']} ({len(doc['text'])} characters)")
            return True
        else:
            print("⚠️  No documents found in data directory")
            return False
            
    except Exception as e:
        print(f"❌ Document loading test failed: {e}")
        return False

def test_llm_provider():
    """Test LLM provider initialization"""
    print("\nTesting LLM provider...")
    try:
        from llm import LLMProvider
        
        llm = LLMProvider()
        print(f"✅ LLM provider initialized: {llm.provider}")
        print(f"   - Embedding model: {llm.embedding_model}")
        print(f"   - LLM model: {llm.llm_model}")
        print(f"   - Reranker: {llm.rerank_provider}")
        
        return True
    except Exception as e:
        print(f"❌ LLM provider test failed: {e}")
        return False

def test_pinecone_client():
    """Test Pinecone client initialization"""
    print("\nTesting Pinecone client...")
    try:
        from pinecone_client import PineconeClient
        
        pc = PineconeClient()
        print(f"✅ Pinecone client initialized")
        print(f"   - Index: {pc.index_name}")
        print(f"   - Cloud: {pc.cloud}")
        print(f"   - Region: {pc.region}")
        
        return True
    except Exception as e:
        print(f"❌ Pinecone client test failed: {e}")
        return False

def test_rag_core():
    """Test RAG core initialization"""
    print("\nTesting RAG core...")
    try:
        from rag_core import RAGCore
        
        rag = RAGCore()
        print("✅ RAG core initialized")
        
        return True
    except Exception as e:
        print(f"❌ RAG core test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🧪 Mini RAG System Test Suite")
    print("=" * 40)
    
    tests = [
        test_imports,
        test_environment,
        test_chunking,
        test_document_loading,
        test_llm_provider,
        test_pinecone_client,
        test_rag_core,
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        if test():
            passed += 1
    
    print("\n" + "=" * 40)
    print(f"Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! System is ready for deployment.")
        return True
    else:
        print("⚠️  Some tests failed. Please fix issues before deployment.")
        return False

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