# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

import os
import tempfile
from collections.abc import Generator
from typing import Annotated

import pytest

from autogen.import_utils import optional_import_block, run_for_optional_imports

# Import the tool itself
# Ensure this path is correct for your AG2 framework structure
from autogen.tools.experimental.reliable import ReliableTool
from autogen.tools.experimental.reliable.reliable import ToolExecutionDetails

# AG2 framework specific imports for testing
# Adjust the path to conftest.py if necessary based on your test file's location
from test.credentials import Credentials

with optional_import_block() as result:
    import openai  # noqa: F401

sub_question_runner_system_message_addition = "You are an assistant that helps break down questions. Your goal is to generate exactly 3 relevant sub-questions based on the main task. Use the provided `generate_sub_questions_list` tool to output the list you generate. Provide the generated list as the 'sub_questions' argument."
sub_question_validator_system_message_addition = """You are a quality control assistant. Your task is to validate the output received, which should be a list of sub-questions generated by another assistant.

**Validation Criteria:**
1.  **Correct Format:** The output MUST be a list of strings.
2.  **Correct Quantity:** The list MUST contain exactly 3 sub-questions.
3.  **Relevance:** Each sub-question MUST be clearly relevant to the original main question described in the initial task.
"""


@pytest.fixture(autouse=True)
def setup_test_environment() -> Generator[None, None, None]:
    """Set up a temporary directory for SQLite database."""
    with tempfile.TemporaryDirectory() as temp_dir:
        # Change to temporary directory so logs.db gets created there
        original_cwd = os.getcwd()
        os.chdir(temp_dir)
        try:
            yield
        finally:
            os.chdir(original_cwd)


class TestReliableTool:
    @run_for_optional_imports("openai", "openai")
    def test_bad_response(self, credentials_gpt_4o_mini: Credentials) -> None:
        # Skip this test in GitHub Actions due to SQLite database permission issues
        if os.getenv("GITHUB_ACTIONS") == "true":
            pytest.skip("Skipping ReliableTool test in GitHub Actions due to SQLite database permission issues")

        should_bad_response = True

        def generate_sub_questions_list(
            sub_questions: Annotated[list[str], "A list of sub-questions related to the main question."],
        ) -> list[str]:
            """
            Receives and returns a list of generated sub-questions.
            """
            nonlocal should_bad_response
            if should_bad_response:
                should_bad_response = False
                return []
            return sub_questions

        sub_question_tool = ReliableTool(
            name="SubQuestionGenerator",
            func_or_tool=generate_sub_questions_list,
            description="Reliably generates exactly 3 relevant sub-questions for a given main question.",
            runner_llm_config=credentials_gpt_4o_mini.llm_config,
            validator_llm_config=credentials_gpt_4o_mini.llm_config,
            system_message_addition_for_tool_calling=sub_question_runner_system_message_addition,
            system_message_addition_for_result_validation=sub_question_validator_system_message_addition,
            max_tool_invocations=5,
        )

        result: ToolExecutionDetails = sub_question_tool.run_and_get_details(
            task="How does photosynthesis work in plants?"
        )

        # Should fail once, then pass because of should_bad_response
        assert result.final_tool_context.attempt_count == 2
        assert not should_bad_response

    @run_for_optional_imports("openai", "openai")
    def test_error(self, credentials_gpt_4o_mini: Credentials) -> None:
        # Skip this test in GitHub Actions due to SQLite database permission issues
        if os.getenv("GITHUB_ACTIONS") == "true":
            pytest.skip("Skipping ReliableTool test in GitHub Actions due to SQLite database permission issues")

        should_error = True

        def generate_sub_questions_list(
            sub_questions: Annotated[list[str], "A list of sub-questions related to the main question."],
        ) -> list[str]:
            """
            Receives and returns a list of generated sub-questions.
            """
            nonlocal should_error
            if should_error:
                should_error = False
                raise Exception("Test Error")
            return sub_questions

        sub_question_tool = ReliableTool(
            name="SubQuestionGenerator",
            func_or_tool=generate_sub_questions_list,
            description="Reliably generates exactly 3 relevant sub-questions for a given main question.",
            runner_llm_config=credentials_gpt_4o_mini.llm_config,
            validator_llm_config=credentials_gpt_4o_mini.llm_config,
            system_message_addition_for_tool_calling=sub_question_runner_system_message_addition,
            system_message_addition_for_result_validation=sub_question_validator_system_message_addition,
            max_tool_invocations=5,
        )

        result: ToolExecutionDetails = sub_question_tool.run_and_get_details(
            task="How does photosynthesis work in plants?"
        )

        # Should fail once, then pass because of should_error
        assert result.final_tool_context.attempt_count == 2
        assert not should_error

    @run_for_optional_imports("openai", "openai")
    def test_return_tuple(self, credentials_gpt_4o_mini: Credentials) -> None:
        # Skip this test in GitHub Actions due to SQLite database permission issues
        if os.getenv("GITHUB_ACTIONS") == "true":
            pytest.skip("Skipping ReliableTool test in GitHub Actions due to SQLite database permission issues")

        should_error = True

        def generate_sub_questions_list(
            sub_questions: Annotated[list[str], "A list of sub-questions related to the main question."],
        ) -> tuple[list[str], str]:
            """
            Receives and returns a list of generated sub-questions.
            """
            nonlocal should_error
            if should_error:
                should_error = False
                raise Exception("Test Error")
            return sub_questions, "Sub Questions are:\n" + "\n".join(sub_questions)

        sub_question_tool = ReliableTool(
            name="SubQuestionGenerator",
            func_or_tool=generate_sub_questions_list,
            description="Reliably generates exactly 3 relevant sub-questions for a given main question.",
            runner_llm_config=credentials_gpt_4o_mini.llm_config,
            validator_llm_config=credentials_gpt_4o_mini.llm_config,
            system_message_addition_for_tool_calling=sub_question_runner_system_message_addition,
            system_message_addition_for_result_validation=sub_question_validator_system_message_addition,
            max_tool_invocations=5,
        )

        result: ToolExecutionDetails = sub_question_tool.run_and_get_details(
            task="How does photosynthesis work in plants?"
        )

        # Should fail once, then pass because of should_error
        assert result.final_tool_context.attempt_count == 2
        assert not should_error
        assert isinstance(result.final_tool_context.get_final_result_data(), list)
        assert isinstance(result.final_tool_context.get_final_result_str(), str)
