import os
import sys
import json
import tempfile
import shutil
from pathlib import Path
from unittest import mock

# Ensure the environment variable is set before importing the module
os.environ["OPENAI_API_KEY"] = "test-key"

# Import the module after setting the environment variable
import src.index as index

# Helper to create a temporary text file
def create_temp_txt(folder: Path, name: str, content: str):
    file_path = folder / name
    file_path.write_text(content, encoding="utf-8")
    return file_path

def test_split_text_basic():
    text = "Sentence one. Sentence two. Sentence three."
    chunks = index._split_text(text, max_chunk_size=50)
    assert len(chunks) == 3
    assert chunks[0] == "Sentence one."
    assert chunks[1] == "Sentence two."
    assert chunks[2] == "Sentence three."

def test_split_text_long_sentence():
    long_sentence = "A" * 200
    text = f"{long_sentence}. Another short sentence."
    chunks = index._split_text(text, max_chunk_size=100)
    # The long sentence should be split into two chunks
    assert len(chunks) == 2
    assert chunks[0].startswith("A" * 100)
    assert chunks[1].startswith("A" * 100)

def test_ingest_folder(monkeypatch):
    # Create a temporary directory with a single .txt file
    temp_dir = Path(tempfile.mkdtemp())
    try:
        content = "Hello world. This is a test."
        create_temp_txt(temp_dir, "test.txt", content)

        # Mock the add_documents method to capture its arguments
        captured = {}
        def mock_add_documents(self, documents, ids):
            captured["documents"] = documents
            captured["ids"] = ids

        monkeypatch.setattr(index.ChromaDBWrapper, "add_documents", mock_add_documents)

        # Instantiate the agent with a dummy db wrapper
        dummy_db = index.ChromaDBWrapper()
        agent = index.RAGAgent(dummy_db)

        # Run ingestion
        agent.ingest_folder(str(temp_dir))

        # Verify that documents were split and added
        assert "documents" in captured
        assert "ids" in captured
        assert len(captured["documents"]) == 2  # two sentences
        assert captured["documents"][0] == "Hello world."
        assert captured["documents"][1] == "This is a test."
        assert len(captured["ids"]) == 2
        assert captured["ids"][0].startswith("test_")
    finally:
        shutil.rmtree(temp_dir)

def test_answer_query_local(monkeypatch):
    # Dummy database that returns a relevant document
    class DummyDB:
        def query(self, query_text, k=5):
            return [("Relevant context about Python.", 0.8)]

    dummy_db = DummyDB()
    agent = index.RAGAgent(dummy_db)

    # Mock the OpenAI ChatCompletion to return a predictable answer
    mock_response = {
        "choices": [
            {"message": {"content": "Python is a programming language."}}
        ]
    }
    monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)

    answer = agent.answer_query("What is Python?")
    assert answer == "Python is a programming language."

def test_answer_query_fallback(monkeypatch):
    # Dummy database that returns no relevant documents
    class DummyDB:
        def query(self, query_text, k=5):
            return []

    dummy_db = DummyDB()
    agent = index.RAGAgent(dummy_db)

    # Mock web_search to return snippets
    monkeypatch.setattr(index, "web_search", lambda query, max_results=3: ["Snippet about AI.", "Another snippet."])

    # Mock the OpenAI ChatCompletion to return a predictable answer
    mock_response = {
        "choices": [
            {"message": {"content": "AI stands for Artificial Intelligence."}}
        ]
    }
    monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)

    answer = agent.answer_query("What does AI stand for?")
    assert answer == "AI stands for Artificial Intelligence."

def test_web_search_mock(monkeypatch):
    # Mock ddg to return predefined results
    mock_results = [
        {"body": "First snippet content."},
        {"body": "Second snippet content."},
    ]
    monkeypatch.setattr(index.ddg, "__call__", lambda query, max_results=3: mock_results)

    snippets = index.web_search("test query")
    assert snippets == ["First snippet content.", "Second snippet content."]

if __name__ == "__main__":
    # Run tests manually if executed as a script
    import pytest
    sys.exit(pytest.main([__file__]))