import os
import sys
from unittest.mock import MagicMock

# Ensure src is in path
sys.path.insert(0, os.path.abspath("src"))

from openai_resptools import ToolRegistry, ToolRunner
from openai_resptools.errors import ToolExecutionError


def test_max_retry():
    # Setup
    registry = ToolRegistry()

    @registry.tool()
    def failing_tool():
        """Always fails."""
        raise ValueError("Something went wrong")

    # Mock client
    client = MagicMock()

    # We need to simulate the loop.
    # The runner calls client.responses.create(**kwargs)
    # We need to return a response that calls 'failing_tool'.

    # We'll make the mock return a response with a tool call.
    # On subsequent calls (retries), it should also return a tool call (simulating model retry).

    call_item = {
        "type": "function_call",
        "call_id": "call_123",
        "name": "failing_tool",
        "arguments": "{}",
    }

    mock_response = MagicMock()
    mock_response.output = [call_item]
    mock_response.output_text = "Calling tool..."

    client.responses.create.return_value = mock_response

    # Test with max_retry = 2
    runner = ToolRunner(
        client, model="gpt-4", registry=registry, max_iters=5, max_retry=2
    )

    print("Running with max_retry=2...")
    try:
        runner.run("Please fail")
    except ToolExecutionError:
        print("Caught expected ToolExecutionError!")
    else:
        print("Error: Did not raise ToolExecutionError!")
        sys.exit(1)

    # Verify calls
    # We expect:
    # 1. Initial call (iteration 0) -> Tool fails -> Returns Error -> consecutive_errors=1
    # 2. Model call (iteration 1) -> Tool fails -> Returns Error -> consecutive_errors=2
    # 3. Model call (iteration 2) -> Tool fails -> Raise

    # client.responses.create should have been called 3 times (initial + 2 retries)
    # Actually:
    # Iteration 0: creates response (mock_response). Processing response calls tool. Tool fails. Output appended.
    # Iteration 1: creates response. Processing calls tool. Tool fails. Output appended.
    # Iteration 2: creates response. Processing calls tool. Tool fails. Raises.

    print(f"Client create called {client.responses.create.call_count} times.")
    if client.responses.create.call_count == 3:
        print("Success: Retried exactly 2 times before failing on 3rd attempt.")
    else:
        print(f"Failure: Expected 3 calls, got {client.responses.create.call_count}")
        sys.exit(1)


def test_recovery():
    # Test that consecutive errors reset on success
    registry = ToolRegistry()

    # We'll use a side effect to fail once then succeed
    fail_count = 0

    @registry.tool()
    def flaky_tool():
        nonlocal fail_count
        fail_count += 1
        if fail_count <= 1:
            raise ValueError("Transient error")
        return "Success"

    client = MagicMock()

    call_item = {
        "type": "function_call",
        "call_id": "call_123",
        "name": "flaky_tool",
        "arguments": "{}",
    }

    # Response 1: calls tool (fails)
    resp1 = MagicMock()
    resp1.output = [call_item]

    # Response 2: calls tool (succeeds)
    resp2 = MagicMock()
    resp2.output = [call_item]

    # Response 3: finished
    resp3 = MagicMock()
    resp3.output = []
    resp3.output_text = "Done"

    client.responses.create.side_effect = [resp1, resp2, resp3]

    runner = ToolRunner(
        client, model="gpt-4", registry=registry, max_iters=5, max_retry=2
    )

    print("\nRunning recovery test...")
    result = runner.run("Try flaky tool")

    print("Recovery test finished successfully.")
    print(f"Final text: {result.text}")
    print(f"Fail count: {fail_count}")

    if fail_count == 2:  # 1 fail, 1 success
        print("Success: Tool failed once then succeeded.")
    else:
        print(f"Unexpected fail count: {fail_count}")


if __name__ == "__main__":
    test_max_retry()
    test_recovery()
