# tests/test_workflow_pipeline.py
"""Tests for Pipeline orchestrator."""
import pytest
from logos_adk.workflow.pipeline import Pipeline
from logos_adk.types import Response


class TestPipelineBuilder:
    def test_create(self):
        p = Pipeline("test")
        assert p.name == "test"

    def test_step_returns_self(self):
        async def fn(ctx):
            return "ok"
        p = Pipeline("test")
        result = p.function(fn, name="f1")
        assert result is p

    def test_chaining(self):
        async def fn(ctx):
            return "ok"
        p = Pipeline("test").function(fn, name="f1").function(fn, name="f2")
        assert p._steps_config  # internal list has entries


class TestPipelineExecution:
    @pytest.mark.asyncio
    async def test_empty_pipeline(self):
        p = Pipeline("empty")
        response = await p.run("hello")
        assert isinstance(response, Response)
        assert response.content == "hello"

    @pytest.mark.asyncio
    async def test_single_function_step(self):
        async def upper(ctx):
            return ctx["_last"].upper()

        p = Pipeline("upper").function(upper, name="up")
        response = await p.run("hello")
        assert response.content == "HELLO"

    @pytest.mark.asyncio
    async def test_chained_steps(self):
        async def add_world(ctx):
            return ctx["_last"] + " world"

        async def upper(ctx):
            return ctx["_last"].upper()

        p = Pipeline("chain").function(add_world, name="add").function(upper, name="up")
        response = await p.run("hello")
        assert response.content == "HELLO WORLD"

    @pytest.mark.asyncio
    async def test_context_carries_named_results(self):
        async def step_a(ctx):
            return "from_a"

        async def step_b(ctx):
            return f"got {ctx['a']}"

        p = Pipeline("ctx").function(step_a, name="a").function(step_b, name="b")
        response = await p.run("")
        assert response.content == "got from_a"

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

        agent = Agent("mock", model=MockProvider(responses=["mock reply"]))
        p = Pipeline("agent_pipe").step(agent, name="a1")
        response = await p.run("test prompt")
        assert "mock reply" in response.content

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

        async def fn_b(ctx):
            return "B"

        p = Pipeline("par").parallel(fn_a, fn_b, name="par")
        response = await p.run("")
        raw = response.raw.get("pipeline_result", {})
        assert raw["outputs"]["par"]["output"]["fn_a"] == "A"
        assert raw["outputs"]["par"]["output"]["fn_b"] == "B"

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

        async def no(ctx):
            return "no"

        async def set_flag(ctx):
            ctx["flag"] = True
            return "flagged"

        p = Pipeline("cond") \
            .function(set_flag, name="setup") \
            .conditional(lambda ctx: ctx["flag"], then=yes, else_=no, name="branch")
        response = await p.run("")
        assert response.content == "yes"

    @pytest.mark.asyncio
    async def test_loop(self):
        async def increment(ctx):
            ctx["count"] = ctx.get("count", 0) + 1
            return ctx["count"]

        p = Pipeline("loop").loop(
            increment,
            until=lambda ctx: ctx.get("count", 0) >= 3,
            max_iterations=10,
            name="counter",
        )
        response = await p.run("")
        assert response.content == "3"

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

        async def fn_b(ctx):
            return ctx["_last"] + "B"

        inner = Pipeline("inner").function(fn_a, name="a").function(fn_b, name="b")
        outer = Pipeline("outer").step(inner, name="sub")
        response = await outer.run("")
        assert response.content == "AB"

    @pytest.mark.asyncio
    async def test_pipeline_result_in_raw(self):
        async def fn(ctx):
            return "done"

        p = Pipeline("raw").function(fn, name="f")
        response = await p.run("input")
        assert "pipeline_result" in response.raw

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

        async def after(ctx):
            return "should not run"

        p = Pipeline("fail").function(fail, name="bad").function(after, name="after")
        response = await p.run("")
        assert response.raw["pipeline_result"]["status"] == "failed"
