"""Optimized step definitions for LangGraph nodes coverage tests."""

import asyncio
from unittest.mock import AsyncMock
from unittest.mock import MagicMock

from behave import given
from behave import then
from behave import when

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


@given("the LangGraph nodes system is initialized")
def step_init_nodes_system(context):
    """Initialize the nodes system."""
    context.test_results = {}
    context.mock_agent = MagicMock(spec=Agent)
    context.mock_agent.name = "test-agent"
    context.agents = {"test-agent": context.mock_agent}

    # Create reusable states
    context.empty_state = GraphState(messages=[])
    context.message_state = GraphState(
        messages=[{"role": "user", "content": "Test message"}]
    )
    context.question_state = GraphState(
        messages=[{"role": "user", "content": "What is this?"}]
    )
    context.multi_message_state = GraphState(
        messages=[
            {"role": "user", "content": "Message 1"},
            {"role": "user", "content": "Message 2"},
            {"role": "user", "content": "Message 3"},
        ]
    )
    context.metadata_state = GraphState(messages=[], metadata={"status": "ready"})


@given("I have various node configurations")
def step_various_configs(context):
    """Create various node configurations for testing."""
    context.node_configs = {
        "agent": NodeConfig(name="agent-node", type=NodeType.AGENT, agent="test-agent"),
        "agent_no_spec": NodeConfig(name="agent-no-spec", type=NodeType.AGENT),
        "agent_missing": NodeConfig(
            name="agent-missing", type=NodeType.AGENT, agent="missing"
        ),
        "function_summarize": NodeConfig(
            name="func-sum", type=NodeType.FUNCTION, function="summarize"
        ),
        "function_route": NodeConfig(
            name="func-route", type=NodeType.FUNCTION, function="route"
        ),
        "function_validate": NodeConfig(
            name="func-val", type=NodeType.FUNCTION, function="validate"
        ),
        "function_unknown": NodeConfig(
            name="func-unk", type=NodeType.FUNCTION, function="unknown"
        ),
        "function_none": NodeConfig(name="func-none", type=NodeType.FUNCTION),
        "tool_with": NodeConfig(
            name="tool-with", type=NodeType.TOOL, tools=["calc", "search"]
        ),
        "tool_empty": NodeConfig(name="tool-empty", type=NodeType.TOOL, tools=[]),
        "conditional": NodeConfig(
            name="cond", type=NodeType.CONDITIONAL, condition={"type": "always"}
        ),
        "subgraph": NodeConfig(name="sub", type=NodeType.SUBGRAPH, subgraph="test-sub"),
        "subgraph_none": NodeConfig(name="sub-none", type=NodeType.SUBGRAPH),
        "start": NodeConfig(name="start", type=NodeType.START),
        "end": NodeConfig(name="end", type=NodeType.END),
        "unknown": NodeConfig(name="unknown", type="unknown_type"),
        "retry": NodeConfig(
            name="retry",
            type=NodeType.FUNCTION,
            function="test",
            retry_policy={"max_retries": 2, "delay": 0},
        ),
        "parallel": NodeConfig(
            name="parallel", type=NodeType.FUNCTION, function="test", parallel=True
        ),
        "timeout": NodeConfig(
            name="timeout", type=NodeType.FUNCTION, function="test", timeout=5.0
        ),
    }


@when("I test node initialization and basic execution")
def step_test_initialization(context):
    """Test node initialization and basic execution for all types."""
    results = {}

    # Test initialization
    for name, config in context.node_configs.items():
        try:
            node = Node(config, context.agents)
            results[f"{name}_init"] = {
                "success": True,
                "name": node.name,
                "type": node.type,
                "execution_count": node.execution_count,
                "last_execution_time": node.last_execution_time,
                "last_error": node.last_error,
            }
        except Exception as e:
            results[f"{name}_init"] = {"success": False, "error": str(e)}

    context.test_results["initialization"] = results


@then("all node types should initialize correctly")
def step_verify_initialization(context):
    """Verify all nodes initialized correctly."""
    init_results = context.test_results["initialization"]

    # Check successful initializations
    success_cases = [
        "agent",
        "function_summarize",
        "function_route",
        "function_validate",
        "function_unknown",
        "tool_with",
        "tool_empty",
        "conditional",
        "subgraph",
        "start",
        "end",
        "unknown",
        "retry",
        "parallel",
        "timeout",
    ]

    for case in success_cases:
        result = init_results[f"{case}_init"]
        assert result["success"], f"Failed to init {case}: {result.get('error')}"
        assert result["execution_count"] == 0
        assert result["last_execution_time"] is None
        assert result["last_error"] is None


