from __future__ import annotations

import asyncio

import pytest
from pydantic import BaseModel

from agentbyte.workflow import (
    CheckpointConfig,
    FileCheckpointStore,
    InMemoryCheckpointStore,
    StepMetadata,
    Workflow,
    WorkflowCheckpoint,
    WorkflowConfig,
    WorkflowFailedError,
    WorkflowRunError,
    WorkflowRunner,
)
from agentbyte.workflow.core.models import (
    StepExecution,
    StepStatus,
    WorkflowContext,
    WorkflowExecution,
    WorkflowStatus,
)
from agentbyte.workflow.steps.function import FunctionStep
from agentbyte.workflow.steps.step import BaseStep


class NumberInput(BaseModel):
    value: int


class NumberOutput(BaseModel):
    result: int


class CommentResponse(BaseModel):
    comment: str


@pytest.mark.asyncio
async def test_workflow_checkpoint_factory_and_inmemory_store() -> None:
    execution = WorkflowExecution(
        workflow_id="wf",
        state={"x": 1},
        step_executions={
            "a": StepExecution(step_id="a", status=StepStatus.SUSPENDED),
        },
    )
    checkpoint = WorkflowCheckpoint.from_execution(
        execution=execution,
        workflow_id="wf",
        workflow_version="1.0.0",
        workflow_structure_hash="abc123",
        all_step_ids=["a", "b"],
    )
    store = InMemoryCheckpointStore()
    await store.save(checkpoint)

    loaded = await store.load(checkpoint.checkpoint_id)
    latest = await store.load_latest("wf")
    metadata = await store.list_metadata("wf")

    assert loaded is not None
    assert latest is not None
    assert latest.checkpoint_id == checkpoint.checkpoint_id
    assert metadata[0].total_steps == 2
    assert checkpoint.pending_step_ids == ["a", "b"]
    assert metadata[0].pending_steps == 2


@pytest.mark.asyncio
async def test_file_checkpoint_store_roundtrip(tmp_path) -> None:
    store = FileCheckpointStore(tmp_path / "checkpoints")
    execution = WorkflowExecution(workflow_id="wf")
    checkpoint = WorkflowCheckpoint.from_execution(
        execution=execution,
        workflow_id="wf",
        workflow_version="1.0.0",
        workflow_structure_hash="hash",
        all_step_ids=["a"],
    )

    await store.save(checkpoint)
    loaded = await store.load(checkpoint.checkpoint_id)
    latest = await store.load_latest("wf")
    metadata = await store.list_metadata("wf")

    assert loaded is not None
    assert latest is not None
    assert loaded.checkpoint_id == checkpoint.checkpoint_id
    assert metadata[0].size_bytes is not None


