# workflow_engine/nodes/data.py
"""
Utility nodes to construct and deconstruct data objects.
"""

from collections.abc import Sequence
from functools import cached_property
from typing import ClassVar, Generic, Self, Type, TypeVar

from overrides import override
from pydantic import Field
from pydantic.fields import FieldInfo

from ..core import (
    Data,
    DataValue,
    Empty,
    ExecutionContext,
    IntegerValue,
    Node,
    NodeTypeInfo,
    Params,
    SequenceValue,
    StringMapValue,
    StringValue,
    ValidationContext,
    Value,
    ValueType,
)
from ..core.values import build_data_type, get_data_dict, get_data_field

V = TypeVar("V", bound=Value)


################################################################################
# Sequences


class SequenceParams(Params):
    length: IntegerValue = Field(
        title="Length", description="The number of elements in the sequence."
    )


class SequenceData(Data, Generic[V]):
    sequence: SequenceValue[V] = Field(
        title="Sequence",
        description="The sequence of values.",
    )

    @classmethod
    def empty(cls) -> Self:
        sequence_type = get_data_field(cls, "sequence")
        assert sequence_type is not None
        assert issubclass(sequence_type, SequenceValue)
        return cls(sequence=sequence_type(root=()))


class GatherSequenceNode(Node[Data, SequenceData, SequenceParams]):
    """
    Creates a new sequence object of a given length.
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="Gather Sequence",
        description="Creates a new sequence object of a given length.",
        version="0.4.0",
        parameter_type=SequenceParams,
    )

    # The type of the elements in the sequence.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this available at runtime
    element_type: ValueType = Field(default=Value, exclude=True)

    @cached_property
    def indices(self) -> Sequence[int]:
        return range(self.params.length.root)

    def key(self, index: int) -> str:
        return f"element_{index}"

    def title(self, index: int) -> str:
        """
        We display the indices as 1-based for readability.
        """
        return f"Item {index + 1}"

    def description(self, index: int) -> str:
        """
        We display the indices as 1-based for readability.
        """
        return f"The {index + 1} item in the sequence."

    @override
    async def dynamic_input_type(self, context: ValidationContext) -> Type[Data]:
        return build_data_type(
            name="GatherSequenceInput",
            fields={
                self.key(i): (
                    self.element_type,
                    FieldInfo(
                        title=self.title(i),
                        description=self.description(i),
                    ),
                )
                for i in self.indices
            },
        )

    @override
    async def dynamic_output_type(
        self, context: ValidationContext
    ) -> Type[SequenceData]:
        return SequenceData[self.element_type]

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[Data],
        output_type: Type[SequenceData],
        input: Data,
    ) -> SequenceData:
        input_dict = get_data_dict(input)
        return output_type(
            sequence=SequenceValue[self.element_type](
                root=[input_dict[self.key(i)] for i in self.indices]
            )
        )


class ExpandSequenceNode(Node[SequenceData, Data, SequenceParams]):
    """
    Extracts a sequence of elements to a data object.
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="Expand Sequence",
        description="Extracts a sequence of elements to a data object.",
        version="0.4.0",
        parameter_type=SequenceParams,
    )

    # The type of the element to extract.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this available at runtime
    element_type: ValueType = Value

    def key(self, index: int) -> str:
        return f"element_{index}"

    def title(self, index: int) -> str:
        """
        We display the indices as 1-based for readability.
        """
        return f"Item {index + 1}"

    def description(self, index: int) -> str:
        """
        We display the indices as 1-based for readability.
        """
        return f"The {index + 1} item in the sequence."

    @cached_property
    def indices(self) -> Sequence[int]:
        return range(self.params.length.root)

    @override
    async def dynamic_input_type(
        self, context: ValidationContext
    ) -> Type[SequenceData]:
        return SequenceData[self.element_type]

    @override
    async def dynamic_output_type(self, context: ValidationContext) -> Type[Data]:
        return build_data_type(
            name="ExpandSequenceOutput",
            fields={
                self.key(i): (
                    self.element_type,
                    FieldInfo(title=self.title(i), description=self.description(i)),
                )
                for i in self.indices
            },
        )

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[SequenceData],
        output_type: Type[Data],
        input: SequenceData,
    ) -> Data:
        N = self.params.length.root
        assert len(input.sequence) == N, (
            f"Expected sequence of length {N}, but got {len(input.sequence)}"
        )
        return output_type(**{self.key(i): input.sequence[i] for i in range(N)})


################################################################################
# Mappings


class MappingParams(Params):
    keys: SequenceValue[StringValue] = Field(
        title="Keys",
        description="The keys of the mapping.",
    )


class MappingData(Data, Generic[V]):
    mapping: StringMapValue[V] = Field(
        title="Mapping",
        description="The mapping of keys to values.",
    )


