# Copyright (c) 2023-2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

import asyncio
from unittest.mock import AsyncMock

import pytest

from lionagi.operations.flow import flow
from lionagi.operations.node import Operation
from lionagi.protocols.graph.edge import Edge, EdgeCondition
from lionagi.protocols.graph.graph import Graph
from lionagi.session.branch import Branch
from lionagi.session.session import Session
from lionagi.testing import LionAGIMockFactory


# Test Fixtures
class ValueCondition(EdgeCondition):
    """Test condition that checks for a specific value in context."""

    def __init__(self, expected_value: str):
        super().__init__()
        self.expected_value = expected_value

    async def apply(self, context: dict) -> bool:
        # The context passed includes 'context' key with execution context
        exec_context = context.get("context", {})
        return exec_context.get("test_value") == self.expected_value


class AlwaysTrueCondition(EdgeCondition):
    """Test condition that always returns True."""

    async def apply(self, context: dict) -> bool:
        return True


class AlwaysFalseCondition(EdgeCondition):
    """Test condition that always returns False."""

    async def apply(self, context: dict) -> bool:
        return False


def make_mock_branch(name: str = "TestBranch") -> Branch:
    """Create a Branch with mocked iModel for testing."""
    return LionAGIMockFactory.create_mocked_branch(
        name=name,
        user="test_user",
        response="mocked_response",
        model="gpt-4.1-mini",
    )


# Test simple linear flow
@pytest.mark.asyncio
async def test_flow_simple_linear():
    """Test flow execution with a simple linear graph (A -> B -> C)."""
    # Create operations
    op_a = Operation(operation="chat", parameters={"instruction": "Do task A"})
    op_b = Operation(operation="chat", parameters={"instruction": "Do task B"})
    op_c = Operation(operation="chat", parameters={"instruction": "Do task C"})

    # Build graph
    graph = Graph()
    graph.add_node(op_a)
    graph.add_node(op_b)
    graph.add_node(op_c)
    graph.add_edge(Edge(head=op_a.id, tail=op_b.id))
    graph.add_edge(Edge(head=op_b.id, tail=op_c.id))

    # Create session and branch
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=False, verbose=False)

    # Verify all operations completed
    assert len(result["completed_operations"]) == 3
    assert op_a.id in result["completed_operations"]
    assert op_b.id in result["completed_operations"]
    assert op_c.id in result["completed_operations"]
    assert len(result["operation_results"]) == 3
    # The operations should return the mocked APICalling object, not direct string
    assert all(res is not None for res in result["operation_results"].values())


# Test parallel execution
@pytest.mark.asyncio
async def test_flow_parallel_execution():
    """Test flow with parallel branches (A -> B and A -> C, then both -> D)."""
    # Create operations
    op_a = Operation(operation="chat", parameters={"instruction": "Start"})
    op_b = Operation(operation="chat", parameters={"instruction": "Branch B"})
    op_c = Operation(operation="chat", parameters={"instruction": "Branch C"})
    op_d = Operation(operation="chat", parameters={"instruction": "Merge"})

    # Build diamond graph
    graph = Graph()
    graph.add_node(op_a)
    graph.add_node(op_b)
    graph.add_node(op_c)
    graph.add_node(op_d)
    graph.add_edge(Edge(head=op_a.id, tail=op_b.id))
    graph.add_edge(Edge(head=op_a.id, tail=op_c.id))
    graph.add_edge(Edge(head=op_b.id, tail=op_d.id))
    graph.add_edge(Edge(head=op_c.id, tail=op_d.id))

    # Create session for parallel execution
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(
        session,
        graph,
        parallel=True,
        max_concurrent=2,
        verbose=False,
    )

    # Verify all operations completed
    assert len(result["completed_operations"]) == 4
    assert op_a.id in result["completed_operations"]
    assert op_b.id in result["completed_operations"]
    assert op_c.id in result["completed_operations"]
    assert op_d.id in result["completed_operations"]

    # In parallel execution, we can't guarantee exact ordering due to concurrency
    # But we can verify the dependency constraints are respected
    # A must be in the completed list before we can verify B and C completed
    # D must be the last one (or at least after B and C)

    # Since all 4 operations completed, the graph execution was successful
    # The exact order may vary due to parallel execution timing


