import pytest
from unittest.mock import AsyncMock, patch, MagicMock

from cognee.modules.retrieval.graph_completion_cot_retriever import (
    GraphCompletionCotRetriever,
    _as_answer_text,
)
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
from cognee.infrastructure.llm.LLMGateway import LLMGateway


@pytest.fixture
def mock_edge():
    """Create a mock edge."""
    edge = MagicMock(spec=Edge)
    return edge


@pytest.mark.asyncio
async def test_get_triplets_inherited(mock_edge):
    """Test that get_triplets is inherited from parent class."""
    retriever = GraphCompletionCotRetriever()

    with patch(
        "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
        return_value=[mock_edge],
    ):
        triplets = await retriever.get_triplets("test query")

    assert len(triplets) == 1
    assert triplets[0] == mock_edge


@pytest.mark.asyncio
async def test_init_custom_params():
    """Test GraphCompletionCotRetriever initialization with custom parameters."""
    retriever = GraphCompletionCotRetriever(
        top_k=10,
        user_prompt_path="custom_user.txt",
        system_prompt_path="custom_system.txt",
        validation_user_prompt_path="custom_validation_user.txt",
        validation_system_prompt_path="custom_validation_system.txt",
        followup_system_prompt_path="custom_followup_system.txt",
        followup_user_prompt_path="custom_followup_user.txt",
    )

    assert retriever.top_k == 10
    assert retriever.user_prompt_path == "custom_user.txt"
    assert retriever.system_prompt_path == "custom_system.txt"
    assert retriever.validation_user_prompt_path == "custom_validation_user.txt"
    assert retriever.validation_system_prompt_path == "custom_validation_system.txt"
    assert retriever.followup_system_prompt_path == "custom_followup_system.txt"
    assert retriever.followup_user_prompt_path == "custom_followup_user.txt"


@pytest.mark.asyncio
async def test_init_defaults():
    """Test GraphCompletionCotRetriever initialization with defaults."""
    retriever = GraphCompletionCotRetriever()

    assert retriever.validation_user_prompt_path == "cot_validation_user_prompt.txt"
    assert retriever.validation_system_prompt_path == "cot_validation_system_prompt.txt"
    assert retriever.followup_system_prompt_path == "cot_followup_system_prompt.txt"
    assert retriever.followup_user_prompt_path == "cot_followup_user_prompt.txt"


@pytest.mark.asyncio
async def test_run_cot_completion_round_zero(mock_edge):
    """Test _run_cot_completion round 0."""
    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever._as_answer_text",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.render_prompt",
            return_value="Rendered prompt",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.read_query_prompt",
            return_value="System prompt",
        ),
        patch.object(retriever, "get_triplets", new_callable=AsyncMock, return_value=[[mock_edge]]),
        patch.object(
            LLMGateway,
            "acreate_structured_output",
            new_callable=AsyncMock,
            side_effect=["validation_result", "followup_question"],
        ),
    ):
        completion, context_text, triplets = await retriever._run_cot_completion(
            query_batch=["test query"],
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert completion[0] == "Generated answer"
    assert isinstance(context_text, list)
    assert len(context_text) == 1
    assert context_text[0] == "Resolved context"
    assert len(triplets) >= 1


@pytest.mark.asyncio
async def test_run_cot_completion_multiple_rounds(mock_edge):
    """Test _run_cot_completion with multiple rounds."""
    retriever = GraphCompletionCotRetriever(max_iter=2)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.render_prompt",
            return_value="Rendered prompt",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.read_query_prompt",
            return_value="System prompt",
        ),
        patch.object(
            LLMGateway,
            "acreate_structured_output",
            new_callable=AsyncMock,
            side_effect=[
                "validation_result",
                "followup_question",
                "validation_result2",
                "followup_question2",
            ],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever._as_answer_text",
            return_value="Generated answer",
        ),
        patch.object(retriever, "get_triplets", new_callable=AsyncMock, return_value=[[mock_edge]]),
    ):
        completion, context_text, triplets = await retriever._run_cot_completion(
            query_batch=["test query"]
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert completion[0] == "Generated answer"
    assert isinstance(context_text, list)
    assert len(context_text) == 1
    assert context_text[0] == "Resolved context"
    assert len(triplets) >= 1


@pytest.mark.asyncio
async def test_run_cot_completion_with_conversation_history(mock_edge):
    """Test _run_cot_completion with conversation history."""
    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ) as mock_generate,
        patch.object(retriever, "get_triplets", new_callable=AsyncMock, return_value=[[mock_edge]]),
    ):
        completion, context_text, triplets = await retriever._run_cot_completion(
            query_batch=["test query"],
            conversation_history="Previous conversation",
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert completion[0] == "Generated answer"
    call_kwargs = mock_generate.call_args[1]
    assert call_kwargs.get("conversation_history") == "Previous conversation"


@pytest.mark.asyncio
async def test_run_cot_completion_with_response_model(mock_edge):
    """Test _run_cot_completion with custom response model."""
    from pydantic import BaseModel

    class TestModel(BaseModel):
        answer: str

    retriever = GraphCompletionCotRetriever(response_model=TestModel, max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value=TestModel(answer="Test answer"),
        ),
        patch.object(retriever, "get_triplets", new_callable=AsyncMock, return_value=[[mock_edge]]),
    ):
        completion, context_text, triplets = await retriever._run_cot_completion(
            query_batch=["test query"]
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert isinstance(completion[0], TestModel)
    assert completion[0].answer == "Test answer"


@pytest.mark.asyncio
async def test_run_cot_completion_empty_conversation_history(mock_edge):
    """Test _run_cot_completion with empty conversation history."""
    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ) as mock_generate,
        patch.object(retriever, "get_triplets", new_callable=AsyncMock, return_value=[[mock_edge]]),
    ):
        completion, context_text, triplets = await retriever._run_cot_completion(
            query_batch=["test query"],
            conversation_history="",
        )

    assert isinstance(completion, list)
    assert completion[0] == "Generated answer"
    # Verify conversation_history was passed as None when empty
    call_kwargs = mock_generate.call_args[1]
    assert call_kwargs.get("conversation_history") is None


