"""Tests for workflow steps."""
import pytest
from logos_adk.event_bus import InMemoryEventBus
from logos_adk.workflow.step import Step
from logos_adk.workflow.context import PipelineContext, StepResult


class TestStepABC:
    def test_cannot_instantiate(self):
        with pytest.raises(TypeError):
            Step("test")

    def test_subclass_works(self):
        class MyStep(Step):
            async def execute(self, ctx: PipelineContext) -> StepResult:
                return StepResult(output="ok", status="completed", duration_ms=0)

        s = MyStep("my_step")
        assert s.name == "my_step"


from logos_adk.workflow.steps import AgentStep, FunctionStep, PublishEventStep, WaitForEventStep


class TestAgentStep:
    @pytest.mark.asyncio
    async def test_runs_agent(self):
        from logos_adk.agent import Agent
        from logos_adk.providers.mock import MockProvider

        agent = Agent("test", model=MockProvider(responses=["agent says hello"]))
        step = AgentStep("greet", agent)
        ctx = PipelineContext(pipeline_id="p1")
        ctx["_last"] = "Say hello"
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert "hello" in result.output.lower()

    @pytest.mark.asyncio
    async def test_uses_last_as_prompt(self):
        from logos_adk.agent import Agent
        from logos_adk.providers.mock import MockProvider

        agent = Agent("echo", model=MockProvider(responses=["echoed"]))
        step = AgentStep("echo", agent)
        ctx = PipelineContext(pipeline_id="p1")
        ctx["_last"] = "test prompt"
        result = await step.execute(ctx)
        assert result.status == "completed"

    @pytest.mark.asyncio
    async def test_empty_last_uses_empty_string(self):
        from logos_adk.agent import Agent
        from logos_adk.providers.mock import MockProvider

        agent = Agent("a", model=MockProvider(responses=["ok"]))
        step = AgentStep("a", agent)
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"


class TestFunctionStep:
    @pytest.mark.asyncio
    async def test_runs_function(self):
        async def transform(ctx):
            return ctx["_last"].upper()

        step = FunctionStep("upper", transform)
        ctx = PipelineContext(pipeline_id="p1")
        ctx["_last"] = "hello"
        result = await step.execute(ctx)
        assert result.output == "HELLO"
        assert result.status == "completed"

    @pytest.mark.asyncio
    async def test_function_error(self):
        async def fail(ctx):
            raise ValueError("boom")

        step = FunctionStep("fail", fail)
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "failed"
        assert "boom" in result.metadata.get("error", "")


from logos_adk.workflow.steps import ParallelStep, ConditionalStep


class TestParallelStep:
    @pytest.mark.asyncio
    async def test_runs_parallel(self):
        async def fn_a(ctx):
            return "A"

        async def fn_b(ctx):
            return "B"

        step = ParallelStep("par", [
            FunctionStep("a", fn_a),
            FunctionStep("b", fn_b),
        ])
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert result.output["a"] == "A"
        assert result.output["b"] == "B"

    @pytest.mark.asyncio
    async def test_empty_parallel(self):
        step = ParallelStep("empty", [])
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert result.output == {}

    @pytest.mark.asyncio
    async def test_parallel_one_fails(self):
        async def ok(ctx):
            return "ok"

        async def fail(ctx):
            raise ValueError("boom")

        step = ParallelStep("par", [
            FunctionStep("ok", ok),
            FunctionStep("fail", fail),
        ])
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert result.output["ok"] == "ok"
        assert result.output["fail"] is None

    @pytest.mark.asyncio
    async def test_parallel_merges_non_conflicting_blackboard_writes(self):
        async def write_a(ctx):
            ctx.blackboard["a"] = "A"
            return "A"

        async def write_b(ctx):
            ctx.blackboard["b"] = "B"
            return "B"

        step = ParallelStep("par", [
            FunctionStep("a", write_a),
            FunctionStep("b", write_b),
        ])
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert ctx.blackboard["a"] == "A"
        assert ctx.blackboard["b"] == "B"

    @pytest.mark.asyncio
    async def test_parallel_conflicting_blackboard_writes_fail(self):
        async def write_first(ctx):
            ctx.blackboard["shared"] = "one"
            return "one"

        async def write_second(ctx):
            ctx.blackboard["shared"] = "two"
            return "two"

        step = ParallelStep("par", [
            FunctionStep("first", write_first),
            FunctionStep("second", write_second),
        ])
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "failed"
        assert result.metadata["conflict_key"] == "shared"
        assert "BlackboardConflictError" in result.metadata["error"]