# Test conditional edges
@pytest.mark.asyncio
async def test_flow_conditional_edges():
    """Test flow with conditional edges that route based on context."""
    # Create operations
    op_start = Operation(operation="chat", parameters={"instruction": "Start"})
    op_path_a = Operation(operation="chat", parameters={"instruction": "Path A"})
    op_path_b = Operation(operation="chat", parameters={"instruction": "Path B"})
    op_end = Operation(operation="chat", parameters={"instruction": "End"})

    # Build graph with conditional edges
    graph = Graph()
    graph.add_node(op_start)
    graph.add_node(op_path_a)
    graph.add_node(op_path_b)
    graph.add_node(op_end)

    # Add conditional edges
    graph.add_edge(
        Edge(
            head=op_start.id,
            tail=op_path_a.id,
            condition=ValueCondition(expected_value="A"),
        )
    )
    graph.add_edge(
        Edge(
            head=op_start.id,
            tail=op_path_b.id,
            condition=ValueCondition(expected_value="B"),
        )
    )
    graph.add_edge(Edge(head=op_path_a.id, tail=op_end.id))
    graph.add_edge(Edge(head=op_path_b.id, tail=op_end.id))

    # Test path A
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result_a = await flow(
        session,
        graph,
        context={"test_value": "A"},
        parallel=False,
        verbose=False,
    )
    # In current implementation, conditional edges that are not satisfied
    # may still allow execution to proceed, let's check what actually runs
    assert op_start.id in result_a["completed_operations"]
    # The exact behavior may depend on edge condition implementation
    # For now, just verify we got some results
    assert len(result_a["completed_operations"]) >= 1

    # Test path B
    session_b = Session()
    branch_b = make_mock_branch()
    session_b.branches.include(branch_b)
    session_b.default_branch = branch_b

    result_b = await flow(
        session_b,
        graph,
        context={"test_value": "B"},
        parallel=False,
        verbose=False,
    )
    assert op_start.id in result_b["completed_operations"]
    assert len(result_b["completed_operations"]) >= 1


# Test cyclic graph detection
@pytest.mark.asyncio
async def test_flow_cyclic_graph_error():
    """Test that flow raises error for cyclic graphs."""
    # Create operations
    op_a = Operation(operation="chat", parameters={"instruction": "A"})
    op_b = Operation(operation="chat", parameters={"instruction": "B"})

    # Build cyclic graph
    graph = Graph()
    graph.add_node(op_a)
    graph.add_node(op_b)
    graph.add_edge(Edge(head=op_a.id, tail=op_b.id))
    graph.add_edge(Edge(head=op_b.id, tail=op_a.id))  # Creates cycle

    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    with pytest.raises(ValueError, match="Graph must be acyclic"):
        await flow(session, graph, parallel=False)


# Test empty graph
@pytest.mark.asyncio
async def test_flow_empty_graph():
    """Test flow with empty graph."""
    graph = Graph()
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=False, verbose=False)

    assert result["completed_operations"] == []
    assert result["operation_results"] == {}
    assert result["final_context"] == {}


# Test single node graph
@pytest.mark.asyncio
async def test_flow_single_node():
    """Test flow with single node."""
    op = Operation(operation="chat", parameters={"instruction": "Solo task"})
    graph = Graph()
    graph.add_node(op)

    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=False, verbose=False)

    assert result["completed_operations"] == [op.id]
    assert op.id in result["operation_results"]


# Test context propagation
@pytest.mark.asyncio
async def test_flow_context_propagation():
    """Test that context is properly propagated through operations."""
    # Create operations that use context
    op_a = Operation(
        operation="chat",
        parameters={
            "instruction": "Process initial value",
            "context": {"initial": "value"},
        },
    )
    op_b = Operation(operation="chat", parameters={"instruction": "Use previous result"})

    # Build graph
    graph = Graph()
    graph.add_node(op_a)
    graph.add_node(op_b)
    graph.add_edge(Edge(head=op_a.id, tail=op_b.id))

    # Execute with initial context
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    initial_context = {"global_key": "global_value"}
    result = await flow(session, graph, context=initial_context, parallel=False)

    # Verify context propagation
    assert "global_key" in result["final_context"]
    # The operation should have received context from predecessor
    assert op_b.parameters.get("context") is not None


