"""Transform step: maps input fields to output fields with optional static values."""

from __future__ import annotations

from typing import Any, Type

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


class TransformStep(BaseStep[InputType, OutputType]):
    """Field-mapping transform between input and output Pydantic models."""

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

    async def execute(self, input_data: InputType, context: WorkflowContext) -> OutputType:
        input_dict = input_data.model_dump()
        output_schema = self.output_type.model_json_schema()
        result: dict[str, Any] = {}

        for output_field, mapping in self.mappings.items():
            target_schema = output_schema.get("properties", {}).get(output_field, {})
            if isinstance(mapping, str):
                if mapping.startswith("static:"):
                    raw_value = mapping.replace("static:", "", 1)
                else:
                    raw_value = input_dict.get(mapping)
            elif isinstance(mapping, dict):
                raw_value = {k: input_dict.get(v) for k, v in mapping.items()}
            else:
                raw_value = mapping

            result[output_field] = coerce_value_to_schema_type(raw_value, target_schema)

        return self.output_type(**result)


__all__ = ["TransformStep"]