@pytest.mark.asyncio
async def test_runner_auto_saves_and_resumes_from_checkpoint() -> None:
    async def first_step(input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
        context.set("seen_first", True)
        await asyncio.sleep(0)
        return NumberOutput(result=input_data.value + 1)

    class FlakyStep(BaseStep[NumberOutput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="second",
                metadata=StepMetadata(name="second"),
                input_type=NumberOutput,
                output_type=NumberOutput,
            )
            self.calls = 0

        async def execute(self, input_data: NumberOutput, context: WorkflowContext) -> NumberOutput:
            self.calls += 1
            if self.calls == 1:
                raise RuntimeError("fail once")
            context.set("resumed", True)
            return NumberOutput(result=input_data.result + 10)

    first = FunctionStep(
        step_id="first",
        metadata=StepMetadata(name="first"),
        input_type=NumberInput,
        output_type=NumberOutput,
        func=first_step,
    )
    second = FlakyStep()
    workflow = Workflow(WorkflowConfig(name="resume-flow"))
    workflow.chain(first, second)

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

    with pytest.raises(RuntimeError, match="fail once"):
        await runner.run(workflow, {"value": 1}, checkpoint_config=checkpoint_config)

    latest = await store.load_latest("resume-flow")
    assert latest is not None
    assert latest.completed_step_ids == ["first"]

    execution = await runner.run(
        workflow,
        {"value": 999},
        checkpoint=latest,
        checkpoint_config=checkpoint_config,
    )

    assert second.calls == 2
    assert execution.step_executions["first"].output_data == {"result": 2}
    assert execution.step_executions["second"].output_data == {"result": 12}
    assert execution.state["resumed"] is True


@pytest.mark.asyncio
async def test_runner_rejects_mismatched_checkpoint() -> None:
    async def passthrough(input_data: NumberInput, _: WorkflowContext) -> NumberOutput:
        return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="wf-a"))
    workflow.set_start_step(
        FunctionStep(
            step_id="only",
            metadata=StepMetadata(name="only"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=passthrough,
        )
    ).add_end_step("only")

    other_workflow = Workflow(WorkflowConfig(name="wf-b"))
    other_workflow.set_start_step(
        FunctionStep(
            step_id="only",
            metadata=StepMetadata(name="only"),
            input_type=NumberInput,
            output_type=NumberOutput,
            func=passthrough,
        )
    ).add_end_step("only")

    checkpoint = WorkflowCheckpoint.from_execution(
        execution=WorkflowExecution(workflow_id="wf-b"),
        workflow_id="wf-b",
        workflow_version="1.0.0",
        workflow_structure_hash=other_workflow.compute_structure_hash(),
        all_step_ids=["only"],
    )

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

    assert "Checkpoint validation failed" in events[-1].error

@pytest.mark.asyncio
async def test_runner_saves_suspended_checkpoint_and_resumes_with_response() -> None:
    class ReviewStep(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("review_comment", response.comment)
            return NumberOutput(result=input_data.value + len(response.comment))

    def build_workflow() -> Workflow:
        step = ReviewStep()
        return (
            Workflow(WorkflowConfig(name="checkpoint-review"))
            .set_start_step(step)
            .add_end_step(step)
        )

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

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

    assert suspended_execution.status == WorkflowStatus.SUSPENDED
    latest = await store.load_latest("checkpoint-review")
    assert latest is not None
    assert latest.execution.status == WorkflowStatus.SUSPENDED
    assert len(latest.execution.pending_requests) == 1

    request_id = next(iter(latest.execution.pending_requests))
    pending_request = latest.execution.pending_requests[request_id]
    assert pending_request.response_type_module == __name__
    assert pending_request.response_type_qualname == "CommentResponse"

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

    assert resumed_execution.status == WorkflowStatus.COMPLETED
    assert resumed_execution.step_executions["review"].output_data == {"result": 4}
    assert resumed_execution.state["review_comment"] == "ok"


@pytest.mark.asyncio
async def test_runner_checkpoint_resume_fails_for_local_response_type() -> None:
    class LocalCommentResponse(BaseModel):
        approved: bool

    class ReviewStep(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=LocalCommentResponse,
                prompt="Approve this review?",
            )
            raise AssertionError("request_input should suspend execution")

        async def resume_from_request(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
            request,
            response: LocalCommentResponse,
        ) -> NumberOutput:
            return NumberOutput(result=input_data.value if response.approved else 0)

    def build_workflow() -> Workflow:
        step = ReviewStep()
        return (
            Workflow(WorkflowConfig(name="checkpoint-local-type"))
            .set_start_step(step)
            .add_end_step(step)
        )

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

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

    assert suspended_execution.status == WorkflowStatus.SUSPENDED
    latest = await store.load_latest("checkpoint-local-type")
    assert latest is not None

    request_id = next(iter(latest.execution.pending_requests))
    with pytest.raises(
        WorkflowFailedError,
        match="Cannot restore a locally defined response type",
    ):
        await WorkflowRunner().run(
            build_workflow(),
            checkpoint=latest,
            responses={request_id: {"approved": True}},
        )


@pytest.mark.asyncio
async def test_runner_rejects_resume_with_both_initial_input_and_responses() -> None:
    """Providing initial_input and responses together must raise ValueError."""

    class PassStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="pass",
                metadata=StepMetadata(name="pass"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="mutual-exclusion")).set_start_step(PassStep()).add_end_step("pass")
    runner = WorkflowRunner()

    with pytest.raises(ValueError, match="Cannot provide both initial_input and responses"):
        await runner.run(workflow, initial_input={"value": 1}, responses={"req-1": {}})


@pytest.mark.asyncio
async def test_fresh_runner_resume_without_checkpoint_raises() -> None:
    """Resuming via responses= on a fresh runner with no suspended execution must raise ValueError."""

    class PassStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="pass",
                metadata=StepMetadata(name="pass"),
                input_type=NumberInput,
                output_type=NumberOutput,
            )

        async def execute(self, input_data: NumberInput, context: WorkflowContext) -> NumberOutput:
            return NumberOutput(result=input_data.value)

    workflow = Workflow(WorkflowConfig(name="fresh-runner-test")).set_start_step(PassStep()).add_end_step("pass")

    # The runner wraps the ValueError into WorkflowFailedError via the event stream.
    with pytest.raises(WorkflowRunError, match="No suspended workflow execution available"):
        await WorkflowRunner().run(workflow, responses={"req-1": {}})


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

        async def execute(
            self,
            input_data: NumberInput,
            context: WorkflowContext,
        ) -> NumberOutput:
            self.cached_base = input_data.value + 5
            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:
            assert self.cached_base is not None
            context.set("cached_base", self.cached_base)
            return NumberOutput(result=self.cached_base + len(response.comment))

        def on_checkpoint_save(self) -> dict[str, int]:
            return {"cached_base": self.cached_base} if self.cached_base is not None else {}

        def on_checkpoint_restore(self, state: dict[str, int]) -> None:
            self.cached_base = state["cached_base"]

    def build_workflow(step: StatefulReviewStep) -> Workflow:
        return (
            Workflow(WorkflowConfig(name="checkpoint-stateful-review"))
            .set_start_step(step)
            .add_end_step(step)
        )

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

    original_step = StatefulReviewStep()
    suspended_execution = await WorkflowRunner().run(
        build_workflow(original_step),
        {"value": 3},
        checkpoint_config=checkpoint_config,
    )

    assert suspended_execution.status == WorkflowStatus.SUSPENDED
    latest = await store.load_latest("checkpoint-stateful-review")
    assert latest is not None
    assert latest.step_states["review"] == {"cached_base": 8}

    restored_step = StatefulReviewStep()
    request_id = next(iter(latest.execution.pending_requests))
    resumed_execution = await WorkflowRunner().run(
        build_workflow(restored_step),
        checkpoint=latest,
        responses={request_id: {"comment": "hey"}},
    )

    assert resumed_execution.step_executions["review"].output_data == {"result": 11}
    assert resumed_execution.state["cached_base"] == 8
