from __future__ import annotations

from typing import Any

from flownet.core.exceptions import FlowDefinitionError
from flownet.core.request import FlowRequest, 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.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,
    _is_by_key_mode,
    _partial_keyed_items,
    _payload_value,
    _release_count_for_merge,
    _sorted_keyed_items,
)


class _MergeReducerState:
    def __init__(self, mode: str, num_partitions: int) -> None:
        self.phase = _LifecyclePhase.OPEN
        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.num_partitions = num_partitions
        self.done_partitions: set[int] = set()
        self.partial_partitions: set[int] = set()
        self.pending_count = 0
        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_partial(self, partition_id: int, payload: PayloadData) -> bool:
        if self.phase == _LifecyclePhase.FINALIZED:
            return False
        if partition_id in self.partial_partitions:
            return False
        self.pending_count += 1
        self.partial_partitions.add(partition_id)
        value = _payload_value(payload)
        if self._by_key_mode:
            for key, partial_value in _partial_keyed_items(value):
                if self._base_mode in {"count", "sum"}:
                    self.acc[key] = self.acc.get(key, 0) + partial_value
                else:
                    bucket = self.acc.setdefault(key, [])
                    if isinstance(partial_value, list):
                        bucket.extend(partial_value)
                    else:
                        bucket.append(partial_value)
            return True

        if self.mode == "count":
            self.acc += value
        elif self.mode == "sum":
            self.acc += value
        else:
            if isinstance(value, list):
                self.acc.extend(value)
            else:
                self.acc.append(value)
        return True

    def mark_done(self, partition_id: int) -> 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)
        entered_draining = False
        if self.phase == _LifecyclePhase.OPEN:
            self.phase = _LifecyclePhase.DRAINING
            entered_draining = True
        return True, entered_draining

    def is_complete(self) -> bool:
        return (
            len(self.done_partitions) >= self.num_partitions
            and len(self.partial_partitions) >= self.num_partitions
        )

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



class MergeReducerOperatorHandler(_StatefulOperatorMixin, OperatorHandler):
    """Merge reducer that combines partition partials across containers."""

    def __init__(
        self,
        *,
        max_active_states: int = 10000,
        max_state_idle_ops: int = 10000,
        state_cleanup_interval: int = 256,
    ) -> None:
        self._states: dict[tuple[str, str], _MergeReducerState] = {}
        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_partial_partitions_max"] = max(
                len(state.partial_partitions) 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"
        )
        merge_partitions = reducer_meta.get("merge_partitions")
        expected_partition_ids: set[int] | None = None
        if merge_partitions is not None:
            if not isinstance(merge_partitions, (list, tuple)):
                raise FlowDefinitionError("reducer.merge_partitions must be a list when provided.")
            try:
                expected_partition_ids = {int(pid) for pid in merge_partitions}
            except (TypeError, ValueError) as exc:
                raise FlowDefinitionError(
                    "reducer.merge_partitions must contain integer partition ids."
                ) from exc
            num_partitions = len(expected_partition_ids)
            if num_partitions == 0:
                raise FlowDefinitionError("reducer.merge_partitions must not be empty.")
        mode = reducer_meta.get("mode", "concat")
        try:
            partition_id = int(meta.get("partition_id", 0))
        except (TypeError, ValueError) as exc:
            raise FlowDefinitionError("partition_id must be an integer >= 0.") from exc
        control = meta.get("control", "partial")
        if control not in {"partial", "done"}:
            raise FlowDefinitionError("Merge reducer control must be one of: partial, done.")

        if expected_partition_ids is not None:
            if partition_id not in expected_partition_ids:
                return OperatorAction.RETURN
        else:
            partition_id = validate_partition_id(
                partition_id, num_partitions, field="partition_id"
            )

        logical_key = (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 = _MergeReducerState(mode, num_partitions)
            self._states[logical_key] = state
        self._track_state_key(logical_key)

        if control == "partial":
            accepted = state.add_partial(partition_id, transformation_frame.payload)
            if not accepted:
                self._stats.inc("duplicate_partial_dropped")
        else:
            accepted_done, entered_draining = state.mark_done(partition_id)
            if entered_draining:
                self._stats.inc("lifecycle_entered_draining")
            accepted = accepted_done
            if not accepted_done:
                self._stats.inc("duplicate_done_dropped")
        if not accepted:
            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). Pending branches = num_partitions.
        release_count = _release_count_for_merge(
            num_partitions=num_partitions,
            pending_count=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
