# Copyright 2026 Slymer-Tech
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
PyTree engine configuration and logic for Node.
"""

import types
from typing import Any, Iterable, cast
from types import MappingProxyType
from slyme.utils.pytree import (
    PyTreeEngine,
    PYTREE_ENGINE_REGISTRY,
    PyTreeAux,
    AttributeKey,
    MappingKey,
)
from slyme.utils.pytree.common import flatten_mapping_proxy, unflatten_mapping_proxy
from .core import (
    NodeDef,
    ExpressionDef,
    WrapperDef,
    NodeExec,
    ExpressionExec,
    WrapperExec,
    AsyncNodeDef,
    AsyncExpressionDef,
    AsyncWrapperDef,
    AsyncNodeExec,
    AsyncExpressionExec,
    AsyncWrapperExec,
)

# 1. Standard Engine: Preserves Types (Def -> Def, List -> List)
# Used for inspection, rendering, and validation.
NODE_ENGINE = PyTreeEngine("node_engine")
PYTREE_ENGINE_REGISTRY.register(NODE_ENGINE, key="node_engine")
# 2. Prepare Engine: Transforms Types (Def -> Exec, List -> Tuple, Dict -> MappingProxy)
# Used for compiling the definition tree into an execution tree.
NODE_PREPARE_ENGINE = PyTreeEngine("node_prepare", register_defaults=False)
PYTREE_ENGINE_REGISTRY.register(NODE_PREPARE_ENGINE, key="node_prepare")


# NodeDef Logic
def _flatten_node_def(obj: NodeDef) -> tuple[Iterable[Any], PyTreeAux]:
    """
    Flatten NodeDef into children (wrappers + kwargs values) and metadata.
    """
    # 1. Wrappers (List[WrapperDef]) treated as a single child container
    children = [obj.wrappers]
    rich_keys = [AttributeKey("wrappers")]

    # 2. Kwargs
    for k, v in obj._kwargs.items():
        children.append(v)
        rich_keys.append(MappingKey(k))

    metadata = {"func": obj._func, "specs": obj._specs}
    return tuple(children), PyTreeAux(
        children_keys=tuple(rich_keys), metadata=metadata, cls=NodeDef
    )


def _unflatten_node_def(children: Iterable[Any], aux: PyTreeAux) -> NodeDef:
    """
    Reconstruct NodeDef (Standard Engine).
    """
    # Use zip for cleaner iteration
    iterator = zip(aux.children_keys, children)
    # 1. Wrappers (Always the first element based on flatten logic)
    _, wrappers = next(iterator)
    # 2. Kwargs (Remaining elements)
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return NodeDef(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


def _unflatten_node_def_to_exec(children: Iterable[Any], aux: PyTreeAux) -> NodeExec:
    """
    Transform NodeDef into NodeExec (Prepare Engine).
    """
    iterator = zip(aux.children_keys, children)
    # 1. Wrappers
    # NOTE: The engine has already recursively transformed the wrappers list into a tuple.
    _, wrappers = next(iterator)
    # 2. Kwargs
    # NOTE: NodeExec.__init__ is responsible for converting this dict to MappingProxy.
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return NodeExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


# ExpressionDef Logic
def _flatten_expression_def(
    obj: ExpressionDef,
) -> tuple[Iterable[Any], PyTreeAux]:
    """
    Flatten ExpressionDef.
    """
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=ExpressionDef
    )


def _unflatten_expression_def(children: Iterable[Any], aux: PyTreeAux) -> ExpressionDef:
    """
    Reconstruct ExpressionDef (Standard Engine).
    """
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return ExpressionDef(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _unflatten_expression_def_to_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> ExpressionExec:
    """
    Transform ExpressionDef into ExpressionExec (Prepare Engine).
    """
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return ExpressionExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        kwargs=kwargs,
    )


# WrapperDef Logic
def _flatten_wrapper_def(obj: WrapperDef) -> tuple[Iterable[Any], PyTreeAux]:
    """
    Flatten WrapperDef.
    """
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=WrapperDef
    )


def _unflatten_wrapper_def(children: Iterable[Any], aux: PyTreeAux) -> WrapperDef:
    """
    Reconstruct WrapperDef (Standard Engine).
    """
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return WrapperDef(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _unflatten_wrapper_def_to_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> WrapperExec:
    """
    Transform WrapperDef into WrapperExec (Prepare Engine).
    """
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return WrapperExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        kwargs=kwargs,
    )


# Exec Types Logic
def _flatten_node_exec(obj: NodeExec) -> tuple[Iterable[Any], PyTreeAux]:
    children = [obj.wrappers]
    rich_keys = [AttributeKey("wrappers")]
    for k, v in obj._kwargs.items():
        children.append(v)
        rich_keys.append(MappingKey(k))
    metadata = {"func": obj._func, "specs": obj._specs}
    return tuple(children), PyTreeAux(
        children_keys=tuple(rich_keys), metadata=metadata, cls=NodeExec
    )


def _unflatten_node_exec(children: Iterable[Any], aux: PyTreeAux) -> NodeExec:
    iterator = zip(aux.children_keys, children)
    _, wrappers = next(iterator)
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return NodeExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


def _flatten_expression_exec(
    obj: ExpressionExec,
) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=ExpressionExec
    )


def _unflatten_expression_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> ExpressionExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return ExpressionExec(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _flatten_wrapper_exec(obj: WrapperExec) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=WrapperExec
    )


def _unflatten_wrapper_exec(children: Iterable[Any], aux: PyTreeAux) -> WrapperExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return WrapperExec(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


# AsyncNodeDef Logic
def _flatten_async_node_def(obj: AsyncNodeDef) -> tuple[Iterable[Any], PyTreeAux]:
    children = [obj.wrappers]
    rich_keys = [AttributeKey("wrappers")]
    for k, v in obj._kwargs.items():
        children.append(v)
        rich_keys.append(MappingKey(k))
    metadata = {"func": obj._func, "specs": obj._specs}
    return tuple(children), PyTreeAux(
        children_keys=tuple(rich_keys), metadata=metadata, cls=AsyncNodeDef
    )


def _unflatten_async_node_def(children: Iterable[Any], aux: PyTreeAux) -> AsyncNodeDef:
    iterator = zip(aux.children_keys, children)
    _, wrappers = next(iterator)
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return AsyncNodeDef(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


def _unflatten_async_node_def_to_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncNodeExec:
    iterator = zip(aux.children_keys, children)
    _, wrappers = next(iterator)
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return AsyncNodeExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


# AsyncExpressionDef Logic
def _flatten_async_expression_def(
    obj: AsyncExpressionDef,
) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=AsyncExpressionDef
    )


def _unflatten_async_expression_def(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncExpressionDef:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncExpressionDef(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _unflatten_async_expression_def_to_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncExpressionExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncExpressionExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        kwargs=kwargs,
    )


# AsyncWrapperDef Logic
def _flatten_async_wrapper_def(obj: AsyncWrapperDef) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=AsyncWrapperDef
    )


def _unflatten_async_wrapper_def(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncWrapperDef:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncWrapperDef(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _unflatten_async_wrapper_def_to_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncWrapperExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncWrapperExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        kwargs=kwargs,
    )


# AsyncExec Types Logic
def _flatten_async_node_exec(obj: AsyncNodeExec) -> tuple[Iterable[Any], PyTreeAux]:
    children = [obj.wrappers]
    rich_keys = [AttributeKey("wrappers")]
    for k, v in obj._kwargs.items():
        children.append(v)
        rich_keys.append(MappingKey(k))
    metadata = {"func": obj._func, "specs": obj._specs}
    return tuple(children), PyTreeAux(
        children_keys=tuple(rich_keys), metadata=metadata, cls=AsyncNodeExec
    )


def _unflatten_async_node_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncNodeExec:
    iterator = zip(aux.children_keys, children)
    _, wrappers = next(iterator)
    kwargs = {cast("MappingKey", k).key: v for k, v in iterator}
    return AsyncNodeExec(
        func=aux.metadata["func"],
        specs=aux.metadata["specs"],
        wrappers=wrappers,
        kwargs=kwargs,
    )


def _flatten_async_expression_exec(
    obj: AsyncExpressionExec,
) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=AsyncExpressionExec
    )


def _unflatten_async_expression_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncExpressionExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncExpressionExec(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


def _flatten_async_wrapper_exec(
    obj: AsyncWrapperExec,
) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(obj._kwargs.keys())
    children = tuple(obj._kwargs.values())
    rich_keys = tuple(MappingKey(k) for k in keys)
    metadata = {"func": obj._func, "specs": obj._specs}
    return children, PyTreeAux(
        children_keys=rich_keys, metadata=metadata, cls=AsyncWrapperExec
    )


def _unflatten_async_wrapper_exec(
    children: Iterable[Any], aux: PyTreeAux
) -> AsyncWrapperExec:
    kwargs = {cast("MappingKey", k).key: v for k, v in zip(aux.children_keys, children)}
    return AsyncWrapperExec(
        func=aux.metadata["func"], specs=aux.metadata["specs"], kwargs=kwargs
    )


# Container Logic (Prepare Transformations)
def _flatten_list(l: list) -> tuple[Iterable[Any], PyTreeAux]:
    return iter(l), PyTreeAux()


def _unflatten_to_tuple(children: Iterable[Any], _: PyTreeAux) -> tuple:
    return tuple(children)


def _flatten_tuple(t: tuple) -> tuple[Iterable[Any], PyTreeAux]:
    return iter(t), PyTreeAux()


def _unflatten_tuple(children: Iterable[Any], _: PyTreeAux) -> tuple:
    return tuple(children)


def _flatten_dict(d: dict) -> tuple[Iterable[Any], PyTreeAux]:
    keys = tuple(d.keys())
    rich_keys = tuple(MappingKey(k) for k in keys)
    children = (d[k] for k in keys)
    return children, PyTreeAux(children_keys=rich_keys)


def _unflatten_to_mapping_proxy(
    children: Iterable[Any], aux: PyTreeAux
) -> types.MappingProxyType:
    raw_keys = [k.key for k in cast("Iterable[MappingKey]", aux.children_keys)]
    return types.MappingProxyType(dict(zip(raw_keys, children)))


# Registrations
# --- NODE_PYTREE_ENGINE (Def -> Def, Exec -> Exec) ---
NODE_ENGINE.register(NodeDef, _flatten_node_def, _unflatten_node_def, strict=True)
NODE_ENGINE.register(
    ExpressionDef, _flatten_expression_def, _unflatten_expression_def, strict=True
)
NODE_ENGINE.register(
    WrapperDef, _flatten_wrapper_def, _unflatten_wrapper_def, strict=True
)
NODE_ENGINE.register(NodeExec, _flatten_node_exec, _unflatten_node_exec, strict=True)
NODE_ENGINE.register(
    ExpressionExec,
    _flatten_expression_exec,
    _unflatten_expression_exec,
    strict=True,
)
NODE_ENGINE.register(
    WrapperExec, _flatten_wrapper_exec, _unflatten_wrapper_exec, strict=True
)
NODE_ENGINE.register(
    AsyncNodeDef, _flatten_async_node_def, _unflatten_async_node_def, strict=True
)
NODE_ENGINE.register(
    AsyncExpressionDef,
    _flatten_async_expression_def,
    _unflatten_async_expression_def,
    strict=True,
)
NODE_ENGINE.register(
    AsyncWrapperDef,
    _flatten_async_wrapper_def,
    _unflatten_async_wrapper_def,
    strict=True,
)
NODE_ENGINE.register(
    AsyncNodeExec, _flatten_async_node_exec, _unflatten_async_node_exec, strict=True
)
NODE_ENGINE.register(
    AsyncExpressionExec,
    _flatten_async_expression_exec,
    _unflatten_async_expression_exec,
    strict=True,
)
NODE_ENGINE.register(
    AsyncWrapperExec,
    _flatten_async_wrapper_exec,
    _unflatten_async_wrapper_exec,
    strict=True,
)
NODE_ENGINE.register(MappingProxyType, flatten_mapping_proxy, unflatten_mapping_proxy)

# --- NODE_PREPARE_PYTREE_ENGINE (Def -> Exec, Mutables -> Immutables) ---
# Custom Containers
NODE_PREPARE_ENGINE.register(list, _flatten_list, _unflatten_to_tuple)
NODE_PREPARE_ENGINE.register(tuple, _flatten_tuple, _unflatten_tuple)
NODE_PREPARE_ENGINE.register(dict, _flatten_dict, _unflatten_to_mapping_proxy)
# Def -> Exec Transformations
NODE_PREPARE_ENGINE.register(
    NodeDef, _flatten_node_def, _unflatten_node_def_to_exec, strict=True
)
NODE_PREPARE_ENGINE.register(
    ExpressionDef,
    _flatten_expression_def,
    _unflatten_expression_def_to_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    WrapperDef,
    _flatten_wrapper_def,
    _unflatten_wrapper_def_to_exec,
    strict=True,
)
# Exec Identity (Exec -> Exec)
NODE_PREPARE_ENGINE.register(
    NodeExec, _flatten_node_exec, _unflatten_node_exec, strict=True
)
NODE_PREPARE_ENGINE.register(
    ExpressionExec,
    _flatten_expression_exec,
    _unflatten_expression_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    WrapperExec, _flatten_wrapper_exec, _unflatten_wrapper_exec, strict=True
)
NODE_PREPARE_ENGINE.register(
    AsyncNodeDef,
    _flatten_async_node_def,
    _unflatten_async_node_def_to_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    AsyncExpressionDef,
    _flatten_async_expression_def,
    _unflatten_async_expression_def_to_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    AsyncWrapperDef,
    _flatten_async_wrapper_def,
    _unflatten_async_wrapper_def_to_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    AsyncNodeExec, _flatten_async_node_exec, _unflatten_async_node_exec, strict=True
)
NODE_PREPARE_ENGINE.register(
    AsyncExpressionExec,
    _flatten_async_expression_exec,
    _unflatten_async_expression_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    AsyncWrapperExec,
    _flatten_async_wrapper_exec,
    _unflatten_async_wrapper_exec,
    strict=True,
)
NODE_PREPARE_ENGINE.register(
    MappingProxyType, flatten_mapping_proxy, unflatten_mapping_proxy
)
