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.meta_validation import require_dict, require_positive_int, validate_partition_id
from flownet.engine.utils.flow_stack import append_frame
from flownet.runtime.flows.replay_checkpoint import parse_replay_checkpoint

from ..base import OperatorAction, OperatorHandler
from .common import (
    _LifecyclePhase,
    _StatefulOperatorMixin,
    _base_by_key_mode,
    _checkpoint_boundary_token,
    _extract_group_key,
    _extract_group_sum_value,
    _is_by_key_mode,
    _parse_non_negative_int,
    _payload_value,
    _release_count_from_pending,
    _sorted_keyed_items,
)


class _ReducerState:
    def __init__(
        self,
        mode: str,
        num_partitions: int,
        *,
        max_reorder_gap: int = 64,
        groupby: dict[str, Any] | None = None,
    ) -> None:
        self.mode = mode
        self._by_key_mode = _is_by_key_mode(mode)
        self._base_mode = _base_by_key_mode(mode) if self._by_key_mode else mode
        self._groupby = dict(groupby or {})
        self.phase = _LifecyclePhase.OPEN
        self.num_partitions = num_partitions
        self.max_reorder_gap = max_reorder_gap
        self.pending_count = 0
        self.done_partitions: set[int] = set()
        self.source_final_seq: int | None = None
        self._next_source_seq = 1
        self._future_source_payloads: dict[int, Any] = {}
        if self._by_key_mode:
            self.acc: dict[Any, Any] = {}
        elif mode == "count":
            self.acc = 0
        elif mode == "sum":
            self.acc = 0
        else:
            self.acc: list[Any] = []

    def add(
        self,
        partition_id: int,
        payload: PayloadData,
        source_seq: int | None = None,
    ) -> tuple[bool, int, int]:
        if self.phase == _LifecyclePhase.FINALIZED:
            return False, 0, 0
        if partition_id in self.done_partitions and not self._allow_post_done_data():
            return False, 0, 0
        value = _payload_value(payload)
        if source_seq is None:
            self.pending_count += 1
            self._apply_value(value)
            return True, 0, 0

        accepted, missing_seq, recovered = self._add_source_sequenced(value, source_seq)
        if not accepted:
            return False, missing_seq, recovered
        self.pending_count += 1
        return True, missing_seq, recovered

    def mark_done(self, partition_id: int, source_final_seq: int | None = None) -> tuple[bool, bool]:
        if self.phase == _LifecyclePhase.FINALIZED:
            return False, False
        if partition_id in self.done_partitions:
            return False, False
        self.pending_count += 1
        self.done_partitions.add(partition_id)
        if source_final_seq is not None:
            if self.source_final_seq is None or source_final_seq > self.source_final_seq:
                self.source_final_seq = source_final_seq
        entered_draining = False
        if self.phase == _LifecyclePhase.OPEN:
            self.phase = _LifecyclePhase.DRAINING
            entered_draining = True
        return True, entered_draining

    def is_complete(self) -> bool:
        if len(self.done_partitions) < self.num_partitions:
            return False
        if self.source_final_seq is None:
            return True
        return self._contiguous_source_seq() >= self.source_final_seq

    def finalize(self) -> Any:
        self.phase = _LifecyclePhase.FINALIZED
        if self._by_key_mode:
            return _sorted_keyed_items(self.acc)
        return self.acc

    def partial(self) -> Any:
        if not self._by_key_mode:
            return self.acc
        if self._base_mode == "concat":
            return {k: list(v) for k, v in self.acc.items()}
        return dict(self.acc)

    def _allow_post_done_data(self) -> bool:
        if self.source_final_seq is None:
            return False
        return self._contiguous_source_seq() < self.source_final_seq

    def _apply_value(self, value: Any) -> None:
        if self._by_key_mode:
            key = _extract_group_key(value, self._groupby)
            if self._base_mode == "count":
                self.acc[key] = int(self.acc.get(key, 0)) + 1
            elif self._base_mode == "sum":
                amount = _extract_group_sum_value(value, self._groupby)
                self.acc[key] = self.acc.get(key, 0) + amount
            else:
                bucket = self.acc.setdefault(key, [])
                bucket.append(value)
            return

        if self.mode == "count":
            self.acc += 1
        elif self.mode == "sum":
            self.acc += value
        else:
            self.acc.append(value)

    def _add_source_sequenced(self, value: Any, source_seq: int) -> tuple[bool, int, int]:
        if source_seq < self._next_source_seq:
            return False, 0, 0
        if source_seq in self._future_source_payloads:
            return False, 0, 0

        missing_seq, recovered = self._advance_reorder_window(source_seq)
        if source_seq < self._next_source_seq:
            return False, missing_seq, recovered

        if source_seq > self._next_source_seq:
            self._future_source_payloads[source_seq] = value
            return True, missing_seq, recovered

        self._apply_value(value)
        self._next_source_seq += 1
        recovered += self._drain_buffered_contiguous()
        return True, missing_seq, recovered

    def _advance_reorder_window(self, target_seq: int) -> tuple[int, int]:
        missing_seq = 0
        recovered = 0
        while (target_seq - self._next_source_seq) > self.max_reorder_gap:
            buffered = self._future_source_payloads.pop(self._next_source_seq, None)
            if buffered is None:
                missing_seq += 1
            else:
                self._apply_value(buffered)
                recovered += 1
            self._next_source_seq += 1
        recovered += self._drain_buffered_contiguous()
        return missing_seq, recovered

    def _drain_buffered_contiguous(self) -> int:
        recovered = 0
        while self._next_source_seq in self._future_source_payloads:
            self._apply_value(self._future_source_payloads.pop(self._next_source_seq))
            self._next_source_seq += 1
            recovered += 1
        return recovered

    def _contiguous_source_seq(self) -> int:
        return self._next_source_seq - 1



