from __future__ import annotations

import asyncio
from typing import Any

import pytest
from pydantic import BaseModel

from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, UserMessage
from agentbyte.orchestration.round_robin import RoundRobinOrchestrator
from agentbyte.termination.max_message import MaxMessageTermination
from agentbyte.workflow import StepMetadata, Workflow, WorkflowConfig, WorkflowContext, WorkflowRunner
from agentbyte.workflow.steps import FunctionStep
from agentbyte.workflow.steps.step import BaseStep


class TaskInput(BaseModel):
    task: str
    history: list[dict[str, Any]] = []


class TaskOutput(BaseModel):
    response: str


def _build_task_workflow(name: str = "task-workflow") -> Workflow:
    async def respond(input_data: TaskInput, _: WorkflowContext) -> TaskOutput:
        await asyncio.sleep(0)
        return TaskOutput(
            response=f"{input_data.task}|history={len(input_data.history)}"
        )

    workflow = Workflow(WorkflowConfig(name=name))
    workflow.set_start_step(
        FunctionStep(
            step_id="respond",
            metadata=StepMetadata(name="respond"),
            input_type=TaskInput,
            output_type=TaskOutput,
            func=respond,
        )
    ).add_end_step("respond")
    return workflow


@pytest.mark.asyncio
async def test_workflow_as_agent_runs_and_returns_agent_response() -> None:
    agent = _build_task_workflow("as-agent-basic").as_agent(name="workflow_agent")

    response = await agent.run("hello")

    assert response.source == "workflow_agent"
    assert response.finish_reason == "stop"
    assert response.final_content == "hello|history=0"


@pytest.mark.asyncio
async def test_workflow_as_agent_full_memory_mode_preserves_history() -> None:
    agent = _build_task_workflow("as-agent-full").as_agent(
        name="workflow_agent",
        memory_mode="full",
    )
    context = AgentContext()

    first = await agent.run("first", context=context)
    second = await agent.run("second", context=context)

    assert first.final_content == "first|history=0"
    assert second.final_content == "second|history=2"
    assert [message.role for message in context.messages] == [
        "user",
        "assistant",
        "user",
        "assistant",
    ]


@pytest.mark.asyncio
async def test_workflow_as_agent_last_result_memory_mode_keeps_latest_assistant_only() -> None:
    agent = _build_task_workflow("as-agent-last").as_agent(
        name="workflow_agent",
        memory_mode="last_result",
    )
    context = AgentContext()

    first = await agent.run("first", context=context)
    second = await agent.run("second", context=context)

    assert first.final_content == "first|history=0"
    assert second.final_content == "second|history=1"
    assert len(context.messages) == 1
    assert isinstance(context.messages[0], AssistantMessage)
    assert context.messages[0].content == "second|history=1"