class GatherMappingNode(Node[Data, MappingData, MappingParams]):
    """
    Creates a new mapping object from the inputs to this node.

    Example:
        >>> node = GatherMappingNode.from_keys("node_id", ["a", "b", "c"])
        >>> node.run(context, input={"a": 1, "b": 2, "c": 3}).model_dump()
        {"mapping": {"a": 1, "b": 2, "c": 3}}
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="GatherMapping",
        description="Creates a new mapping object from the inputs to this node.",
        version="0.4.0",
        parameter_type=MappingParams,
    )

    # The type of the values in the mapping.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this serializable/deserializable
    value_type: ValueType = Value

    @override
    async def dynamic_input_type(self, context: ValidationContext) -> Type[Data]:
        return build_data_type(
            name="GatherMappingInput",
            fields={
                key.root: (
                    self.value_type,
                    FieldInfo(title=key.root),
                )
                for key in self.params.keys
            },
        )

    @override
    async def dynamic_output_type(
        self, context: ValidationContext
    ) -> Type[MappingData]:
        return MappingData[self.value_type]

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[Data],
        output_type: Type[MappingData],
        input: Data,
    ) -> MappingData:
        return output_type(
            mapping=StringMapValue[self.value_type](
                {key.root: getattr(input, key.root) for key in self.params.keys}
            )
        )


class ExpandMappingNode(Node[MappingData, Data, MappingParams]):
    """
    Extracts values from a mapping object at specific keys.

    Example:
        >>> node = ExpandMappingNode.from_keys("node_id", ["a", "b", "c"])
        >>> node.run(context, input={"mapping": {"a": 1, "b": 2, "c": 3}}).model_dump()
        {"a": 1, "b": 2, "c": 3}
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="Expand Mapping",
        description="Extracts values from a mapping object at specific keys.",
        version="0.4.0",
        parameter_type=MappingParams,
    )

    # The type of the values in the mapping.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this serializable/deserializable
    value_type: ValueType = Value

    @override
    async def dynamic_input_type(self, context: ValidationContext) -> Type[MappingData]:
        return MappingData[self.value_type]

    @override
    async def dynamic_output_type(self, context: ValidationContext) -> Type[Data]:
        return build_data_type(
            name="ExpandMappingOutput",
            fields={
                key.root: (
                    self.value_type,
                    FieldInfo(title=key.root),
                )
                for key in self.params.keys
            },
        )

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[MappingData],
        output_type: Type[Data],
        input: MappingData,
    ) -> Data:
        return output_type(**{key.root: input.mapping[key] for key in self.params.keys})


################################################################################
# Data nodes


D = TypeVar("D", bound=Data)


class NestedData(Data, Generic[D]):
    """
    A data type that contains a nested data object.
    """

    data: DataValue[D] = Field(title="Data", description="The nested data object.")


class GatherDataNode(Node[Data, NestedData, Empty]):
    """
    A node that gathers a data object from the inputs to this node.

    Example:
        >>> node = GatherDataNode.from_data_type("node_id", Data)
        >>> node.run(context, input={"a": 1, "b": 2, "c": 3}).model_dump()
        {"data": {"a": 1, "b": 2, "c": 3}}
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="GatherData",
        description="A node that gathers a data object from the inputs to this node.",
        version="0.4.0",
        parameter_type=Empty,
    )

    # The type of the element to extract.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this serializable/deserializable
    data_type: Type[Data] = Field(default=Data, exclude=True)

    @override
    async def dynamic_input_type(self, context: ValidationContext) -> Type[Data]:
        return self.data_type

    @override
    async def dynamic_output_type(self, context: ValidationContext) -> Type[NestedData]:
        return NestedData[self.data_type]

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[Data],
        output_type: Type[NestedData],
        input: Data,
    ) -> NestedData:
        return NestedData[self.data_type](data=DataValue[self.data_type](root=input))


class ExpandDataNode(Node[NestedData, Data, Empty]):
    """
    A node that expands a nested data object into its individual fields.

    Example:
        >>> node = ExpandDataNode.from_data_type("node_id", Data)
        >>> node.run(context, input={"data": {"a": 1, "b": 2, "c": 3}}).model_dump()
        {"a": 1, "b": 2, "c": 3}
    """

    TYPE_INFO: ClassVar[NodeTypeInfo] = NodeTypeInfo.from_parameter_type(
        display_name="ExpandData",
        description="A node that expands a nested data object into its individual fields.",
        version="0.4.0",
        parameter_type=Empty,
    )

    # The type of the nested data object.
    # For now, this field is only available when the node is constructed
    # programmatically.
    # TODO: make this serializable/deserializable
    data_type: Type[Data] = Field(default=Data, exclude=True)

    @override
    async def dynamic_input_type(self, context: ValidationContext) -> Type[NestedData]:
        return NestedData[self.data_type]

    @override
    async def dynamic_output_type(self, context: ValidationContext) -> Type[Data]:
        return self.data_type

    @override
    async def run(
        self,
        *,
        context: ExecutionContext,
        input_type: Type[NestedData],
        output_type: Type[Data],
        input: NestedData,
    ) -> Data:
        return input.data.root


__all__ = [
    "ExpandDataNode",
    "ExpandMappingNode",
    "ExpandSequenceNode",
    "GatherDataNode",
    "GatherMappingNode",
    "GatherSequenceNode",
]
