from __future__ import annotations

import asyncio
from typing import Any

import httpx
import pytest
from pydantic import BaseModel

from agentbyte.agents import BaseAgent
from agentbyte.agents.types import AgentResponse
from agentbyte.context import AgentContext
from agentbyte.messages import AssistantMessage, Usage
from agentbyte.workflow.core.models import StepMetadata, WorkflowContext
from agentbyte.workflow.schema_utils import coerce_value_to_schema_type
from agentbyte.workflow.steps.agentbyte_agent import AgentStep, AgentbyteAgentInput
from agentbyte.workflow.steps.echo import EchoInput, EchoStep
from agentbyte.workflow.steps.function import FunctionStep
from agentbyte.workflow.steps.http import HttpRequestInput, HttpStep
from agentbyte.workflow.steps.step import BaseStep
from agentbyte.workflow.steps.transform import TransformStep


class NumberInput(BaseModel):
    value: int


class NumberOutput(BaseModel):
    result: int


@pytest.mark.asyncio
async def test_function_step_handles_sync_and_async() -> None:
    def sync_double(input_data: NumberInput, _: WorkflowContext) -> dict[str, int]:
        return {"result": input_data.value * 2}

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

    metadata = StepMetadata(name="math")

    step_sync = FunctionStep("double", metadata, NumberInput, NumberOutput, sync_double)
    step_async = FunctionStep("square", metadata, NumberInput, NumberOutput, async_square)

    sync_out = await step_sync.run({"value": 3}, {})
    async_out = await step_async.run({"value": 3}, {})

    assert sync_out["result"] == 6
    assert async_out["result"] == 9


@pytest.mark.asyncio
async def test_base_step_retry_and_timeout() -> None:
    class FlakyStep(BaseStep[NumberInput, NumberOutput]):
        def __init__(self) -> None:
            super().__init__(
                step_id="flaky",
                metadata=StepMetadata(name="flaky", max_retries=1),
                input_type=NumberInput,
                output_type=NumberOutput,
            )
            self.attempts = 0

        async def execute(self, input_data: NumberInput, context: WorkflowContext) -> NumberOutput:  # noqa: ARG002
            self.attempts += 1
            if self.attempts == 1:
                raise RuntimeError("fail once")
            return NumberOutput(result=input_data.value)

    step = FlakyStep()
    output = await step.run({"value": 5}, {})
    assert output["result"] == 5
    assert step.attempts == 2

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

        async def execute(self, input_data: NumberInput, context: WorkflowContext) -> NumberOutput:  # noqa: ARG002
            await asyncio.sleep(0.1)
            return NumberOutput(result=input_data.value)

    slow_step = SlowStep()
    with pytest.raises(asyncio.TimeoutError):
        await slow_step.run({"value": 1}, {})


@pytest.mark.asyncio
async def test_echo_step() -> None:
    metadata = StepMetadata(name="echo")
    step = EchoStep("echo", metadata, prefix="[", suffix="]")

    output = await step.run(EchoInput(message="hi").model_dump(), {})
    assert output["result"] == "[hi]"


@pytest.mark.asyncio
async def test_transform_step_mapping_and_static() -> None:
    class InputModel(BaseModel):
        name: str
        age: int

    class OutputModel(BaseModel):
        full_name: str
        age_str: str
        note: str

    metadata = StepMetadata(name="transform")
    mappings = {
        "full_name": "name",
        "age_str": "age",
        "note": "static:ok",
    }
    step = TransformStep("transform", metadata, InputModel, OutputModel, mappings)

    output = await step.run({"name": "Ada", "age": 30}, {})
    assert output["full_name"] == "Ada"
    assert output["age_str"] == "30"
    assert output["note"] == "ok"

    # Coercion helper sanity check
    assert coerce_value_to_schema_type("5", {"type": "integer"}) == 5


@pytest.mark.asyncio
async def test_http_step_with_mock_transport() -> None:
    async def handler(request: httpx.Request) -> httpx.Response:  # noqa: ARG001
        return httpx.Response(200, json={"ok": True})

    transport = httpx.MockTransport(handler)
    async with httpx.AsyncClient(transport=transport) as client:
        metadata = StepMetadata(name="http")
        step = HttpStep("http", metadata, client=client)
        output = await step.run(HttpRequestInput(url="https://example.com").model_dump(), {})

    assert output["status_code"] == 200
    assert "ok" in output["content"]


class _FakeAgent(BaseAgent):  # type: ignore[misc]
    """Minimal stub satisfying BaseAgent interface for testing."""

    def __init__(self) -> None:
        self.last_task: str | None = None

    async def run(self, task: Any, context: Any | None = None, cancellation_token: Any | None = None):  # noqa: ANN401, ARG002
        self.last_task = task
        ctx = AgentContext(messages=[AssistantMessage(content=f"done: {task}", source="assistant")])
        return AgentResponse(
            context=ctx,
            source="fake-agent",
            finish_reason="stop",
            usage=Usage(),
        )

    def run_stream(self, *args: Any, **kwargs: Any):  # noqa: ANN401
        raise NotImplementedError


@pytest.mark.asyncio
async def test_agentbyte_agent_step() -> None:
    agent = _FakeAgent()
    metadata = StepMetadata(name="agent")
    step = AgentStep("agent_step", metadata, agent)

    output = await step.run(AgentbyteAgentInput(task="hello").model_dump(), {})

    assert agent.last_task == "hello"
    assert output["response"] == "done: hello"
    assert output["messages"][0]["content"] == "done: hello"
    assert output["usage"] is not None

