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,
    _parse_non_negative_int,
    _StatefulOperatorMixin,
    _checkpoint_boundary_token,
    _payload_value,
    _release_count_from_pending,
)


class _JoinReducerState:
    def __init__(self, num_partitions: int) -> None:
        self.phase = _LifecyclePhase.OPEN
        self.num_partitions = num_partitions
        self.pending_count = 0
        self.left: dict[int, list[Any]] = {}
        self.right: dict[int, list[Any]] = {}
        self.done_left: set[int] = set()
        self.done_right: set[int] = set()
        self.source_final_seq_left: int | None = None
        self.source_final_seq_right: int | None = None
        self._next_source_seq_left = 1
        self._next_source_seq_right = 1
        self._future_source_seq_left: set[int] = set()
        self._future_source_seq_right: set[int] = set()

    def add(
        self,
        side: str,
        partition_id: int,
        payload: PayloadData,
        *,
        source_seq: int | None = None,
    ) -> bool:
        if self.phase == _LifecyclePhase.FINALIZED:
            return False
        done_bucket = self.done_left if side == "left" else self.done_right
        if partition_id in done_bucket and not self._allow_post_done_data(side):
            return False
        self.pending_count += 1
        value = _payload_value(payload)
        bucket = self.left if side == "left" else self.right
        bucket.setdefault(partition_id, []).append(value)
        if source_seq is not None:
            self._track_source_seq(side, source_seq)
        return True

    def mark_done(
        self,
        side: str,
        partition_id: int,
        *,
        source_final_seq: int | None = None,
    ) -> tuple[bool, bool]:
        if self.phase == _LifecyclePhase.FINALIZED:
            return False, False
        done_bucket = self.done_left if side == "left" else self.done_right
        if partition_id in done_bucket:
            return False, False
        self.pending_count += 1
        if side == "left":
            self.done_left.add(partition_id)
            if source_final_seq is not None:
                current = self.source_final_seq_left
                if current is None or source_final_seq > current:
                    self.source_final_seq_left = source_final_seq
        else:
            self.done_right.add(partition_id)
            if source_final_seq is not None:
                current = self.source_final_seq_right
                if current is None or source_final_seq > current:
                    self.source_final_seq_right = 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_left) < self.num_partitions or len(self.done_right) < self.num_partitions:
            return False
        if (
            self.source_final_seq_left is not None
            and self._contiguous_source_seq("left") < self.source_final_seq_left
        ):
            return False
        if (
            self.source_final_seq_right is not None
            and self._contiguous_source_seq("right") < self.source_final_seq_right
        ):
            return False
        return True

    def _allow_post_done_data(self, side: str) -> bool:
        source_final_seq = (
            self.source_final_seq_left if side == "left" else self.source_final_seq_right
        )
        if source_final_seq is None:
            return False
        return self._contiguous_source_seq(side) < source_final_seq

    def _track_source_seq(self, side: str, source_seq: int) -> None:
        if side == "left":
            self._next_source_seq_left = self._advance_source_seq(
                self._next_source_seq_left,
                self._future_source_seq_left,
                source_seq,
            )
            return
        self._next_source_seq_right = self._advance_source_seq(
            self._next_source_seq_right,
            self._future_source_seq_right,
            source_seq,
        )

    @staticmethod
    def _advance_source_seq(next_seq: int, future: set[int], source_seq: int) -> int:
        if source_seq < 1 or source_seq < next_seq:
            return next_seq
        if source_seq == next_seq:
            next_seq += 1
            while next_seq in future:
                future.remove(next_seq)
                next_seq += 1
            return next_seq
        future.add(source_seq)
        return next_seq

    def _contiguous_source_seq(self, side: str) -> int:
        if side == "left":
            return self._next_source_seq_left - 1
        return self._next_source_seq_right - 1

    def finalize(self) -> list[tuple[Any, Any]]:
        self.phase = _LifecyclePhase.FINALIZED
        pairs: list[tuple[Any, Any]] = []
        for pid in range(self.num_partitions):
            left_vals = self.left.get(pid, [])
            right_vals = self.right.get(pid, [])
            for l in left_vals:
                for r in right_vals:
                    pairs.append((l, r))
        return pairs



class JoinReducerOperatorHandler(_StatefulOperatorMixin, OperatorHandler):
    """Internal join reducer (pseudo) that pairs left/right partitions."""

    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], _JoinReducerState] = {}
        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_left_max"] = max(len(state.done_left) for state in self._states.values())
            snapshot["state_done_right_max"] = max(len(state.done_right) for state in self._states.values())
            snapshot["state_left_values_max"] = max(
                sum(len(v) for v in state.left.values()) for state in self._states.values()
            )
            snapshot["state_right_values_max"] = max(
                sum(len(v) for v in state.right.values()) 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 {}
        join_meta = require_dict(meta.get("join"), "join")

        stage_id = join_meta.get("stage_id") or transformation_frame.transformation.trans_id
        num_partitions = require_positive_int(
            join_meta.get("num_partitions", 1), "join.num_partitions"
        )
        partition_id = validate_partition_id(
            meta.get("partition_id", 0), num_partitions, field="partition_id"
        )
        side = meta.get("side")
        if side not in {"left", "right"}:
            raise FlowDefinitionError("Join reducer meta requires side=left/right.")

        control = meta.get("control", "data")
        if control not in {"data", "done"}:
            raise FlowDefinitionError("Join reducer control must be one of: data, done.")
        source_seq = _parse_non_negative_int(meta.get("source_seq"))
        source_final_seq = _parse_non_negative_int(meta.get("source_final_seq"))

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

        if control == "done":
            accepted_done, entered_draining = state.mark_done(
                side,
                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:
            if not state.add(
                side,
                partition_id,
                transformation_frame.payload,
                source_seq=source_seq,
            ):
                self._stats.inc("late_data_dropped")
                return OperatorAction.RETURN

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

        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
