from __future__ import annotations

from typing import Any

from flownet.core.exceptions import FlowDefinitionError
from flownet.core.request import FlowRequest, FlowStack, TransformationFrame
from flownet.core.types import PayloadData
from flownet.engine.context import OperatorContext
from flownet.engine.utils.flow_stack import append_frame

from .base import OperatorAction, OperatorHandler


class RouterOperatorHandler(OperatorHandler):
    """Internal router operator that emits continuation FlowRequests."""

    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        meta = ctx.current_frame_meta or {}
        routes = meta.get("routes") or meta.get("continuations")
        if routes is None:
            raise FlowDefinitionError("Router operator requires frame meta 'routes'.")
        if not isinstance(routes, (list, tuple)):
            raise TypeError("Router operator routes must be a list.")

        branch_count = len(routes)
        await ctx.request_add_delta(request, branch_count - 1)

        for route in routes:
            target_address, next_request = self._build_route_request(
                ctx, request, transformation_frame, route
            )
            await ctx.emit_flow_request(target_address, next_request)

        return OperatorAction.RETURN

    def _build_route_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
        route: Any,
    ) -> tuple[str, FlowRequest]:
        if isinstance(route, FlowRequest):
            if route.request_id != request.request_id:
                raise FlowDefinitionError("Router continuation must keep request_id.")
            if route.origin_address != request.origin_address:
                raise FlowDefinitionError("Router continuation must keep origin_address.")
            if route.ack_route != request.ack_route:
                route = FlowRequest(
                    flow_stack=route.flow_stack,
                    request_id=route.request_id,
                    origin_address=route.origin_address,
                    ack_route=request.ack_route,
                )
            return ctx.local_address(), route

        if not isinstance(route, dict):
            raise TypeError("Router route must be a dict or FlowRequest.")

        address = route.get("address") or ctx.local_address()
        flow_stack = route.get("flow_stack")
        payload = route.get("payload", transformation_frame.payload)
        meta = route.get("meta")
        if meta is None:
            meta = self._clone_meta(getattr(transformation_frame, "meta", None))
        else:
            meta = self._clone_meta(meta)

        if isinstance(flow_stack, FlowStack):
            return address, FlowRequest(
                flow_stack=flow_stack,
                request_id=request.request_id,
                origin_address=request.origin_address,
                ack_route=request.ack_route,
            )

        next_transformation = route.get("transformation")
        if next_transformation is None:
            downstreams = transformation_frame.transformation.downstreams
            if len(downstreams) != 1:
                raise FlowDefinitionError(
                    "Router route missing transformation and has multiple downstreams."
                )
            next_transformation = downstreams[0]

        if not isinstance(payload, PayloadData):
            payload = PayloadData(args=[payload], kwargs={})

        base_stack = FlowStack(request.flow_stack.stack[:-1])
        new_frame = TransformationFrame(
            transformation=next_transformation,
            payload=payload,
            meta=meta,
        )
        new_request = FlowRequest(
            flow_stack=append_frame(base_stack, new_frame),
            request_id=request.request_id,
            origin_address=request.origin_address,
            ack_route=request.ack_route,
        )
        return address, new_request

    @staticmethod
    def _clone_meta(value: Any) -> Any:
        if isinstance(value, dict):
            return {k: RouterOperatorHandler._clone_meta(v) for k, v in value.items()}
        if isinstance(value, list):
            return [RouterOperatorHandler._clone_meta(item) for item in value]
        return value