@then("basic execution should work for each type")
def step_verify_basic_execution(context):
    """Verify basic execution works for each node type."""
    # Create a shared event loop for efficiency
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        execution_results = {}

        # Test each node type with appropriate state
        test_cases = [
            ("agent", context.message_state),
            ("function_summarize", context.multi_message_state),
            ("function_route", context.question_state),
            ("function_validate", context.message_state),
            ("function_unknown", context.message_state),
            ("tool_with", context.message_state),
            ("tool_empty", context.message_state),
            ("conditional", context.message_state),
            ("subgraph", context.message_state),
            ("start", context.empty_state),
            ("end", context.empty_state),
            ("unknown", context.empty_state),
        ]

        for config_name, state in test_cases:
            config = context.node_configs[config_name]
            node = Node(config, context.agents)

            try:
                result = loop.run_until_complete(node.execute(state))
                execution_results[config_name] = {"success": True, "result": result}

                # Verify basic properties
                assert "current_node" in result
                assert result["current_node"] == config.name
                assert node.execution_count > 0
                assert node.last_execution_time is not None

            except Exception as e:
                execution_results[config_name] = {"success": False, "error": str(e)}

        context.test_results["basic_execution"] = execution_results

        # Verify successful cases
        assert execution_results["agent"]["success"]
        assert execution_results["function_summarize"]["success"]
        assert execution_results["start"]["success"]
        assert execution_results["end"]["success"]

        # Verify start/end specific results
        assert execution_results["start"]["result"].get("started") is True
        assert execution_results["end"]["result"].get("completed") is True

    finally:
        loop.close()


@given("I have agent node test configurations")
def step_agent_configs(context):
    """Prepare agent node configurations."""
    # Already done in various_configs
    pass


@when("I test agent node execution paths")
def step_test_agent_execution(context):
    """Test various agent node execution scenarios."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        results = {}

        # Test valid agent execution
        agent_config = context.node_configs["agent"]
        agent_node = Node(agent_config, context.agents)
        result = loop.run_until_complete(agent_node.execute(context.message_state))
        results["valid_agent"] = result

        # Test empty state
        result_empty = loop.run_until_complete(agent_node.execute(context.empty_state))
        results["empty_state"] = result_empty

        # Test missing agent spec (should return error in result)
        no_spec_config = context.node_configs["agent_no_spec"]
        no_spec_node = Node(no_spec_config, context.agents)
        result_no_spec = loop.run_until_complete(
            no_spec_node.execute(context.message_state)
        )
        results["no_agent_spec"] = result_no_spec

        # Test missing agent (should return error in result)
        missing_config = context.node_configs["agent_missing"]
        missing_node = Node(missing_config, context.agents)
        result_missing = loop.run_until_complete(
            missing_node.execute(context.message_state)
        )
        results["missing_agent"] = result_missing

        context.test_results["agent_execution"] = results

    finally:
        loop.close()


@then("agent execution should handle valid agents correctly")
def step_verify_valid_agents(context):
    """Verify valid agent execution."""
    results = context.test_results["agent_execution"]

    # Valid agent should succeed
    valid = results["valid_agent"]
    assert "current_node" in valid
    assert "messages" in valid
    assert len(valid["messages"]) > 0
    assert valid["messages"][0]["role"] == "assistant"

    # Empty state should still work
    empty = results["empty_state"]
    assert "current_node" in empty
    assert "messages" in empty


@then("error cases should be handled properly")
def step_verify_agent_errors(context):
    """Verify agent error handling."""
    results = context.test_results["agent_execution"]

    # No agent specified should return error
    no_spec = results["no_agent_spec"]
    assert "error" in no_spec
    assert "no agent specified" in no_spec["error"]

    # Missing agent should return error
    missing = results["missing_agent"]
    assert "error" in missing
    assert "not found" in missing["error"]


@given("I have function node test configurations")
def step_function_configs(context):
    """Prepare function node configurations."""
    # Already done in various_configs
    pass


@when("I test function node execution paths")
def step_test_function_execution(context):
    """Test function node execution scenarios."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        results = {}

        # Test summarize function
        sum_node = Node(context.node_configs["function_summarize"], context.agents)
        results["summarize"] = loop.run_until_complete(
            sum_node.execute(context.multi_message_state)
        )

        # Test route function with question
        route_node = Node(context.node_configs["function_route"], context.agents)
        results["route_question"] = loop.run_until_complete(
            route_node.execute(context.question_state)
        )
        results["route_statement"] = loop.run_until_complete(
            route_node.execute(context.message_state)
        )
        results["route_empty"] = loop.run_until_complete(
            route_node.execute(context.empty_state)
        )

        # Test validate function
        val_node = Node(context.node_configs["function_validate"], context.agents)
        results["validate_valid"] = loop.run_until_complete(
            val_node.execute(context.message_state)
        )
        error_state = GraphState(messages=[], error="test error")
        results["validate_invalid"] = loop.run_until_complete(
            val_node.execute(error_state)
        )

        # Test unknown function
        unk_node = Node(context.node_configs["function_unknown"], context.agents)
        results["unknown"] = loop.run_until_complete(
            unk_node.execute(context.message_state)
        )

        # Test no function specified
        none_node = Node(context.node_configs["function_none"], context.agents)
        results["no_function"] = loop.run_until_complete(
            none_node.execute(context.message_state)
        )

        context.test_results["function_execution"] = results

    finally:
        loop.close()