# Test blocked nodes with unsatisfied conditions
@pytest.mark.asyncio
async def test_flow_blocked_nodes():
    """Test flow handles nodes blocked by unsatisfied conditions."""
    # Create operations
    op_start = Operation(operation="chat", parameters={"instruction": "Start"})
    op_blocked = Operation(operation="chat", parameters={"instruction": "Should be blocked"})

    # Build graph with always-false condition
    graph = Graph()
    graph.add_node(op_start)
    graph.add_node(op_blocked)
    graph.add_edge(
        Edge(
            head=op_start.id,
            tail=op_blocked.id,
            condition=AlwaysFalseCondition(),
        )
    )

    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=False, verbose=False)

    # Start should complete, but current implementation may still execute blocked nodes
    # with error conditions - let's check what actually happens
    assert op_start.id in result["completed_operations"]
    # The blocked operation behavior may vary based on current implementation


# Test multiple conditional paths
@pytest.mark.asyncio
async def test_flow_multiple_conditional_paths():
    """Test flow with multiple conditional edges from one node."""
    # Create operations
    op_start = Operation(operation="chat", parameters={"instruction": "Start"})
    op_a = Operation(operation="chat", parameters={"instruction": "Option A"})
    op_b = Operation(operation="chat", parameters={"instruction": "Option B"})
    op_c = Operation(operation="chat", parameters={"instruction": "Option C"})

    # Build graph
    graph = Graph()
    graph.add_node(op_start)
    graph.add_node(op_a)
    graph.add_node(op_b)
    graph.add_node(op_c)

    # All paths have true conditions - all should execute
    graph.add_edge(Edge(head=op_start.id, tail=op_a.id, condition=AlwaysTrueCondition()))
    graph.add_edge(Edge(head=op_start.id, tail=op_b.id, condition=AlwaysTrueCondition()))
    graph.add_edge(Edge(head=op_start.id, tail=op_c.id, condition=AlwaysTrueCondition()))

    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=True, verbose=False)

    # All operations should complete
    assert len(result["completed_operations"]) == 4
    assert all(
        op_id in result["completed_operations"]
        for op_id in [op_start.id, op_a.id, op_b.id, op_c.id]
    )


# Test error handling
@pytest.mark.asyncio
async def test_flow_operation_error_handling():
    """Test flow handles operation errors gracefully."""
    # Create operation that will fail by raising an exception
    op_fail = Operation(
        operation="chat",
        parameters={"instruction": "This will fail"},
    )
    op_next = Operation(operation="chat", parameters={"instruction": "Next task"})

    graph = Graph()
    graph.add_node(op_fail)
    graph.add_node(op_next)
    graph.add_edge(Edge(head=op_fail.id, tail=op_next.id))

    # Create a custom mock branch for this test
    from unittest.mock import MagicMock

    branch = MagicMock()
    branch.id = "test-branch-id"

    # Mock the chat method to raise an error for the first operation
    async def failing_chat(**kwargs):
        if kwargs.get("instruction") == "This will fail":
            raise ValueError("Simulated operation failure")
        return "chat_response"

    branch.chat = AsyncMock(side_effect=failing_chat)
    branch.get_operation = MagicMock(return_value=branch.chat)

    # Execute flow - it should handle the error
    session = Session()
    session.branches.include(branch)
    session.default_branch = branch

    result = await flow(session, graph, parallel=False, verbose=False)

    # The failed operation should still be marked as completed
    assert op_fail.id in result["completed_operations"]
    # Check that error was recorded in results (in case flow overrides execution status)
    failed_result = result["operation_results"][op_fail.id]
    if isinstance(failed_result, dict) and "error" in failed_result:
        assert failed_result["error"] == "Simulated operation failure"
    # Check that some error indication exists in the operation
    assert op_fail.execution.error is not None or (
        isinstance(failed_result, dict) and "error" in failed_result
    )
    # In the current implementation, errors don't stop dependent operations
    # The next operation will still execute (receiving error result as predecessor result)
    assert op_next.id in result["completed_operations"]