@pytest.mark.asyncio
async def test_get_completion(mock_edge):
    """Test get_completion (CoT loop + final completion in get_completion_from_context)."""
    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.generate_completion_batch",
            new_callable=AsyncMock,
            return_value=["Generated answer"],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.generate_completion",
            new_callable=AsyncMock,
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever._as_answer_text",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.render_prompt",
            return_value="Rendered prompt",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.read_query_prompt",
            return_value="System prompt",
        ),
        patch.object(
            LLMGateway,
            "acreate_structured_output",
            new_callable=AsyncMock,
            side_effect=["validation_result", "followup_question"],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.CacheConfig"
        ) as mock_cache_config,
    ):
        mock_config = MagicMock()
        mock_config.caching = False
        mock_cache_config.return_value = mock_config

        objects = await retriever.get_retrieved_objects(query="test query")
        context = await retriever.get_context_from_objects(
            query="test query", retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query="test query", retrieved_objects=objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert completion[0] == "Generated answer"


@pytest.mark.asyncio
async def test_get_completion_with_session(mock_edge):
    """Test get_completion with session (SessionManager path for final completion)."""
    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(session_id="test_session", max_iter=1)

    mock_user = MagicMock()
    mock_user.id = "test-user-id"

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_session_manager",
        ) as mock_get_sm,
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.generate_completion_batch",
            new_callable=AsyncMock,
            return_value=["Generated answer"],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever._as_answer_text",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.render_prompt",
            return_value="Rendered prompt",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_cot_retriever.read_query_prompt",
            return_value="System prompt",
        ),
        patch.object(
            LLMGateway,
            "acreate_structured_output",
            new_callable=AsyncMock,
            side_effect=["validation_result", "followup_question"],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.CacheConfig"
        ) as mock_cache_config,
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.session_user"
        ) as mock_session_user,
    ):
        mock_config = MagicMock()
        mock_config.caching = True
        mock_cache_config.return_value = mock_config
        mock_session_user.get.return_value = mock_user
        mock_sm = MagicMock()
        mock_sm.get_session = AsyncMock(return_value="")
        mock_sm.session_history_last_n = 10
        mock_sm.generate_completion_with_session = AsyncMock(return_value="Generated answer")
        mock_get_sm.return_value = mock_sm

        retrieved_objects = await retriever.get_retrieved_objects("test query")
        context = await retriever.get_context_from_objects(
            query="test query", retrieved_objects=retrieved_objects
        )
        completion = await retriever.get_completion_from_context(
            query="test query", retrieved_objects=retrieved_objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert completion[0] == "Generated answer"
    mock_sm.generate_completion_with_session.assert_awaited_once()


@pytest.mark.asyncio
async def test_get_completion_with_response_model(mock_edge):
    """Test get_completion with custom response model."""
    from pydantic import BaseModel

    class TestModel(BaseModel):
        answer: str

    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(response_model=TestModel, max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value=TestModel(answer="Test answer"),
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.CacheConfig"
        ) as mock_cache_config,
    ):
        mock_config = MagicMock()
        mock_config.caching = False
        mock_cache_config.return_value = mock_config

        objects = await retriever.get_retrieved_objects("test query")
        context = await retriever.get_context_from_objects(
            query="test query", retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query="test query", retrieved_objects=objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 1
    assert isinstance(completion[0], TestModel)


@pytest.mark.asyncio
async def test_get_completion_with_session_no_user_id(mock_edge):
    """Test get_completion with session config but no user ID."""
    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.CacheConfig"
        ) as mock_cache_config,
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.session_user"
        ) as mock_session_user,
    ):
        mock_config = MagicMock()
        mock_config.caching = True
        mock_cache_config.return_value = mock_config
        mock_session_user.get.return_value = None  # No user

        objects = await retriever.get_retrieved_objects("test query")
        context = await retriever.get_context_from_objects(
            query="test query", retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query="test query", retrieved_objects=objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 1


@pytest.mark.asyncio
async def test_as_answer_text_with_typeerror():
    """Test _as_answer_text handles TypeError when json.dumps fails."""
    non_serializable = {1, 2, 3}

    result = _as_answer_text(non_serializable)

    assert isinstance(result, str)
    assert result == str(non_serializable)


@pytest.mark.asyncio
async def test_as_answer_text_with_string():
    """Test _as_answer_text with string input."""
    result = _as_answer_text("test string")
    assert result == "test string"


@pytest.mark.asyncio
async def test_as_answer_text_with_dict():
    """Test _as_answer_text with dictionary input."""
    test_dict = {"key": "value", "number": 42}
    result = _as_answer_text(test_dict)
    assert isinstance(result, str)
    assert "key" in result
    assert "value" in result


@pytest.mark.asyncio
async def test_as_answer_text_with_basemodel():
    """Test _as_answer_text with Pydantic BaseModel input."""
    from pydantic import BaseModel

    class TestModel(BaseModel):
        answer: str

    test_model = TestModel(answer="test answer")
    result = _as_answer_text(test_model)

    assert isinstance(result, str)
    assert "[Structured Response]" in result
    assert "test answer" in result


@pytest.mark.asyncio
async def test_get_completion_batch_queries(mock_edge):
    """Test get_completion batch queries without provided context."""
    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge], [mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ),
    ):
        objects = await retriever.get_retrieved_objects(
            query_batch=["test query 1", "test query 2"]
        )
        context = await retriever.get_context_from_objects(
            query_batch=["test query 1", "test query 2"], retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query_batch=["test query 1", "test query 2"], retrieved_objects=objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 2
    assert completion[0] == "Generated answer" and completion[1] == "Generated answer"


#
@pytest.mark.asyncio
async def test_get_completion_batch_queries_with_response_model(mock_edge):
    """Test get_completion of batch queries with custom response model."""
    from pydantic import BaseModel

    class TestModel(BaseModel):
        answer: str

    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever(response_model=TestModel, max_iter=1)

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge], [mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value=TestModel(answer="Test answer"),
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.CacheConfig"
        ) as mock_cache_config,
    ):
        mock_config = MagicMock()
        mock_config.caching = False
        mock_cache_config.return_value = mock_config

        objects = await retriever.get_retrieved_objects(
            query_batch=["test query 1", "test query 2"]
        )
        context = await retriever.get_context_from_objects(
            query_batch=["test query 1", "test query 2"], retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query_batch=["test query 1", "test query 2"], retrieved_objects=objects, context=context
        )

        assert isinstance(completion, list)
        assert len(completion) == 2
        assert isinstance(completion[0], TestModel) and isinstance(completion[1], TestModel)


