"""
Unit tests for the RAG agent and auto-check graph.
"""

import os
import json
import tempfile
import shutil
import pytest

from src.index import RAGAgent, auto_check_graph

# Helper to create a temporary agent with fake embeddings/LLM
def create_temp_agent(tmp_path):
    # Ensure no OpenAI key
    os.environ.pop("OPENAI_API_KEY", None)
    agent = RAGAgent(
        vector_store_path=tmp_path / "vector_store.faiss",
        documents_dir=tmp_path / "docs",
    )
    return agent

def test_add_and_query():
    tmp_dir = tempfile.mkdtemp()
    try:
        agent = create_temp_agent(tmp_dir)
        docs = [
            "The capital of France is Paris.",
            "William Shakespeare wrote Hamlet.",
        ]
        agent.add_documents(docs)
        # Query for first doc
        answer = agent.query("What is the capital of France?")
        assert "Paris" in answer
        # Query for second doc
        answer2 = agent.query("Who wrote Hamlet?")
        assert "Shakespeare" in answer2
    finally:
        shutil.rmtree(tmp_dir)

def test_auto_check_pass():
    tmp_dir = tempfile.mkdtemp()
    try:
        agent = create_temp_agent(tmp_dir)
        docs = [
            "The capital of France is Paris.",
            "William Shakespeare wrote Hamlet.",
        ]
        agent.add_documents(docs)
        ground_truth = {
            "What is the capital of France?": "Paris",
            "Who wrote Hamlet?": "William Shakespeare",
        }
        result = auto_check_graph(
            "What is the capital of France?", agent, ground_truth
        )
        assert result["verdict_row"] == "PASS"
        assert "Paris" in result["answer"]
    finally:
        shutil.rmtree(tmp_dir)

def test_auto_check_fail():
    tmp_dir = tempfile.mkdtemp()
    try:
        agent = create_temp_agent(tmp_dir)
        docs = [
            "The capital of France is Paris.",
        ]
        agent.add_documents(docs)
        ground_truth = {
            "What is the capital of France?": "Berlin",
        }
        result = auto_check_graph(
            "What is the capital of France?", agent, ground_truth
        )
        assert result["verdict_row"] == "FAIL"
    finally:
        shutil.rmtree(tmp_dir)

def test_auto_check_unknown():
    tmp_dir = tempfile.mkdtemp()
    try:
        agent = create_temp_agent(tmp_dir)
        docs = [
            "The capital of France is Paris.",
        ]
        agent.add_documents(docs)
        ground_truth = {
            "What is the capital of Germany?": "Berlin",
        }
        result = auto_check_graph(
            "What is the capital of Germany?", agent, ground_truth
        )
        assert result["verdict_row"] == "UNKNOWN"
    finally:
        shutil.rmtree(tmp_dir)

if __name__ == "__main__":
    pytest.main([__file__])