class TestConditionalStep:
    @pytest.mark.asyncio
    async def test_then_branch(self):
        async def yes(ctx):
            return "yes"

        async def no(ctx):
            return "no"

        step = ConditionalStep(
            "cond",
            condition=lambda ctx: ctx["flag"] is True,
            then=FunctionStep("yes", yes),
            else_=FunctionStep("no", no),
        )
        ctx = PipelineContext(pipeline_id="p1")
        ctx["flag"] = True
        result = await step.execute(ctx)
        assert result.output == "yes"

    @pytest.mark.asyncio
    async def test_else_branch(self):
        async def yes(ctx):
            return "yes"

        async def no(ctx):
            return "no"

        step = ConditionalStep(
            "cond",
            condition=lambda ctx: ctx["flag"] is True,
            then=FunctionStep("yes", yes),
            else_=FunctionStep("no", no),
        )
        ctx = PipelineContext(pipeline_id="p1")
        ctx["flag"] = False
        result = await step.execute(ctx)
        assert result.output == "no"

    @pytest.mark.asyncio
    async def test_no_else_skips(self):
        async def yes(ctx):
            return "yes"

        step = ConditionalStep(
            "cond",
            condition=lambda ctx: False,
            then=FunctionStep("yes", yes),
        )
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "skipped"


from logos_adk.workflow.steps import LoopStep, ApprovalStep


class TestLoopStep:
    @pytest.mark.asyncio
    async def test_loop_until_condition(self):
        call_count = 0

        async def increment(ctx):
            nonlocal call_count
            call_count += 1
            ctx["counter"] = call_count
            return call_count

        step = LoopStep(
            "loop",
            step=FunctionStep("inc", increment),
            until=lambda ctx: ctx.get("counter", 0) >= 3,
            max_iterations=10,
        )
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert result.output == 3
        assert call_count == 3

    @pytest.mark.asyncio
    async def test_loop_max_iterations(self):
        async def never_done(ctx):
            return "again"

        step = LoopStep(
            "loop",
            step=FunctionStep("nd", never_done),
            until=lambda ctx: False,
            max_iterations=5,
        )
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "completed"
        assert result.metadata.get("iterations") == 5

    @pytest.mark.asyncio
    async def test_loop_step_fails(self):
        async def fail(ctx):
            raise ValueError("boom")

        step = LoopStep(
            "loop",
            step=FunctionStep("fail", fail),
            until=lambda ctx: True,
            max_iterations=3,
        )
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "failed"

    @pytest.mark.asyncio
    async def test_loop_propagates_paused_status(self):
        step = LoopStep(
            "loop",
            step=ApprovalStep("review", message="need approval"),
            until=lambda ctx: False,
            max_iterations=3,
        )
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "paused"
        assert result.metadata["termination_reason"] == "paused"
        assert result.metadata["iterations"] == 1


class TestApprovalStep:
    @pytest.mark.asyncio
    async def test_approval_pauses(self):
        step = ApprovalStep("approve", message="Please review")
        ctx = PipelineContext(pipeline_id="p1")
        result = await step.execute(ctx)
        assert result.status == "paused"
        assert result.metadata.get("message") == "Please review"


class TestEventSteps:
    @pytest.mark.asyncio
    async def test_publish_event_step_publishes_payload(self):
        bus = InMemoryEventBus()
        step = PublishEventStep(
            "publish",
            bus,
            topic="need_seo",
            payload=lambda ctx: {"prompt": ctx["_last"]},
        )
        ctx = PipelineContext(pipeline_id="p1", metadata={"correlation_id": "corr-1"})
        ctx["_last"] = "optimize this"

        result = await step.execute(ctx)

        assert result.status == "completed"
        assert result.output == {"prompt": "optimize this"}
        events = bus.list_events(topic="need_seo")
        assert len(events) == 1
        assert events[0].payload == {"prompt": "optimize this"}
        assert events[0].correlation_id == "corr-1"

    @pytest.mark.asyncio
    async def test_wait_for_event_step_pauses_until_event_arrives(self):
        bus = InMemoryEventBus()
        step = WaitForEventStep("wait", bus, topic="seo_done", message="waiting")
        ctx = PipelineContext(pipeline_id="p1", metadata={"correlation_id": "corr-2"})

        paused = await step.execute(ctx)
        assert paused.status == "paused"
        assert paused.metadata["message"] == "waiting"

        bus.publish("seo_done", {"optimized": "ready"}, correlation_id="corr-2", run_id="p1")
        completed = await step.execute(ctx)
        assert completed.status == "completed"
        assert completed.output == {"optimized": "ready"}
        events = bus.list_events(topic="seo_done")
        assert events[0].status == "completed"