class ReducerOperatorHandler(_StatefulOperatorMixin, OperatorHandler):
    """Internal reducer that waits for partition completion signals.

    This handler runs inside whichever runtime receives the continuation
    FlowRequest (often a container runtime). There is no explicit "spawn":
    the remote runtime simply executes this operator when the flowrequest
    arrives.
    """

    def __init__(
        self,
        *,
        max_active_states: int = 10000,
        max_state_idle_ops: int = 10000,
        state_cleanup_interval: int = 256,
    ) -> None:
        self._states: dict[tuple[Any, ...], _ReducerState] = {}
        self._init_stateful_runtime(
            max_active_states=max_active_states,
            max_state_idle_ops=max_state_idle_ops,
            state_cleanup_interval=state_cleanup_interval,
        )

    def stats_snapshot(self) -> dict[str, int]:
        snapshot = self._base_stats_snapshot()
        if self._states:
            snapshot["state_pending_max"] = max(state.pending_count for state in self._states.values())
            snapshot["state_done_partitions_max"] = max(
                len(state.done_partitions) for state in self._states.values()
            )
            snapshot["state_reorder_buffer_max"] = max(
                len(state._future_source_payloads) for state in self._states.values()
            )
            snapshot["state_acc_keys_max"] = max(
                len(state.acc) if isinstance(state.acc, dict) else 0
                for state in self._states.values()
            )
        return snapshot

    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        self._tick_and_maybe_cleanup()
        meta = ctx.current_frame_meta or {}
        reducer_meta = require_dict(meta.get("reducer"), "reducer")

        stage_id = reducer_meta.get("stage_id") or transformation_frame.transformation.trans_id
        num_partitions = require_positive_int(
            reducer_meta.get("num_partitions", 1), "reducer.num_partitions"
        )
        max_reorder_gap = require_positive_int(
            reducer_meta.get("max_reorder_gap", 64), "reducer.max_reorder_gap"
        )
        mode = reducer_meta.get("mode", "concat")
        groupby_meta = reducer_meta.get("groupby")
        if groupby_meta is not None and not isinstance(groupby_meta, dict):
            raise FlowDefinitionError("reducer.groupby must be a dict when provided.")
        partition_id = validate_partition_id(
            meta.get("partition_id", 0), num_partitions, field="partition_id"
        )
        control = meta.get("control", "data")
        if control not in {"data", "done"}:
            raise FlowDefinitionError("Reducer control must be one of: data, done.")
        # TODO(sequence-watermark):
        # `source_seq` / `source_final_seq` are used for v1 watermark-close.
        # This slice only covers reducer close convergence; bounded reorder
        # buffering is tracked as a follow-up task.
        source_seq = _parse_non_negative_int(meta.get("source_seq"))
        source_final_seq = _parse_non_negative_int(meta.get("source_final_seq"))
        merge_address = reducer_meta.get("merge_address")
        merge_transformation = reducer_meta.get("merge_transformation")
        if merge_address and merge_transformation is None:
            raise FlowDefinitionError("merge_address requires merge_transformation.")

        is_partition_reducer = merge_address and merge_transformation is not None
        logical_key = (
            (request.request_id, stage_id, partition_id)
            if is_partition_reducer
            else (request.request_id, stage_id)
        )
        checkpoint = parse_replay_checkpoint(
            meta,
            request_id=request.request_id,
            stage_id=str(stage_id),
        )
        if checkpoint.is_replay:
            self._stats.inc("replay_event_seen")
        boundary_token = _checkpoint_boundary_token(checkpoint)
        epoch = checkpoint.epoch
        latest_boundary = self._latest_boundary.get(logical_key)
        latest_epoch = latest_boundary[0] if latest_boundary is not None else None
        if latest_epoch is not None and epoch < latest_epoch:
            self._stats.inc("stale_epoch_dropped")
            return OperatorAction.RETURN
        if latest_boundary is not None and epoch == latest_epoch and boundary_token[1] != latest_boundary[1]:
            self._stats.inc("mixed_plan_boundary_dropped")
            return OperatorAction.RETURN
        if checkpoint.is_replay and latest_epoch is not None and epoch > latest_epoch:
            if checkpoint.checkpoint_epoch is not None and checkpoint.checkpoint_epoch != latest_epoch:
                self._stats.inc("invalid_replay_boundary_dropped")
                return OperatorAction.RETURN
        if latest_boundary is None or epoch > latest_epoch:
            if checkpoint.is_replay:
                self._stats.inc("replay_epoch_activated")
            self._latest_boundary.set(logical_key, boundary_token)
            self._states.pop(logical_key, None)
            self._discard_state_key(logical_key)

        completion_key = (*logical_key, checkpoint.epoch, checkpoint.plan_id)
        if self._completed.contains(completion_key):
            self._stats.inc("completed_event_dropped")
            return OperatorAction.RETURN

        state = self._states.get(logical_key)
        if state is None:
            state_num_partitions = 1 if is_partition_reducer else num_partitions
            state = _ReducerState(
                mode,
                state_num_partitions,
                max_reorder_gap=max_reorder_gap,
                groupby=groupby_meta,
            )
            self._states[logical_key] = state
        self._track_state_key(logical_key)

        if control == "done":
            accepted_done, entered_draining = state.mark_done(
                partition_id, source_final_seq=source_final_seq
            )
            if entered_draining:
                self._stats.inc("lifecycle_entered_draining")
            if not accepted_done:
                self._stats.inc("duplicate_done_dropped")
                return OperatorAction.RETURN
        else:
            accepted, missing_seq, recovered = state.add(
                partition_id,
                transformation_frame.payload,
                source_seq=source_seq,
            )
            if missing_seq:
                self._stats.inc("missing_seq", missing_seq)
            if recovered:
                self._stats.inc("out_of_order_recovered", recovered)
            if not accepted:
                self._stats.inc("late_data_dropped")
                return OperatorAction.RETURN

        if is_partition_reducer:
            # Partition reducer: finalize only after done + watermark convergence.
            # This allows `done` to arrive before the last sequenced data.
            if not state.is_complete():
                return OperatorAction.RETURN
            was_finalized = state.phase == _LifecyclePhase.FINALIZED
            partial_result = state.partial()
            if not was_finalized:
                state.phase = _LifecyclePhase.FINALIZED
                self._stats.inc("lifecycle_finalized")
            self._states.pop(logical_key, None)
            self._discard_state_key(logical_key)
            self._completed.add(completion_key)

            partial_payload = (
                partial_result
                if isinstance(partial_result, PayloadData)
                else PayloadData(args=[partial_result], kwargs={})
            )
            base_stack = FlowStack(request.flow_stack.stack[:-1])

            partial_meta = {
                "reducer": reducer_meta,
                "partition_id": partition_id,
                "control": "partial",
            }
            if isinstance(meta.get("sys"), dict):
                partial_meta["sys"] = meta["sys"]
            partial_request = FlowRequest(
                flow_stack=append_frame(
                    base_stack,
                    TransformationFrame(
                        transformation=merge_transformation,
                        payload=partial_payload,
                        meta=partial_meta,
                    ),
                ),
                request_id=request.request_id,
                origin_address=request.origin_address,
                ack_route=request.ack_route,
            )
            await ctx.emit_flow_request(merge_address, partial_request)

            done_meta = {
                "reducer": reducer_meta,
                "partition_id": partition_id,
                "control": "done",
            }
            if state.source_final_seq is not None:
                done_meta["source_final_seq"] = state.source_final_seq
            if isinstance(meta.get("sys"), dict):
                done_meta["sys"] = meta["sys"]
            done_request = FlowRequest(
                flow_stack=append_frame(
                    base_stack,
                    TransformationFrame(
                        transformation=merge_transformation,
                        payload=PayloadData(args=[None], kwargs={}),
                        meta=done_meta,
                    ),
                ),
                request_id=request.request_id,
                origin_address=request.origin_address,
                ack_route=request.ack_route,
            )
            await ctx.emit_flow_request(merge_address, done_request)
            return OperatorAction.RETURN

        if not state.is_complete():
            return OperatorAction.RETURN

        # Release pending branches except the current one (which will be acked
        # via post_process_result).
        release_count = _release_count_from_pending(state.pending_count)
        if release_count:
            await ctx.request_add_delta(request, -release_count)

        was_finalized = state.phase == _LifecyclePhase.FINALIZED
        result = state.finalize()
        if not was_finalized:
            self._stats.inc("lifecycle_finalized")
        self._states.pop(logical_key, None)
        self._discard_state_key(logical_key)
        self._completed.add(completion_key)
        await ctx.post_process_result(request, result, transformation_frame.payload)
        return OperatorAction.RETURN