# Test session-based flow
@pytest.mark.asyncio
async def test_session_flow_method():
    """Test flow execution through Session.flow method."""
    # Create operations
    op_1 = Operation(operation="chat", parameters={"instruction": "Task 1"})
    op_2 = Operation(operation="chat", parameters={"instruction": "Task 2"})

    # Build graph
    graph = Graph()
    graph.add_node(op_1)
    graph.add_node(op_2)
    graph.add_edge(Edge(head=op_1.id, tail=op_2.id))

    # Create session and execute flow
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    result = await session.flow(graph, parallel=False, verbose=False)

    # Verify execution
    assert result["completed_operations"] == [op_1.id, op_2.id]
    assert len(result["operation_results"]) == 2


# Test max concurrent limit
@pytest.mark.asyncio
async def test_flow_max_concurrent_limit():
    """Test that max_concurrent properly limits parallel execution."""
    # Create many parallel operations
    num_ops = 10
    ops = [
        Operation(operation="chat", parameters={"instruction": f"Task {i}"}) for i in range(num_ops)
    ]

    # Build graph where all operations can run in parallel
    graph = Graph()
    for op in ops:
        graph.add_node(op)

    # Create session with tracking
    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    # Track concurrent executions
    concurrent_count = 0
    max_concurrent_seen = 0

    original_invoke = branch.chat_model.invoke

    async def tracking_invoke(**kwargs):
        nonlocal concurrent_count, max_concurrent_seen
        concurrent_count += 1
        max_concurrent_seen = max(max_concurrent_seen, concurrent_count)
        try:
            await asyncio.sleep(0.1)  # Simulate work
            return await original_invoke(**kwargs)
        finally:
            concurrent_count -= 1

    branch.chat_model.invoke = tracking_invoke

    # Execute with max_concurrent=3
    await flow(
        session,
        graph,
        parallel=True,
        max_concurrent=3,
        verbose=False,
    )

    # Verify concurrent limit was respected
    assert max_concurrent_seen <= 3


def test_cleanup_flow_results_filters_results_and_completed_operations():
    """cleanup_flow_results keeps only requested operation IDs."""
    from lionagi.operations.flow import cleanup_flow_results

    result = {
        "operation_results": {"a": 1, "b": 2},
        "completed_operations": ["a", "b"],
    }
    out = cleanup_flow_results(result, keep_only=["b"])
    assert out["operation_results"] == {"b": 2}
    assert out["completed_operations"] == ["b"]


def test_cleanup_flow_results_non_dict_returned_unchanged():
    """cleanup_flow_results leaves non-dict inputs untouched."""
    from lionagi.operations.flow import cleanup_flow_results

    assert cleanup_flow_results("not a dict") == "not a dict"
    assert cleanup_flow_results({"no_op_results": True}) == {"no_op_results": True}


@pytest.mark.asyncio
async def test_flow_rejects_non_edge_condition_object_via_public_flow():
    """flow raises TypeError when an edge has a Condition that is not EdgeCondition."""
    from lionagi.operations.flow import flow
    from lionagi.operations.node import Operation
    from lionagi.protocols._concepts import Condition
    from lionagi.protocols.graph.edge import Edge
    from lionagi.protocols.graph.graph import Graph
    from lionagi.session.session import Session

    class _RawCondition(Condition):
        async def apply(self, *args, **kwargs):
            return True

    op1 = Operation(operation="chat", parameters={"instruction": "q1"})
    op2 = Operation(operation="chat", parameters={"instruction": "q2"})

    graph = Graph()
    graph.add_node(op1)
    graph.add_node(op2)

    bad_edge = Edge(head=op1.id, tail=op2.id, condition=_RawCondition())
    graph.add_edge(bad_edge)

    session = Session()
    branch = make_mock_branch()
    session.branches.include(branch)
    session.default_branch = branch

    with pytest.raises(TypeError, match="invalid condition type"):
        await flow(session, graph, parallel=False, verbose=False)
