from __future__ import annotations

from collections.abc import AsyncIterator, Iterator
from typing import Any, Callable

from flownet.compiler.transformation import Transformation
from flownet.core.request import (
    FlowAckRoute,
    FlowRequest,
    FlowStack,
    TransformationFrame,
    coerce_flow_ack_route,
    resolve_flow_ack_coordinator_address,
    resolve_flow_ack_coordinator_target,
)
from flownet.core.types import PayloadData
from flownet.engine.result_normalizer import ResultNormalizer
from flownet.engine.utils.flow_stack import append_frame


class FlowAdvancer:
    """Handle flow advancement and request delta accounting."""

    def __init__(
        self,
        *,
        send_ack: Callable[[str, str, str, int], Any],
        submit_runtime_coro: Callable[[Any], Any],
        handle_flow_request: Callable[[FlowRequest], Any],
        result_normalizer: ResultNormalizer | None = None,
    ) -> None:
        self._send_ack = send_ack
        self._submit_runtime_coro = submit_runtime_coro
        self._handle_flow_request = handle_flow_request
        self._normalizer = result_normalizer or ResultNormalizer()

    async def post_process_result(
        self,
        request: FlowRequest,
        result: Any,
        current_payload: PayloadData,
    ) -> None:
        top_frame = request.flow_stack.top()
        if top_frame is None:
            await self._apply_request_delta(request, -1)
            return
        if not isinstance(top_frame, TransformationFrame):
            raise TypeError(f"Unexpected flow frame: {type(top_frame)}")

        if top_frame.transformation.is_sink:
            await self._apply_request_delta(request, -1)
            return

        transformation = top_frame.transformation
        (
            internal_downstreams,
            base_stack_internal,
            outer_downstreams,
            base_stack_outer,
        ) = self._resolve_downstream_bindings(request, transformation)
        propagated_meta = self._next_frame_meta(getattr(top_frame, "meta", None))

        if transformation.type == "flatmap":
            if isinstance(result, AsyncIterator):
                await self._dispatch_incremental_async(
                    request=request,
                    iterator=result,
                    internal_downstreams=internal_downstreams,
                    base_stack_internal=base_stack_internal,
                    outer_downstreams=outer_downstreams,
                    base_stack_outer=base_stack_outer,
                    propagated_meta=propagated_meta,
                )
                return
            if isinstance(result, Iterator):
                await self._dispatch_incremental_sync(
                    request=request,
                    iterator=result,
                    internal_downstreams=internal_downstreams,
                    base_stack_internal=base_stack_internal,
                    outer_downstreams=outer_downstreams,
                    base_stack_outer=base_stack_outer,
                    propagated_meta=propagated_meta,
                )
                return

        expand_iterable = transformation.type == "flatmap"
        payloads = self._normalizer.normalize(result, expand_iterable)
        children_per_payload = len(internal_downstreams) + len(outer_downstreams)
        total_children = children_per_payload * len(payloads)
        await self._apply_request_delta(request, total_children - 1)

        if total_children == 0:
            return

        for payload in payloads:
            self._dispatch_payload(
                request=request,
                payload=payload,
                internal_downstreams=internal_downstreams,
                base_stack_internal=base_stack_internal,
                outer_downstreams=outer_downstreams,
                base_stack_outer=base_stack_outer,
                propagated_meta=propagated_meta,
            )

    async def _apply_request_delta(self, request: FlowRequest, delta: int) -> None:
        if delta == 0:
            return
        ack_route = self.resolve_ack_route(request)
        if ack_route is None:
            return
        coordinator_target = self.resolve_ack_coordinator_target(request)
        if coordinator_target is None:
            return
        coordinator_address = self.resolve_ack_coordinator_address(request)
        await self._send_ack(
            coordinator_address or request.origin_address,
            coordinator_target,
            request.request_id,
            delta,
        )

    async def apply_request_delta(self, request: FlowRequest, delta: int) -> None:
        await self._apply_request_delta(request, delta)

    @staticmethod
    def resolve_ack_route(request: FlowRequest) -> FlowAckRoute | None:
        return coerce_flow_ack_route(getattr(request, "ack_route", None))

    @staticmethod
    def resolve_ack_coordinator_target(request: FlowRequest) -> str | None:
        ack_route = FlowAdvancer.resolve_ack_route(request)
        return resolve_flow_ack_coordinator_target(ack_route)

    @staticmethod
    def resolve_ack_coordinator_address(request: FlowRequest) -> str | None:
        ack_route = FlowAdvancer.resolve_ack_route(request)
        return resolve_flow_ack_coordinator_address(ack_route)

    @staticmethod
    def _next_frame_meta(frame_meta: Any) -> dict[str, Any] | None:
        if not isinstance(frame_meta, dict):
            return None
        propagated: dict[str, Any] = {}
        next_meta = frame_meta.get("next_meta")
        if isinstance(next_meta, dict) and next_meta:
            propagated["next_meta"] = FlowAdvancer._clone_meta(next_meta)
        # Keep replay/checkpoint boundary metadata available to downstream stages.
        sys_meta = frame_meta.get("sys")
        if isinstance(sys_meta, dict) and sys_meta:
            propagated["sys"] = FlowAdvancer._clone_meta(sys_meta)
        return propagated or None

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

    @staticmethod
    def _resolve_downstream_bindings(
        request: FlowRequest,
        transformation: Transformation,
    ) -> tuple[
        list[Transformation],
        FlowStack,
        list[Transformation],
        FlowStack | None,
    ]:
        internal_downstreams = transformation.downstreams
        base_stack_internal = FlowStack(request.flow_stack.stack[:-1])

        outer_downstreams: list[Transformation] = []
        base_stack_outer: FlowStack | None = None
        if transformation.is_return:
            stack = request.flow_stack.stack
            idx = len(stack) - 2
            while idx >= 0:
                frame = stack[idx]
                if isinstance(frame, TransformationFrame) and frame.transformation.is_return:
                    idx -= 1
                    continue
                break
            if idx >= 0:
                frame = stack[idx]
                if isinstance(frame, TransformationFrame):
                    outer_downstreams = frame.transformation.downstreams
                    base_stack_outer = FlowStack(stack[:idx])
        return internal_downstreams, base_stack_internal, outer_downstreams, base_stack_outer

    async def _dispatch_incremental_sync(
        self,
        *,
        request: FlowRequest,
        iterator: Iterator[Any],
        internal_downstreams: list[Transformation],
        base_stack_internal: FlowStack,
        outer_downstreams: list[Transformation],
        base_stack_outer: FlowStack | None,
        propagated_meta: dict[str, Any] | None,
    ) -> None:
        children_per_payload = len(internal_downstreams) + len(outer_downstreams)
        try:
            for item in iterator:
                payload = self._item_to_payload(item)
                if children_per_payload > 0:
                    await self._apply_request_delta(request, children_per_payload)
                self._dispatch_payload(
                    request=request,
                    payload=payload,
                    internal_downstreams=internal_downstreams,
                    base_stack_internal=base_stack_internal,
                    outer_downstreams=outer_downstreams,
                    base_stack_outer=base_stack_outer,
                    propagated_meta=propagated_meta,
                )
        except Exception:
            # Keep partial emissions valid; current frame completion is handled
            # by the exception resolver via post_process_result(None, ...).
            raise
        await self._apply_request_delta(request, -1)

    async def _dispatch_incremental_async(
        self,
        *,
        request: FlowRequest,
        iterator: AsyncIterator[Any],
        internal_downstreams: list[Transformation],
        base_stack_internal: FlowStack,
        outer_downstreams: list[Transformation],
        base_stack_outer: FlowStack | None,
        propagated_meta: dict[str, Any] | None,
    ) -> None:
        children_per_payload = len(internal_downstreams) + len(outer_downstreams)
        try:
            async for item in iterator:
                payload = self._item_to_payload(item)
                if children_per_payload > 0:
                    await self._apply_request_delta(request, children_per_payload)
                self._dispatch_payload(
                    request=request,
                    payload=payload,
                    internal_downstreams=internal_downstreams,
                    base_stack_internal=base_stack_internal,
                    outer_downstreams=outer_downstreams,
                    base_stack_outer=base_stack_outer,
                    propagated_meta=propagated_meta,
                )
        except Exception:
            # Keep partial emissions valid; current frame completion is handled
            # by the exception resolver via post_process_result(None, ...).
            raise
        await self._apply_request_delta(request, -1)

    def _dispatch_payload(
        self,
        *,
        request: FlowRequest,
        payload: PayloadData,
        internal_downstreams: list[Transformation],
        base_stack_internal: FlowStack,
        outer_downstreams: list[Transformation],
        base_stack_outer: FlowStack | None,
        propagated_meta: dict[str, Any] | None,
    ) -> None:
        for downstream_transformation in internal_downstreams:
            new_frame = TransformationFrame(
                transformation=downstream_transformation,
                payload=payload,
                meta=self._clone_meta(propagated_meta),
            )

            new_request = FlowRequest(
                flow_stack=append_frame(base_stack_internal, new_frame),
                request_id=request.request_id,
                origin_address=request.origin_address,
                ack_route=request.ack_route,
            )
            self._submit_runtime_coro(self._handle_flow_request(new_request))

        if base_stack_outer is None:
            return

        for downstream_transformation in outer_downstreams:
            new_frame = TransformationFrame(
                transformation=downstream_transformation,
                payload=payload,
                meta=self._clone_meta(propagated_meta),
            )

            new_request = FlowRequest(
                flow_stack=append_frame(base_stack_outer, new_frame),
                request_id=request.request_id,
                origin_address=request.origin_address,
                ack_route=request.ack_route,
            )
            self._submit_runtime_coro(self._handle_flow_request(new_request))

    @staticmethod
    def _item_to_payload(item: Any) -> PayloadData:
        if isinstance(item, PayloadData):
            return item
        return PayloadData(args=[item], kwargs={})


__all__ = ["FlowAdvancer"]
