"""FunctionStep wraps a Python callable as a workflow step."""

from __future__ import annotations

import asyncio
from typing import Any, Callable, Type

from pydantic import BaseModel

from agentbyte.workflow.core.models import InputType, OutputType, StepMetadata, WorkflowContext
from agentbyte.workflow.steps.step import BaseStep


class FunctionStep(BaseStep[InputType, OutputType]):
    """Adapts a callable into a workflow step."""

    def __init__(
        self,
        step_id: str,
        metadata: StepMetadata,
        input_type: Type[InputType],
        output_type: Type[OutputType],
        func: Callable[..., Any],
    ) -> None:
        super().__init__(step_id, metadata, input_type, output_type)
        self.func = func

    async def execute(self, input_data: InputType, context: WorkflowContext) -> OutputType:
        result = self.func(input_data, context)
        if asyncio.iscoroutine(result):
            result = await result

        if isinstance(result, self.output_type):
            return result
        if isinstance(result, BaseModel):
            return self.output_type(**result.model_dump())
        if isinstance(result, dict):
            return self.output_type(**result)
        return self.output_type(result=result)


__all__ = ["FunctionStep"]
