"""Echo step: returns the incoming message with optional prefix/suffix and delay."""

from __future__ import annotations

import asyncio
from typing import Any, Type

from pydantic import BaseModel

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


class EchoInput(BaseModel):
    message: str | None = None
    result: str | None = None
    text: str | None = None
    content: str | None = None
    data: Any | None = None


class EchoOutput(BaseModel):
    result: str


class EchoStep(BaseStep[InputType, OutputType]):
    """Simple echo with optional prefix/suffix and artificial delay."""

    def __init__(
        self,
        step_id: str,
        metadata: StepMetadata,
        input_type: Type[InputType] = EchoInput,
        output_type: Type[OutputType] = EchoOutput,
        prefix: str = "",
        suffix: str = "",
        delay_seconds: float = 0.0,
    ) -> None:
        super().__init__(step_id, metadata, input_type, output_type)
        self.prefix = prefix
        self.suffix = suffix
        self.delay_seconds = delay_seconds

    async def execute(self, input_data: InputType, context: WorkflowContext) -> OutputType:
        if self.delay_seconds > 0:
            await asyncio.sleep(self.delay_seconds)

        message = None
        for field_name in ("message", "result", "text", "content"):
            value = getattr(input_data, field_name, None)
            if value is not None:
                message = value
                break
        if message is None:
            message = str(getattr(input_data, "data", ""))

        output_text = f"{self.prefix}{message}{self.suffix}"
        context.set(f"{self.step_id}_output", output_text)
        return self.output_type(result=output_text)


__all__ = ["EchoStep", "EchoInput", "EchoOutput"]