@then("all built-in functions should execute correctly")
def step_verify_functions(context):
    """Verify function execution."""
    results = context.test_results["function_execution"]

    # Summarize should work
    summary = results["summarize"]
    assert "metadata" in summary
    assert "summary" in summary["metadata"]

    # Route should work correctly
    assert results["route_question"]["metadata"]["route"] == "question"
    assert results["route_statement"]["metadata"]["route"] == "statement"
    assert results["route_empty"]["metadata"]["route"] == "default"

    # Validate should work
    assert results["validate_valid"]["metadata"]["valid"] is True
    assert results["validate_invalid"]["metadata"]["valid"] is False


@then("unknown functions should return empty results")
def step_verify_unknown_functions(context):
    """Verify unknown function handling."""
    results = context.test_results["function_execution"]

    # Unknown function should return minimal result
    unknown = results["unknown"]
    assert "current_node" in unknown
    # Should only have current_node, no other keys

    # No function should return error
    no_func = results["no_function"]
    assert "error" in no_func
    assert "no function specified" in no_func["error"]


@given("I have conditional node test configurations")
def step_conditional_configs(context):
    """Prepare conditional node configurations."""
    context.conditional_configs = {
        "always": NodeConfig(
            name="always", type=NodeType.CONDITIONAL, condition={"type": "always"}
        ),
        "never": NodeConfig(
            name="never", type=NodeType.CONDITIONAL, condition={"type": "never"}
        ),
        "has_messages": NodeConfig(
            name="has_msg",
            type=NodeType.CONDITIONAL,
            condition={"type": "has_messages"},
        ),
        "msg_count_gt": NodeConfig(
            name="gt",
            type=NodeType.CONDITIONAL,
            condition={"type": "message_count", "operator": "gt", "value": 2},
        ),
        "msg_count_eq": NodeConfig(
            name="eq",
            type=NodeType.CONDITIONAL,
            condition={"type": "message_count", "operator": "eq", "value": 3},
        ),
        "metadata_check": NodeConfig(
            name="meta",
            type=NodeType.CONDITIONAL,
            condition={"type": "metadata_check", "key": "status", "value": "ready"},
        ),
        "custom": NodeConfig(
            name="custom",
            type=NodeType.CONDITIONAL,
            condition={"type": "custom", "function": lambda s: True},
        ),
        "unknown": NodeConfig(
            name="unknown", type=NodeType.CONDITIONAL, condition={"type": "unknown"}
        ),
        "no_condition": NodeConfig(name="no_cond", type=NodeType.CONDITIONAL),
    }


