import asyncio
from copy import deepcopy

from behave import given, then, when

from cleveragents.agents.base import Agent
from cleveragents.langgraph.nodes import Node, NodeConfig, NodeType, ToolAgent
from cleveragents.langgraph.state import GraphState


def _make_state(messages=None, metadata=None):
    return GraphState(messages=messages or [], metadata=metadata or {})


class RecordingToolAgent(ToolAgent):
    def __init__(self, name: str, response: str):
        super().__init__(name)
        self.response = response
        self.received_inputs = []

    async def process_message(self, message, context=None):
        self.received_inputs.append(message)
        return self.response

    def get_capabilities(self):
        return []


class RecordingAgent(Agent):
    def __init__(self, name: str, response: str):
        super().__init__(name)
        self.response = response
        self.received_inputs = []
        self.received_contexts = []

    async def process_message(self, message, context=None):
        self.received_inputs.append(message)
        if context is not None:
            self.received_contexts.append(deepcopy(context))
        return self.response

    def get_capabilities(self):
        return []


class ClearingAgent(Agent):
    def __init__(self, name: str, response: str):
        super().__init__(name)
        self.response = response

    async def process_message(self, message, context=None):
        if context is not None:
            context.clear()
        return self.response

    def get_capabilities(self):
        return []


class MutatingAgent(Agent):
    def __init__(self, name: str, response: str):
        super().__init__(name)
        self.response = response

    async def process_message(self, message, context=None):
        if context is not None:
            context["mutated"] = "yes"
        return self.response

    def get_capabilities(self):
        return []


@given("a tool node configured")
def step_tool_node_configured(context):
    config = NodeConfig(name="tool_node", type=NodeType.TOOL)
    context.node = Node(config)
    context.state = _make_state()


@when("I execute the tool node")
def step_execute_tool_node(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("it should mark tool executed")
def step_assert_tool_executed(context):
    assert context.result.get("tool_executed") is True
    assert context.result.get("node") == "tool_node"


@given('a tool agent node using ToolAgent with current_message "{text}"')
def step_toolagent_branch(context, text):
    tool_agent = RecordingToolAgent("tool", response="ok")
    config = NodeConfig(name="tool_agent", type=NodeType.AGENT, agent="tool")
    context.node = Node(config, agents={"tool": tool_agent})
    context.state = _make_state(
        messages=[{"role": "assistant", "content": "fallback"}],
        metadata={"current_message": text},
    )


@when("I execute the tool agent node with ToolAgent branch")
def step_execute_toolagent_branch(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then('the ToolAgent should receive input "{text}"')
def step_assert_toolagent_input(context, text):
    agent = context.node.agents["tool"]
    assert agent.received_inputs[-1] == text


@given("an agent node with only assistant history")
def step_agent_assistant_only(context):
    agent = RecordingAgent("rec", response="ok")
    config = NodeConfig(name="assistant_only", type=NodeType.AGENT, agent="rec")
    context.node = Node(config, agents={"rec": agent})
    context.state = _make_state(messages=[{"role": "assistant", "content": "assist"}])


@when("I execute the agent node for assistant fallback")
def step_execute_agent_assistant_only(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("the agent should receive the assistant fallback content")
def step_assert_assistant_fallback(context):
    agent = context.node.agents["rec"]
    assert agent.received_inputs[-1] == "assist"


@given("an agent that clears context during processing")
def step_agent_clears_context(context):
    agent = ClearingAgent("clear", response="done")
    config = NodeConfig(name="clear_agent", type=NodeType.AGENT, agent="clear")
    context.node = Node(config, agents={"clear": agent})
    context.state = _make_state(
        messages=[{"role": "user", "content": "hi"}], metadata={"foo": "bar"}
    )


@when("I execute the clearing agent node")
def step_execute_clearing_agent(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("only last_agent_node metadata should remain")
def step_assert_cleared_metadata(context):
    metadata = context.result.get("metadata", {})
    assert metadata == {"last_agent_node": "clear_agent"}


@given("an agent that mutates context to add metadata")
def step_agent_mutates_context(context):
    agent = MutatingAgent("mut", response="done")
    config = NodeConfig(name="mut_agent", type=NodeType.AGENT, agent="mut")
    context.node = Node(config, agents={"mut": agent})
    context.state = _make_state(messages=[{"role": "user", "content": "hello"}])


@when("I execute the mutating agent node")
def step_execute_mutating_agent(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("the metadata should include the mutated key")
def step_assert_mutated_metadata(context):
    metadata = context.result.get("metadata", {})
    assert metadata.get("mutated") == "yes"


@given('a subgraph node named "{subgraph_name}"')
def step_subgraph_node(context, subgraph_name):
    config = NodeConfig(name="subgraph", type=NodeType.SUBGRAPH, subgraph=subgraph_name)
    context.node = Node(config)
    context.state = _make_state()


@when("I execute the subgraph node")
def step_execute_subgraph(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then('it should report the invoked subgraph "{subgraph_name}"')
def step_assert_subgraph(context, subgraph_name):
    assert context.result.get("subgraph_invoked") == subgraph_name


@given("a start node")
def step_start_node(context):
    config = NodeConfig(name="start", type=NodeType.START)
    context.node = Node(config)
    context.state = _make_state()


@when("I execute the start node")
def step_execute_start(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("it should indicate graph start")
def step_assert_started(context):
    assert context.result.get("started") is True


@given("an end node")
def step_end_node(context):
    config = NodeConfig(name="end", type=NodeType.END)
    context.node = Node(config)
    context.state = _make_state()


@when("I execute the end node")
def step_execute_end(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("it should indicate graph completion")
def step_assert_completed(context):
    assert context.result.get("completed") is True


@given("a tool node patched to return non dict")
def step_tool_node_patched(context):
    config = NodeConfig(name="patched_tool", type=NodeType.TOOL)
    context.node = Node(config)

    async def _patched_execute_tool():
        return "raw-value"

    context.node._execute_tool = _patched_execute_tool  # type: ignore[attr-defined]
    context.state = _make_state()


@when("I execute the patched tool node")
def step_execute_patched_tool(context):
    loop = asyncio.new_event_loop()
    try:
        context.result = loop.run_until_complete(context.node.execute(context.state))
    finally:
        loop.close()


@then("the execute result should only include current_node")
def step_assert_current_node_only(context):
    assert context.result == {"current_node": "patched_tool"}