import os
import sys
import torch
import time

# Ensure the app directory is in the path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(project_root)

from app import (
    index_files, 
    answer_with_rag, 
    get_embedding_function, 
    get_llm, 
    clear_index
)

def run_wiki_rag_test():
    print("--- Wikipedia Top 10 Multi-Dokument RAG Test ---")
    
    # 1. Testdatenverzeichnis
    wiki_dir = os.path.join(project_root, "test", "tests", "test_data", "wiki_top10")
    if not os.path.exists(wiki_dir):
        print(f"Fehler: Verzeichnis {wiki_dir} nicht gefunden.")
        return
        
    all_files = [os.path.join(wiki_dir, f) for f in os.listdir(wiki_dir) if f.endswith(".txt")]
    print(f"Gefundene Artikel: {len(all_files)}")

    # 2. Modelle laden
    print("Lade Modelle...")
    get_embedding_function()
    get_llm()
    
    # 3. Index vorbereiten
    print("Bereite Index vor...")
    clear_index()
    
    # 4. Alle Artikel indexieren
    print(f"Indexiere {len(all_files)} Artikel...")
    status = index_files(all_files, "mongodb://localhost:27017/", "wiki_rag_db", "wiki_chunks", False)
    print(f"Status: {status}")

    # 5. RAG Abfragen (Cross-Article / Multi-Dokument)
    test_queries = [
        {
            "query": "Erkläre den Zusammenhang zwischen Physik und Chemie. Wie wird Materie und Energie in beiden Disziplinen behandelt?",
            "expected_keywords": ["Physik", "Chemie", "Materie", "Energie", "Naturwissenschaft"]
        },
        {
            "query": "Vergleiche die Beiträge von Isaac Newton und Albert Einstein zur modernen Wissenschaft. Inwiefern hat Einstein Newtons Ideen erweitert?",
            "expected_keywords": ["Newton", "Einstein", "Gravitation", "Relativität", "Bewegungsgesetze"]
        },
        {
            "query": "Welche Bedeutung haben Wissenschaft und Bildung (wie sie in den Artikeln 'Science' und 'Education' beschrieben werden) für die Entwicklung von Nationen wie den USA, dem Vereinigten Königreich oder Indien?",
            "expected_keywords": ["USA", "Indien", "Education", "Science", "Universitäten", "Entwicklung"]
        }
    ]

    for i, t in enumerate(test_queries, 1):
        print(f"\n[Test {i}] Frage: {t['query']}")
        print("Antwort:\n" + "="*20)
        
        full_answer = ""
        for token in answer_with_rag(t['query'], []):
            full_answer += token
            sys.stdout.write(token)
            sys.stdout.flush()
        
        print("\n" + "="*20)
        
        # Validierung
        found = [kw for kw in t['expected_keywords'] if kw.lower() in full_answer.lower()]
        print(f"Gefundene Schlüsselwörter: {found}")
        if len(found) >= 2:
            print("ERGEBNIS: PASS")
        else:
            print("ERGEBNIS: FAIL")

    print("\n--- Alle Wiki-RAG Tests abgeschlossen ---")

if __name__ == "__main__":
    run_wiki_rag_test()