@when("I test conditional logic execution")
def step_test_conditionals(context):
    """Test conditional node execution."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        results = {}

        for name, config in context.conditional_configs.items():
            node = Node(config, context.agents)

            # Test with appropriate state
            if name in ["has_messages", "msg_count_gt", "msg_count_eq"]:
                state = context.multi_message_state
            elif name == "metadata_check":
                state = context.metadata_state
            else:
                state = context.message_state

            result = loop.run_until_complete(node.execute(state))
            results[name] = result

        context.test_results["conditional_execution"] = results

    finally:
        loop.close()


@then("all condition types should evaluate correctly")
def step_verify_conditionals(context):
    """Verify conditional execution."""
    results = context.test_results["conditional_execution"]

    # Verify expected results
    assert results["always"]["metadata"]["condition_result"] is True
    assert results["never"]["metadata"]["condition_result"] is False
    assert results["has_messages"]["metadata"]["condition_result"] is True
    assert results["msg_count_gt"]["metadata"]["condition_result"] is True
    assert results["msg_count_eq"]["metadata"]["condition_result"] is True
    assert results["metadata_check"]["metadata"]["condition_result"] is True
    assert results["custom"]["metadata"]["condition_result"] is True


@then("conditional edge cases should be handled properly")
def step_verify_conditional_edges(context):
    """Verify conditional edge cases."""
    results = context.test_results["conditional_execution"]

    # Unknown conditions should default to True
    assert results["unknown"]["metadata"]["condition_result"] is True
    assert results["no_condition"]["metadata"]["condition_result"] is True


@given("I have node configurations for utility testing")
def step_utility_configs(context):
    """Prepare configurations for utility method testing."""
    # Already prepared in various_configs
    pass


@when("I test node utility methods")
def step_test_utilities(context):
    """Test node utility methods."""
    results = {}

    # Test timeout
    timeout_node = Node(context.node_configs["timeout"], context.agents)
    results["timeout"] = timeout_node.get_timeout()

    # Test parallel execution
    parallel_node = Node(context.node_configs["parallel"], context.agents)
    results["parallel"] = parallel_node.can_execute_parallel()

    non_parallel_config = NodeConfig(
        name="non-parallel", type=NodeType.FUNCTION, function="test", parallel=False
    )
    non_parallel_node = Node(non_parallel_config, context.agents)
    results["non_parallel"] = non_parallel_node.can_execute_parallel()

    # Test edge operations
    test_node = Node(context.node_configs["agent"], context.agents)
    edges = [
        Edge(source="agent-node", target="target1"),
        Edge(source="agent-node", target="target2"),
        Edge(source="other", target="target3"),
    ]
    results["outgoing_edges"] = test_node.get_edges(edges)

    # Test edge condition evaluation with proper event loop management
    # Create a new event loop for this test since we need asyncio for edge evaluation
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        edge_no_condition = Edge(source="test", target="test2")
        results["edge_no_condition"] = test_node.evaluate_edge_condition(
            edge_no_condition, context.message_state
        )

        edge_with_condition = Edge(
            source="test", target="test2", condition={"type": "always"}
        )
        results["edge_with_condition"] = test_node.evaluate_edge_condition(
            edge_with_condition, context.message_state
        )
    finally:
        loop.close()

    context.test_results["utilities"] = results


@then("timeout and parallel settings should work")
def step_verify_settings(context):
    """Verify timeout and parallel settings."""
    results = context.test_results["utilities"]

    assert results["timeout"] == 5.0
    assert results["parallel"] is True
    assert results["non_parallel"] is False


@then("edge operations should function correctly")
def step_verify_edges(context):
    """Verify edge operations."""
    results = context.test_results["utilities"]

    # Should get 2 outgoing edges
    edges = results["outgoing_edges"]
    assert len(edges) == 2
    assert all(e.source == "agent-node" for e in edges)

    # Edge conditions should evaluate
    assert results["edge_no_condition"] is True
    assert results["edge_with_condition"] is True


@given("I have error and retry test configurations")
def step_error_configs(context):
    """Prepare error and retry configurations."""
    # Use existing retry config but ensure no delays
    pass


@when("I test error handling scenarios")
def step_test_error_handling(context):
    """Test error handling and retry scenarios."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    try:
        results = {}

        # Test subgraph without specification (should return error)
        subgraph_none_node = Node(context.node_configs["subgraph_none"], context.agents)
        result = loop.run_until_complete(
            subgraph_none_node.execute(context.message_state)
        )
        results["subgraph_error"] = result

        # Test retry mechanism with mocked failure (no actual delays)
        retry_config = context.node_configs["retry"]
        retry_node = Node(retry_config, context.agents)

        # Mock the function to fail then succeed without delays
        call_count = [0]
        original_func = retry_node._execute_function

        async def mock_failing_func(state):
            call_count[0] += 1
            if call_count[0] == 1:
                raise Exception("First attempt fails")
            return {"success": True, "attempt": call_count[0]}

        retry_node._execute_function = mock_failing_func
        result = loop.run_until_complete(retry_node.execute(context.message_state))
        results["retry_success"] = result
        results["retry_attempts"] = call_count[0]

        # Test retry failure (all attempts fail)
        fail_config = NodeConfig(
            name="fail",
            type=NodeType.FUNCTION,
            function="test",
            retry_policy={"max_retries": 1, "delay": 0},
        )
        fail_node = Node(fail_config, context.agents)

        async def mock_always_fail(state):
            raise Exception("Always fails")

        fail_node._execute_function = mock_always_fail
        result = loop.run_until_complete(fail_node.execute(context.message_state))
        results["retry_failure"] = result

        context.test_results["error_handling"] = results

    finally:
        loop.close()


@then("errors should be captured and returned properly")
def step_verify_error_capture(context):
    """Verify error capture."""
    results = context.test_results["error_handling"]

    # Subgraph error should be captured
    subgraph_error = results["subgraph_error"]
    assert "error" in subgraph_error
    assert "no subgraph specified" in subgraph_error["error"]

    # Retry failure should be captured
    retry_failure = results["retry_failure"]
    assert "error" in retry_failure
    assert "Always fails" in retry_failure["error"]
    assert "failed_node" in retry_failure


@then("retry policies should work without delays")
def step_verify_retry(context):
    """Verify retry policies work efficiently."""
    results = context.test_results["error_handling"]

    # Retry should have succeeded after failure
    retry_success = results["retry_success"]
    assert "success" in retry_success
    assert retry_success["success"] is True
    assert results["retry_attempts"] == 2  # Failed once, succeeded on retry