@pytest.mark.asyncio
async def test_get_completion_batch_queries_duplicate_queries(mock_edge):
    """Test get_completion batch queries without provided context."""
    mock_graph_engine = AsyncMock()
    mock_graph_engine.is_empty = AsyncMock(return_value=False)

    unified_mock = AsyncMock()
    unified_mock.graph = mock_graph_engine
    unified_mock.vector = MagicMock()

    retriever = GraphCompletionCotRetriever()

    with (
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.get_unified_engine",
            return_value=unified_mock,
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.brute_force_triplet_search",
            return_value=[[mock_edge], [mock_edge]],
        ),
        patch(
            "cognee.modules.retrieval.graph_completion_retriever.resolve_edges_to_text",
            return_value="Resolved context",
        ),
        patch(
            "cognee.modules.retrieval.utils.completion.generate_completion",
            return_value="Generated answer",
        ),
    ):
        objects = await retriever.get_retrieved_objects(
            query_batch=["test query 1", "test query 1"]
        )
        context = await retriever.get_context_from_objects(
            query_batch=["test query 1", "test query 1"], retrieved_objects=objects
        )
        completion = await retriever.get_completion_from_context(
            query_batch=["test query 1", "test query 1"], retrieved_objects=objects, context=context
        )

    assert isinstance(completion, list)
    assert len(completion) == 2
    assert completion[0] == "Generated answer" and completion[1] == "Generated answer"
