from __future__ import annotations

from typing import Any

from flownet.compiler.transformation import Transformation
from flownet.core.request import BaseFlowFrame, FlowRequest, FlowStack, TransformationFrame
from flownet.core.types import PayloadData


_UNSET: Any = object()


def clone_flow_stack(flow_stack: FlowStack) -> FlowStack:
    return FlowStack(list(flow_stack.stack))


def replace_top_frame(flow_stack: FlowStack, frame: BaseFlowFrame) -> FlowStack:
    stack = list(flow_stack.stack)
    if not stack:
        return FlowStack([frame])
    stack[-1] = frame
    return FlowStack(stack)


def append_frame(flow_stack: FlowStack, frame: BaseFlowFrame) -> FlowStack:
    return FlowStack(list(flow_stack.stack) + [frame])


def clone_transformation_frame(
    frame: TransformationFrame,
    *,
    transformation: Transformation | Any = _UNSET,
    payload: PayloadData | Any = _UNSET,
    meta: dict | None | Any = _UNSET,
) -> TransformationFrame:
    if transformation is _UNSET:
        transformation = frame.transformation
    if payload is _UNSET:
        payload = frame.payload
    if meta is _UNSET:
        meta = frame.meta
    return TransformationFrame(transformation, payload, meta=meta)


def with_frame_meta(frame: TransformationFrame, meta: dict | None) -> TransformationFrame:
    return clone_transformation_frame(frame, meta=meta)


def clone_request_with_frame_meta(
    request: FlowRequest, meta: dict | None
) -> FlowRequest:
    top = request.flow_stack.top()
    if not isinstance(top, TransformationFrame):
        raise TypeError(f"Unexpected flow frame: {type(top)}")
    new_frame = clone_transformation_frame(top, meta=meta)
    new_stack = replace_top_frame(request.flow_stack, new_frame)
    return FlowRequest(
        flow_stack=new_stack,
        request_id=request.request_id,
        origin_address=request.origin_address,
        ack_route=request.ack_route,
    )


__all__ = [
    "clone_flow_stack",
    "replace_top_frame",
    "append_frame",
    "clone_transformation_frame",
    "with_frame_meta",
    "clone_request_with_frame_meta",
]