@pytest.mark.asyncio
async def test_workflow_as_agent_maps_suspension_to_agent_response() -> None:
    class ReviewResponse(BaseModel):
        comment: str

    class SuspendStep(BaseStep[TaskInput, TaskOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=TaskInput,
                output_type=TaskOutput,
            )

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

    workflow = Workflow(WorkflowConfig(name="as-agent-suspend"))
    workflow.set_start_step(SuspendStep()).add_end_step("review")
    agent = workflow.as_agent(name="workflow_agent")
    context = AgentContext()

    response = await agent.run("hello", context=context)

    assert response.finish_reason == "suspended"
    assert response.context.metadata["workflow_agent"]["workflow_status"] == "suspended"
    assert len(response.context.metadata["workflow_agent"]["pending_requests"]) == 1
    assert response.context.metadata["workflow_agent"]["checkpoint_id"] is not None
    assert "checkpoint" not in response.context.metadata["workflow_agent"]
    assert "suspended" in response.final_content.lower()


@pytest.mark.asyncio
async def test_workflow_as_agent_can_resume_suspended_workflow() -> None:
    class ReviewResponse(BaseModel):
        comment: str

    class SuspendStep(BaseStep[TaskInput, TaskOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="review",
                metadata=StepMetadata(name="review"),
                input_type=TaskInput,
                output_type=TaskOutput,
            )

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

        async def resume_from_request(
            self,
            input_data: TaskInput,
            context: WorkflowContext,
            request,
            response: ReviewResponse,
        ) -> TaskOutput:
            del context, request
            return TaskOutput(response=f"resumed:{response.comment}")

    workflow = Workflow(WorkflowConfig(name="as-agent-resume"))
    workflow.set_start_step(SuspendStep()).add_end_step("review")
    agent = workflow.as_agent(name="workflow_agent")
    context = AgentContext()

    first = await agent.run("hello", context=context)

    assert first.finish_reason == "suspended"
    checkpoint_id = first.context.metadata["workflow_agent"]["checkpoint_id"]
    assert checkpoint_id

    resumed = await agent.run({"comment": "ok"}, context=context)

    assert resumed.finish_reason == "stop"
    assert resumed.final_content == "resumed:ok"
    assert resumed.context.metadata["workflow_agent"]["workflow_status"] == "completed"
    assert resumed.context.metadata["workflow_agent"]["checkpoint_id"] is None


@pytest.mark.asyncio
async def test_workflow_as_agent_maps_failure_to_error_response() -> None:
    async def explode(input_data: TaskInput, _: WorkflowContext) -> TaskOutput:
        del input_data
        raise RuntimeError("boom")

    workflow = Workflow(WorkflowConfig(name="as-agent-failure"))
    workflow.set_start_step(
        FunctionStep(
            step_id="explode",
            metadata=StepMetadata(name="explode"),
            input_type=TaskInput,
            output_type=TaskOutput,
            func=explode,
        )
    ).add_end_step("explode")
    agent = workflow.as_agent(name="workflow_agent")

    response = await agent.run("hello")

    assert response.finish_reason == "error"
    assert response.context.metadata["workflow_agent"]["workflow_status"] == "failed"
    assert "boom" in response.final_content


@pytest.mark.asyncio
async def test_workflow_as_agent_custom_memory_mode_uses_history_selector() -> None:
    agent = _build_task_workflow("as-agent-custom").as_agent(
        name="workflow_agent",
        memory_mode="custom",
        history_selector=lambda context: [
            message
            for message in context.messages
            if isinstance(message, AssistantMessage)
        ],
    )
    context = AgentContext(
        messages=[
            UserMessage(content="u1", source="user"),
            AssistantMessage(content="a1", source="workflow_agent"),
            UserMessage(content="u2", source="user"),
            AssistantMessage(content="a2", source="workflow_agent"),
        ]
    )

    response = await agent.run("next", context=context)

    assert response.finish_reason == "stop"
    assert response.final_content == "next|history=2"


@pytest.mark.asyncio
async def test_workflow_as_agent_supports_input_and_output_mappers() -> None:
    agent = _build_task_workflow("as-agent-mappers").as_agent(
        name="workflow_agent",
        input_mapper=lambda task, context, history: {
            "task": f"{task}:{context.metadata['prefix']}",
            "history": [message.model_dump(mode="json") for message in history],
        },
        output_mapper=lambda execution, workflow, context: {
            "response": (
                execution.step_executions[workflow.end_step_ids[0]].output_data["response"]
                + f":{context.metadata['suffix']}"
            )
        },
    )
    context = AgentContext(metadata={"prefix": "in", "suffix": "out"})

    response = await agent.run("hello", context=context)

    assert response.finish_reason == "stop"
    assert response.final_content == "hello:in|history=0:out"


@pytest.mark.asyncio
async def test_workflow_as_agent_maps_pre_run_exceptions_to_error_response() -> None:
    workflow = Workflow(WorkflowConfig(name="as-agent-pre-run-error"))
    agent = workflow.as_agent(name="workflow_agent")

    response = await agent.run("hello")

    assert response.finish_reason == "error"
    assert "start step" in response.final_content.lower()


def test_workflow_as_agent_forwards_custom_runner() -> None:
    runner = WorkflowRunner(max_concurrent_steps=1)
    agent = _build_task_workflow("as-agent-runner").as_agent(
        name="workflow_agent",
        runner=runner,
    )

    assert agent._runner is runner


@pytest.mark.asyncio
async def test_workflow_as_agent_works_in_round_robin_orchestrator() -> None:
    agent = _build_task_workflow("as-agent-orchestrator").as_agent(
        name="workflow_agent",
        memory_mode="last_result",
    )
    orchestrator = RoundRobinOrchestrator(
        [agent],
        MaxMessageTermination(2),
        max_iterations=1,
    )

    response = await orchestrator.run(UserMessage(content="start", source="user"))

    assert response.messages[-1].source == "workflow_agent"
    assert "start" in response.messages[-1].content
