from __future__ import annotations

import asyncio

import pytest
from pydantic import BaseModel

from agentbyte.workflow import (
    CheckpointConfig,
    InMemoryCheckpointStore,
    StepFailedEvent,
    StepMetadata,
    Workflow,
    WorkflowConfig,
    WorkflowContext,
    WorkflowFailedError,
    WorkflowFailedEvent,
    WorkflowRunner,
    WorkflowStatus,
)
from agentbyte.workflow.steps import FunctionStep, SubWorkflowStep
from agentbyte.workflow.steps.step import BaseStep


class NumberInput(BaseModel):
    value: int


class NumberOutput(BaseModel):
    result: int


class CommentResponse(BaseModel):
    comment: str


def _build_child_math_workflow() -> Workflow:
    async def double(input_data: NumberInput, _: WorkflowContext) -> NumberOutput:
        await asyncio.sleep(0)
        return NumberOutput(result=input_data.value * 2)

    async def add_ten(input_data: NumberOutput, _: WorkflowContext) -> NumberOutput:
        await asyncio.sleep(0)
        return NumberOutput(result=input_data.result + 10)

    return Workflow(WorkflowConfig(name="child-math")).chain(
        FunctionStep(
            step_id="double",
            metadata=StepMetadata(name="double"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=double,
        ),
        FunctionStep(
            step_id="add_ten",
            metadata=StepMetadata(name="add_ten"),
            input_type=NumberOutput,
            output_type=NumberOutput,
            func=add_ten,
        ),
    )


@pytest.mark.asyncio
async def test_subworkflow_step_runs_child_workflow_to_completion() -> None:
    parent = Workflow(WorkflowConfig(name="parent-subworkflow"))
    parent.set_start_step(
        SubWorkflowStep(
            step_id="child",
            metadata=StepMetadata(name="child"),
            workflow=_build_child_math_workflow,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    ).add_end_step("child")

    execution = await WorkflowRunner().run(parent, {"value": 3})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["child"].output_data == {"result": 16}
    assert execution.state["child_output"] == {"result": 16}


@pytest.mark.asyncio
async def test_subworkflow_step_propagates_child_request_through_parent_checkpoint() -> None:
    class ChildReviewStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
        ) -> NumberOutput:
            await context.request_input(
                key="review_comment",
                response_type=CommentResponse,
                prompt="Provide review comment",
            )
            raise AssertionError("request_input should suspend execution")

        async def resume_from_request(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
            request,
            response: CommentResponse,
        ) -> NumberOutput:
            context.set("child_review_comment", response.comment)
            return NumberOutput(result=input_data.value + len(response.comment))

    def build_child_review_workflow() -> Workflow:
        step = ChildReviewStep()
        return (
            Workflow(WorkflowConfig(name="child-review"))
            .set_start_step(step)
            .add_end_step(step)
        )

    parent = Workflow(WorkflowConfig(name="parent-review"))
    parent.set_start_step(
        SubWorkflowStep(
            step_id="child",
            metadata=StepMetadata(name="child"),
            workflow=build_child_review_workflow,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    ).add_end_step("child")

    store = InMemoryCheckpointStore()
    checkpoint_config = CheckpointConfig(store=store, auto_save=True, save_interval_steps=1)

    suspended_execution = await WorkflowRunner().run(
        parent,
        {"value": 2},
        checkpoint_config=checkpoint_config,
    )

    assert suspended_execution.status == WorkflowStatus.SUSPENDED
    assert len(suspended_execution.pending_requests) == 1
    parent_request_id = next(iter(suspended_execution.pending_requests))
    parent_request = suspended_execution.pending_requests[parent_request_id]
    assert parent_request.metadata["subworkflow"]["child_request_key"] == "review_comment"

    latest = await store.load_latest("parent-review")
    assert latest is not None

    resumed_execution = await WorkflowRunner().run(
        parent,
        checkpoint=latest,
        responses={parent_request_id: {"comment": "ok"}},
    )

    assert resumed_execution.status == WorkflowStatus.COMPLETED
    assert resumed_execution.step_executions["child"].output_data == {"result": 4}
    assert resumed_execution.state["child_output"] == {"result": 4}


@pytest.mark.asyncio
async def test_two_subworkflow_steps_sharing_a_workflow_builder_remain_isolated() -> None:
    class CountingStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="count",
                metadata=StepMetadata(name="count"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )
            self.calls = 0

        async def execute(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
        ) -> NumberOutput:
            del context
            self.calls += 1
            await asyncio.sleep(0)
            return NumberOutput(result=input_data.value + self.calls)

    shared_child = (
        Workflow(WorkflowConfig(name="shared-child"))
        .set_start_step(CountingStep())
        .add_end_step("count")
    )

    async def start(input_data: NumberInput, _: WorkflowContext) -> NumberInput:
        return input_data

    parent = Workflow(
        WorkflowConfig(
            name="parallel-subworkflow",
            completion_mode="all_end",
        )
    )
    parent.set_start_step(
        FunctionStep(
            step_id="start",
            metadata=StepMetadata(name="start"),
            input_type=NumberInput,
            output_type=NumberInput,
            func=start,
        )
    )
    parent.add_step(
        SubWorkflowStep(
            step_id="child_a",
            metadata=StepMetadata(name="child_a"),
            workflow=shared_child,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    )
    parent.add_step(
        SubWorkflowStep(
            step_id="child_b",
            metadata=StepMetadata(name="child_b"),
            workflow=shared_child,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    )
    parent.add_edge("start", "child_a")
    parent.add_edge("start", "child_b")
    parent.add_end_step("child_a").add_end_step("child_b")

    execution = await WorkflowRunner(max_concurrent_steps=3).run(parent, {"value": 10})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["child_a"].output_data == {"result": 11}
    assert execution.step_executions["child_b"].output_data == {"result": 11}


@pytest.mark.asyncio
async def test_subworkflow_step_fails_when_child_workflow_fails() -> None:
    async def explode(input_data: NumberInput, _: WorkflowContext) -> NumberOutput:
        del input_data
        raise RuntimeError("child boom")

    child = Workflow(WorkflowConfig(name="child-failure")).set_start_step(
        FunctionStep(
            step_id="explode",
            metadata=StepMetadata(name="explode"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=explode,
        )
    ).add_end_step("explode")

    parent = Workflow(WorkflowConfig(name="parent-failure"))
    parent.set_start_step(
        SubWorkflowStep(
            step_id="child",
            metadata=StepMetadata(name="child"),
            workflow=child,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    ).add_end_step("child")

    events = [event async for event in WorkflowRunner().run_stream(parent, {"value": 1})]

    step_failed_events = [event for event in events if isinstance(event, StepFailedEvent)]
    assert len(step_failed_events) == 1
    assert step_failed_events[0].step_id == "child"
    assert "child boom" in step_failed_events[0].error
    assert isinstance(events[-1], WorkflowFailedEvent)
    assert "child boom" in events[-1].error


@pytest.mark.asyncio
async def test_subworkflow_step_supports_input_and_output_mappers() -> None:
    child = _build_child_math_workflow()

    parent = Workflow(
        WorkflowConfig(
            name="mapped-subworkflow",
            initial_state={"offset": 2},
        )
    )
    parent.set_start_step(
        SubWorkflowStep(
            step_id="child",
            metadata=StepMetadata(name="child"),
            workflow=child,
            input_type=NumberInput,
            output_type=NumberOutput,
            input_mapper=lambda input_data, context: {
                "value": input_data.value + context.get("offset", 0)
            },
            output_mapper=lambda execution, workflow, context: {
                "result": execution.step_executions[
                    workflow.end_step_ids[0]
                ].output_data["result"]
                + context.get("offset", 0)
            },
        )
    ).add_end_step("child")

    execution = await WorkflowRunner().run(parent, {"value": 3})

    assert execution.status == WorkflowStatus.COMPLETED
    assert execution.step_executions["child"].output_data == {"result": 22}


@pytest.mark.asyncio
async def test_subworkflow_step_rejects_multiple_child_pending_requests() -> None:
    class ReviewStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self, step_id: str) -> None:
            super().__init__(
                step_id=step_id,
                metadata=StepMetadata(name=step_id),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
        ) -> NumberOutput:
            await context.request_input(
                key=f"{self.step_id}_comment",
                response_type=CommentResponse,
                prompt=f"Provide comment for {self.step_id}",
            )
            raise AssertionError("request_input should suspend execution")

    async def start(input_data: NumberInput, _: WorkflowContext) -> NumberInput:
        return input_data

    child = Workflow(
        WorkflowConfig(
            name="child-multi-request",
            completion_mode="all_end",
        )
    )
    child.set_start_step(
        FunctionStep(
            step_id="start",
            metadata=StepMetadata(name="start"),
            input_type=NumberInput,
            output_type=NumberInput,
            func=start,
        )
    )
    child.add_step(ReviewStep("left"))
    child.add_step(ReviewStep("right"))
    child.add_edge("start", "left")
    child.add_edge("start", "right")
    child.add_end_step("left").add_end_step("right")

    parent = Workflow(WorkflowConfig(name="parent-multi-request"))
    parent.set_start_step(
        SubWorkflowStep(
            step_id="child",
            metadata=StepMetadata(name="child"),
            workflow=child,
            input_type=NumberInput,
            output_type=NumberOutput,
        )
    ).add_end_step("child")

    with pytest.raises(
        WorkflowFailedError,
        match="does not yet support propagating multiple concurrent child pending requests",
    ):
        await WorkflowRunner().run(parent, {"value": 1})


@pytest.mark.asyncio
async def test_subworkflow_step_clears_stale_state_on_fresh_rerun() -> None:
    class ConditionalReviewStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
        ) -> NumberOutput:
            if input_data.value == 0:
                await context.request_input(
                    key="review_comment",
                    response_type=CommentResponse,
                    prompt="Provide review comment",
                )
                raise AssertionError("request_input should suspend execution")
            return NumberOutput(result=input_data.value)

        async def resume_from_request(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
            request,
            response: CommentResponse,
        ) -> NumberOutput:
            del context, request
            return NumberOutput(result=input_data.value + len(response.comment))

    child = (
        Workflow(WorkflowConfig(name="child-conditional-review"))
        .set_start_step(ConditionalReviewStep())
        .add_end_step("review")
    )
    step = SubWorkflowStep(
        step_id="child",
        metadata=StepMetadata(name="child"),
        workflow=child,
        input_type=NumberInput,
        output_type=NumberOutput,
    )
    parent = Workflow(WorkflowConfig(name="parent-rerun"))
    parent.set_start_step(step).add_end_step("child")

    first_execution = await WorkflowRunner().run(parent, {"value": 0})

    assert first_execution.status == WorkflowStatus.SUSPENDED
    assert step._active_invocations
    assert step._request_routes

    second_execution = await WorkflowRunner().run(parent, {"value": 5})

    assert second_execution.status == WorkflowStatus.COMPLETED
    assert second_execution.step_executions["child"].output_data == {"result": 5}
    assert step._active_invocations == {}
    assert step._request_routes == {}
