from __future__ import annotations

import unittest
from unittest.mock import patch

from agents import executor as executor_module
from agents.executor import Executor


class FailingMemory:
    async def save_episode(self, *_args, **_kwargs) -> None:
        raise RuntimeError("memory unavailable")


class RecordingMemory:
    def __init__(self) -> None:
        self.calls = 0

    async def save_episode(self, *_args, **_kwargs) -> None:
        self.calls += 1


class ExecutorSideEffectRetryTests(unittest.IsolatedAsyncioTestCase):
    async def test_memory_failure_does_not_retry_completed_side_effect(self) -> None:
        calls = 0

        async def non_idempotent_tool(**_inputs):
            nonlocal calls
            calls += 1
            return {"created_id": "resource-1"}

        with patch.dict(
            executor_module.TOOL_REGISTRY,
            {"non_idempotent_tool": {"required_inputs": [], "fallbacks": [], "_fn": non_idempotent_tool}},
            clear=False,
        ):
            result = await Executor(llm_client=object(), memory=FailingMemory(), max_retries=2).run_tool(
                "non_idempotent_tool", {}, timeout=2, worker_hint="test",
            )

        self.assertTrue(result["success"])
        self.assertEqual(result["attempt"], 1)
        self.assertEqual(calls, 1)
        self.assertFalse(result["memory_persisted"])
        self.assertIn("memory unavailable", result["memory_error"])

    async def test_tool_failure_still_retries_before_any_side_effect_result(self) -> None:
        calls = 0

        async def flaky_tool(**_inputs):
            nonlocal calls
            calls += 1
            if calls == 1:
                raise RuntimeError("transient tool failure")
            return {"ok": True}

        with patch.dict(
            executor_module.TOOL_REGISTRY,
            {"flaky_tool": {"required_inputs": [], "fallbacks": [], "_fn": flaky_tool}},
            clear=False,
        ):
            result = await Executor(llm_client=object(), memory=None, max_retries=2).run_tool(
                "flaky_tool", {}, timeout=2, worker_hint="test",
            )

        self.assertTrue(result["success"])
        self.assertEqual(result["attempt"], 2)
        self.assertEqual(calls, 2)

    async def test_successful_memory_persistence_is_reported(self) -> None:
        memory = RecordingMemory()

        async def safe_tool(**_inputs):
            return "done"

        with patch.dict(
            executor_module.TOOL_REGISTRY,
            {"safe_tool": {"required_inputs": [], "fallbacks": [], "_fn": safe_tool}},
            clear=False,
        ):
            result = await Executor(llm_client=object(), memory=memory, max_retries=0).run_tool(
                "safe_tool", {}, timeout=2, worker_hint="test",
            )

        self.assertTrue(result["success"])
        self.assertTrue(result["memory_persisted"])
        self.assertEqual(memory.calls, 1)


if __name__ == "__main__":
    unittest.main()