from pydantic import BaseModel

from agentbyte.workflow.core.models import (
    EdgeCondition,
    JoinFailureMode,
    StepExecution,
    StepStatus,
    WorkflowCompletionMode,
    WorkflowStateMode,
    WorkflowExecution,
)
from agentbyte.workflow.core.workflow import BaseWorkflow, WorkflowConfig
from agentbyte.workflow.steps.echo import EchoOutput, EchoStep
from agentbyte.workflow.steps.function import FunctionStep
from agentbyte.workflow.steps.step import StepMetadata


def test_basic_construction_and_hash():
    wf = BaseWorkflow(WorkflowConfig(name="demo"))
    s1 = EchoStep(step_id="start", metadata=StepMetadata(name="start"))
    s2 = EchoStep(
        step_id="end",
        metadata=StepMetadata(name="end"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.chain(s1, s2)

    validation = wf.validate_workflow()
    assert validation.is_valid
    assert not validation.errors
    assert wf.compute_structure_hash()


def test_detect_cycle():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(step_id="b", metadata=StepMetadata(name="b"))
    wf.add_edge(a, b).add_edge(b, a).set_start_step(a)

    result = wf.validate_workflow()
    assert not result.is_valid
    assert result.has_cycles
    assert "Workflow contains cycles" in result.errors


def test_detect_cycle_in_disconnected_component():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(step_id="b", metadata=StepMetadata(name="b"))
    c = EchoStep(step_id="c", metadata=StepMetadata(name="c"))
    d = EchoStep(step_id="d", metadata=StepMetadata(name="d"))

    wf.add_edge(a, b).set_start_step(a)
    wf.add_edge(c, d).add_edge(d, c)

    result = wf.validate_workflow()
    assert not result.is_valid
    assert result.has_cycles
    assert "Workflow contains cycles" in result.errors


def test_unreachable_warning():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    c = EchoStep(step_id="c", metadata=StepMetadata(name="c"))
    wf.add_edge(a, b).set_start_step(a).add_step(c)

    result = wf.validate_workflow()
    assert result.is_valid
    assert "Unreachable steps: c" in result.warnings
    assert result.unreachable_steps == ["c"]


def test_ready_steps_always_edges():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_edge(a, b).set_start_step(a)

    execution = WorkflowExecution(workflow_id="wf", state={}, step_executions={})
    ready = wf.get_ready_steps(execution)
    assert [s.step_id for s in ready] == ["a"]

    execution.step_executions["a"] = StepExecution(
        step_id="a",
        status=StepStatus.COMPLETED,
        output_data={"result": "hi"},
    )
    ready = wf.get_ready_steps(execution)
    assert [s.step_id for s in ready] == ["b"]


def test_ready_steps_conditional_output_based():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    ).set_start_step(a)

    execution = WorkflowExecution(workflow_id="wf", state={}, step_executions={})
    ready = wf.get_ready_steps(execution)
    assert [s.step_id for s in ready] == ["a"]

    execution.step_executions["a"] = StepExecution(
        step_id="a",
        status=StepStatus.COMPLETED,
        output_data={"result": "stop"},
    )
    ready = wf.get_ready_steps(execution)
    assert ready == []

    execution.step_executions["a"].output_data = {"result": "go"}
    ready = wf.get_ready_steps(execution)
    assert [s.step_id for s in ready] == ["b"]


def test_type_mismatch_validation():
    class Out(BaseModel):
        x: int

    class In(BaseModel):
        x: str

    producer = FunctionStep(
        step_id="prod",
        metadata=StepMetadata(name="prod", description="prod"),
        func=lambda _ctx, _state, _input: Out(x=1),
        input_type=In,
        output_type=Out,
    )
    consumer = FunctionStep(
        step_id="cons",
        metadata=StepMetadata(name="cons", description="cons"),
        func=lambda _ctx, _state, _input: _input,
        input_type=In,
        output_type=EchoOutput,
    )

    wf = BaseWorkflow()
    wf.add_edge(producer, consumer).set_start_step(producer)
    result = wf.validate_workflow()
    assert not result.is_valid
    assert "Type mismatch" in result.errors[0]


def test_conditional_conflict_warning():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    c = EchoStep(
        step_id="c",
        metadata=StepMetadata(name="c"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_step(a).add_step(b).add_step(c)
    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="echo", operator="==", value="x"),
    )
    wf.add_edge(
        a,
        c,
        EdgeCondition(type="output_based", field="echo", operator="==", value="y"),
    )
    wf.set_start_step(a)
    result = wf.validate_workflow()
    assert result.is_valid
    assert any("Multiple conditional branches" in w for w in result.warnings)


def test_invalid_condition_type_validation_error():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(a, b, EdgeCondition(type="custom")).set_start_step(a)
    result = wf.validate_workflow()

    assert not result.is_valid
    assert any("unknown condition type" in err for err in result.errors)


def test_output_based_condition_requires_field_and_operator():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(a, b, EdgeCondition(type="output_based", value="go")).set_start_step(a)
    result = wf.validate_workflow()

    assert not result.is_valid
    assert any("must define field" in err for err in result.errors)
    assert any("must define operator" in err for err in result.errors)


def test_mixed_incoming_edges_require_unconditional_dependencies():
    wf = BaseWorkflow()
    gate = EchoStep(step_id="gate", metadata=StepMetadata(name="gate"))
    branch = EchoStep(step_id="branch", metadata=StepMetadata(name="branch"))
    target = EchoStep(
        step_id="target",
        metadata=StepMetadata(name="target"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(gate, target)
    wf.add_edge(
        branch,
        target,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.set_start_step(gate)
    wf.add_step(branch)

    execution = WorkflowExecution(workflow_id="wf", state={}, step_executions={})
    execution.step_executions["branch"] = StepExecution(
        step_id="branch",
        status=StepStatus.COMPLETED,
        output_data={"result": "go"},
    )

    ready = wf.get_ready_steps(execution)
    assert all(step.step_id != "target" for step in ready)

    execution.step_executions["gate"] = StepExecution(
        step_id="gate",
        status=StepStatus.COMPLETED,
        output_data={"result": "ok"},
    )
    ready = wf.get_ready_steps(execution)
    assert [step.step_id for step in ready] == ["target"]


def test_invalid_comparison_operator_evaluation_is_safe():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator=">", value=10),
    ).set_start_step(a)

    execution = WorkflowExecution(workflow_id="wf", state={}, step_executions={})
    execution.step_executions["a"] = StepExecution(
        step_id="a",
        status=StepStatus.COMPLETED,
        output_data={"result": "not-a-number"},
    )
    ready = wf.get_ready_steps(execution)
    assert ready == []


def test_warning_for_step_without_outgoing_not_marked_end():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_edge(a, b).set_start_step(a)

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("has no outgoing edges but is not marked as an end step" in w for w in result.warnings)


def test_warning_for_conditional_only_fanin_or_semantics():
    wf = BaseWorkflow()
    start = EchoStep(step_id="start", metadata=StepMetadata(name="start"))
    left = EchoStep(
        step_id="left",
        metadata=StepMetadata(name="left"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    right = EchoStep(
        step_id="right",
        metadata=StepMetadata(name="right"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    target = EchoStep(
        step_id="target",
        metadata=StepMetadata(name="target"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(start, left).add_edge(start, right).set_start_step(start)
    wf.add_edge(
        left,
        target,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.add_edge(
        right,
        target,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("has only conditional fan-in edges; readiness uses OR semantics" in w for w in result.warnings)


def test_warning_for_duplicate_conditional_branch_signature():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.set_start_step(a)

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("Duplicate conditional branch" in w for w in result.warnings)


def test_warning_for_mixed_equal_not_equal_overlap():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    c = EchoStep(
        step_id="c",
        metadata=StepMetadata(name="c"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.add_edge(
        a,
        c,
        EdgeCondition(type="output_based", field="result", operator="!=", value="go"),
    )
    wf.set_start_step(a)

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("may overlap between '==' and '!=' operators" in w for w in result.warnings)


def test_warning_for_multiple_not_equal_branch_overlap():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    c = EchoStep(
        step_id="c",
        metadata=StepMetadata(name="c"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(
        a,
        b,
        EdgeCondition(type="output_based", field="result", operator="!=", value="x"),
    )
    wf.add_edge(
        a,
        c,
        EdgeCondition(type="output_based", field="result", operator="!=", value="y"),
    )
    wf.set_start_step(a)

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("Multiple '!=' branches" in w for w in result.warnings)


def test_warning_for_mixed_always_and_conditional_fanout():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    c = EchoStep(
        step_id="c",
        metadata=StepMetadata(name="c"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )

    wf.add_edge(a, b)
    wf.add_edge(
        a,
        c,
        EdgeCondition(type="output_based", field="result", operator="==", value="go"),
    )
    wf.set_start_step(a)

    result = wf.validate_workflow()
    assert result.is_valid
    assert any("mixes always and conditional fan-out edges" in w for w in result.warnings)


def test_structure_hash_changes_with_topology():
    wf = BaseWorkflow()
    a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
    b = EchoStep(
        step_id="b",
        metadata=StepMetadata(name="b"),
        input_type=EchoOutput,
        output_type=EchoOutput,
    )
    wf.add_edge(a, b).set_start_step(a)
    h1 = wf.compute_structure_hash()

    wf.add_end_step(b)
    h2 = wf.compute_structure_hash()
    assert h1 != h2


def test_structure_hash_changes_with_completion_mode():
    """Changing WorkflowCompletionMode must produce a different hash."""
    wf_any = BaseWorkflow(WorkflowConfig(name="demo", completion_mode=WorkflowCompletionMode.ANY_END))
    wf_all = BaseWorkflow(WorkflowConfig(name="demo", completion_mode=WorkflowCompletionMode.ALL_END))
    for wf in (wf_any, wf_all):
        a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
        b = EchoStep(step_id="b", metadata=StepMetadata(name="b"), input_type=EchoOutput, output_type=EchoOutput)
        wf.add_edge(a, b).set_start_step(a).add_end_step(b)

    assert wf_any.compute_structure_hash() != wf_all.compute_structure_hash()


def test_structure_hash_changes_with_state_mode() -> None:
    wf_immediate = BaseWorkflow(
        WorkflowConfig(name="demo", state_mode=WorkflowStateMode.IMMEDIATE)
    )
    wf_staged = BaseWorkflow(
        WorkflowConfig(name="demo", state_mode=WorkflowStateMode.STAGED)
    )
    for wf in (wf_immediate, wf_staged):
        a = EchoStep(step_id="a", metadata=StepMetadata(name="a"))
        b = EchoStep(
            step_id="b",
            metadata=StepMetadata(name="b"),
            input_type=EchoOutput,
            output_type=EchoOutput,
        )
        wf.add_edge(a, b).set_start_step(a).add_end_step(b)

    assert wf_immediate.compute_structure_hash() != wf_staged.compute_structure_hash()


def test_configure_join_sets_failure_mode_on_step_metadata():
    """configure_join() must update the join step's join_failure_mode."""
    wf = BaseWorkflow()
    branch_a = EchoStep(step_id="branch_a", metadata=StepMetadata(name="branch_a"))
    branch_b = EchoStep(step_id="branch_b", metadata=StepMetadata(name="branch_b"))
    join = EchoStep(step_id="join", metadata=StepMetadata(name="join"))
    wf.add_step(branch_a).add_step(branch_b).add_step(join)
    wf.set_start_step(branch_a)
    wf.add_edge(branch_a, join)
    wf.add_edge(branch_b, join)

    assert join.metadata.join_failure_mode == JoinFailureMode.FAIL_FAST
    assert not wf.is_step_failure_tolerated("join")

    wf.configure_join("join", failure_mode=JoinFailureMode.SKIP_FAILED)

    assert join.metadata.join_failure_mode == JoinFailureMode.SKIP_FAILED
    assert wf.is_step_failure_tolerated("join")


def test_fan_in_join_step_skips_type_check_in_validation():
    """Fan-in join steps with >=2 always-edge inputs must not raise type-mismatch errors."""
    class BranchOutput(BaseModel):
        value: int

    class JoinInput(BaseModel):
        branch_outputs: dict
        failed_branches: list

    wf = BaseWorkflow(WorkflowConfig(name="fanin-typecheck"))
    branch_a = FunctionStep(
        step_id="branch_a",
        metadata=StepMetadata(name="branch_a"),
        input_type=EchoOutput,
        output_type=BranchOutput,
        func=lambda i, c: BranchOutput(value=1),  # type: ignore[arg-type]
    )
    branch_b = FunctionStep(
        step_id="branch_b",
        metadata=StepMetadata(name="branch_b"),
        input_type=EchoOutput,
        output_type=BranchOutput,
        func=lambda i, c: BranchOutput(value=2),  # type: ignore[arg-type]
    )
    join = FunctionStep(
        step_id="join",
        metadata=StepMetadata(name="join"),
        input_type=JoinInput,
        output_type=EchoOutput,
        func=lambda i, c: EchoOutput(result="done"),  # type: ignore[arg-type]
    )
    start = EchoStep(step_id="start", metadata=StepMetadata(name="start"))
    wf.add_edge(start, branch_a).add_edge(start, branch_b)
    wf.add_edge(branch_a, join).add_edge(branch_b, join)
    wf.set_start_step(start).add_end_step(join)

    result = wf.validate_workflow()
    # No type-mismatch errors for the fan-in edges
    assert all("Type mismatch" not in err for err in result.errors)
