import asyncio
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given, then, when

from cleveragents.agents.base import Agent, AgentWithMemory
from cleveragents.core.exceptions import ExecutionError


class DefaultAgent(Agent):
    def __init__(
        self, name: str, config: dict[str, Any] | None = None, raise_error: bool = False
    ):
        self.raise_error = raise_error
        super().__init__(name, config)
        self.last_message = None
        self.last_context = None

    async def process_message(
        self, message: Any, context: dict[str, Any] | None = None
    ) -> str:
        self.last_message = message
        self.last_context = context
        if self.raise_error:
            raise RuntimeError("boom")
        return f"processed:{message}"

    def get_capabilities(self) -> list[str]:
        return ["default"]


class MemoryAgent(AgentWithMemory):
    async def process_message(
        self, message: Any, context: dict[str, Any] | None = None
    ) -> Any:
        await self.remember("last", message)
        return f"mem:{message}"

    def get_capabilities(self) -> list[str]:
        return ["memory"]


@given("a basic test agent setup")
def step_basic_agent(context):
    context.loop = asyncio.new_event_loop()
    asyncio.set_event_loop(context.loop)
    context.agent = DefaultAgent("base-test")
    context.captured_output = []


@when("I patch asyncio create_task and send a message")
def step_patch_create_task(context):
    with patch("asyncio.create_task") as mock_task:
        context.agent.send_message("hello", {"foo": "bar"})
        context.create_task_called_with = None
        if mock_task.call_args:
            context.create_task_called_with = mock_task.call_args[0][0]


@then("the processing pipeline should schedule the process wrapper task")
def step_verify_create_task_called(context):
    assert context.create_task_called_with is not None
    assert context.create_task_called_with.__qualname__.endswith("_process_wrapper")


@when("I process a tuple message through the wrapper")
def step_process_tuple(context):
    context.captured_output = []
    context.agent.output_stream.subscribe(
        lambda value: context.captured_output.append(value)
    )
    asyncio.run(context.agent._process_wrapper(("tuple-msg", {"ctx": 1})))


@then("the tuple message and context should be passed and emitted")
def step_verify_tuple_processing(context):
    assert context.agent.last_message == "tuple-msg"
    assert context.agent.last_context == {"ctx": 1}
    assert context.captured_output == ["processed:tuple-msg"]


@when("I process a non tuple message through the wrapper")
def step_process_non_tuple(context):
    context.captured_output = []
    context.agent.output_stream.subscribe(
        lambda value: context.captured_output.append(value)
    )
    asyncio.run(context.agent._process_wrapper("plain-msg"))


@then("the non tuple message should be processed with empty context")
def step_verify_non_tuple_processing(context):
    assert context.agent.last_message == "plain-msg"
    assert context.agent.last_context == {}
    assert context.captured_output == ["processed:plain-msg"]


@when("process_message raises an error inside the wrapper")
def step_wrapper_raises(context):
    error_agent = DefaultAgent("err-agent", raise_error=True)
    context.wrapper_error = None
    try:
        asyncio.run(error_agent._process_wrapper("fail"))
    except Exception as exc:
        context.wrapper_error = exc


@then("an ExecutionError should be raised with the agent name")
def step_verify_execution_error(context):
    assert isinstance(context.wrapper_error, ExecutionError)
    assert "err-agent" in str(context.wrapper_error)


@when("I call process_message_sync on the agent")
def step_call_process_message_sync(context):
    loop = context.loop
    asyncio.set_event_loop(loop)
    context.sync_result = context.agent.process_message_sync("sync-msg", {"sync": True})


@then("the synchronous result should match the async process output")
def step_verify_sync_result(context):
    assert context.sync_result == "processed:sync-msg"
    assert context.agent.last_context == {"sync": True}


@when("I call dispose on an agent with disposable streams")
def step_call_dispose(context):
    context.agent.input_stream = MagicMock(dispose=MagicMock())
    context.agent.output_stream = MagicMock(dispose=MagicMock())
    context.agent.dispose()


@then("both input and output streams should be disposed if available")
def step_verify_dispose(context):
    context.agent.input_stream.dispose.assert_called_once()
    context.agent.output_stream.dispose.assert_called_once()


@when("I remember a value and recall it later")
def step_memory_remember_recall(context):
    context.memory_agent = MemoryAgent("mem-agent")
    asyncio.run(context.memory_agent.remember("answer", 42))
    context.recalled = asyncio.run(context.memory_agent.recall("answer"))


@then("the remembered value should be returned from memory")
def step_verify_memory_recall(context):
    assert context.recalled == 42