from __future__ import annotations

from collections import OrderedDict
from collections import defaultdict
from collections import deque
from enum import Enum
from typing import Any

from flownet.core.exceptions import FlowDefinitionError
from flownet.core.types import PayloadData
from flownet.runtime.flows.replay_checkpoint import ReplayCheckpoint


class _LifecyclePhase(Enum):
    OPEN = "open"
    DRAINING = "draining"
    FINALIZED = "finalized"



def _payload_value(payload: PayloadData) -> Any:
    if isinstance(payload, PayloadData):
        if len(payload.args) == 1 and not payload.kwargs:
            return payload.args[0]
        return payload
    return payload


def _is_by_key_mode(mode: str) -> bool:
    return mode in {"count_by_key", "sum_by_key", "concat_by_key"}


def _base_by_key_mode(mode: str) -> str:
    if mode == "count_by_key":
        return "count"
    if mode == "sum_by_key":
        return "sum"
    if mode == "concat_by_key":
        return "concat"
    raise FlowDefinitionError(f"Unsupported by-key reducer mode: {mode}")


def _extract_group_key(value: Any, groupby: dict[str, Any]) -> Any:
    key_field = groupby.get("key_field")
    key_index = int(groupby.get("key_index", 0))

    if isinstance(value, dict):
        if key_field is not None:
            if key_field not in value:
                raise FlowDefinitionError(f"groupby.key_field '{key_field}' is missing in payload.")
            return value[key_field]
        if "key" in value:
            return value["key"]

    if isinstance(value, (list, tuple)):
        if key_index < 0 or key_index >= len(value):
            raise FlowDefinitionError("groupby.key_index out of range.")
        return value[key_index]

    raise FlowDefinitionError(
        "reduce_by_key expects dict payload with key_field/key or sequence payload with key_index."
    )


def _extract_group_sum_value(value: Any, groupby: dict[str, Any]) -> Any:
    value_field = groupby.get("value_field")
    value_index = groupby.get("value_index")
    key_index = int(groupby.get("key_index", 0))

    if isinstance(value, dict):
        if value_field is not None:
            if value_field not in value:
                raise FlowDefinitionError(f"groupby.value_field '{value_field}' is missing in payload.")
            return value[value_field]
        if "value" in value:
            return value["value"]

    if isinstance(value, (list, tuple)):
        if value_index is not None:
            idx = int(value_index)
            if idx < 0 or idx >= len(value):
                raise FlowDefinitionError("groupby.value_index out of range.")
            return value[idx]
        if len(value) >= 2:
            idx = 1 if key_index == 0 else 0
            if idx < len(value):
                return value[idx]

    raise FlowDefinitionError(
        "sum_by_key requires value_field/value_index or pair-like payload (key, value)."
    )


def _sorted_keyed_items(acc: dict[Any, Any]) -> list[tuple[Any, Any]]:
    return sorted(acc.items(), key=lambda kv: repr(kv[0]))


def _partial_keyed_items(value: Any) -> list[tuple[Any, Any]]:
    if isinstance(value, dict):
        return list(value.items())
    if isinstance(value, (list, tuple)):
        items: list[tuple[Any, Any]] = []
        for item in value:
            if not isinstance(item, (list, tuple)) or len(item) != 2:
                raise FlowDefinitionError(
                    "by-key merge partial list entries must be (key, value) pairs."
                )
            items.append((item[0], item[1]))
        return items
    raise FlowDefinitionError("by-key merge partial must be dict or list of (key, value) pairs.")


def _release_count_from_pending(pending_count: int) -> int:
    # Deferred-ack policy for stateful operators:
    # - keep current terminal event for post_process_result ack
    # - release all prior pending branches at completion
    return max(pending_count - 1, 0)


def _release_count_for_merge(num_partitions: int, pending_count: int) -> int:
    # Merge protocol emits one terminal result and releases one branch per
    # completed partition except the current terminal event branch.
    by_partitions = max(num_partitions - 1, 0)
    by_pending = _release_count_from_pending(pending_count)
    return min(by_partitions, by_pending)


def _parse_non_negative_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        return None
    if parsed < 0:
        return None
    return parsed


def _checkpoint_boundary_token(checkpoint: ReplayCheckpoint) -> tuple[int, str]:
    return checkpoint.epoch, (checkpoint.plan_id or "-")



class _CompletionTracker:
    """Bounded dedupe for terminal signals per request/stage key."""

    def __init__(self, *, max_entries: int = 10000) -> None:
        self._max_entries = max_entries
        self._order: deque[tuple[Any, ...]] = deque()
        self._seen: set[tuple[Any, ...]] = set()
        self._evicted_total = 0

    def contains(self, key: tuple[Any, ...]) -> bool:
        return key in self._seen

    def add(self, key: tuple[Any, ...]) -> None:
        if key in self._seen:
            return
        self._seen.add(key)
        self._order.append(key)
        while len(self._order) > self._max_entries:
            old_key = self._order.popleft()
            self._seen.discard(old_key)
            self._evicted_total += 1

    @property
    def size(self) -> int:
        return len(self._seen)

    @property
    def evicted_total(self) -> int:
        return self._evicted_total


class _BoundaryTracker:
    """Bounded map for latest replay boundary tokens per logical key."""

    def __init__(self, *, max_entries: int = 10000) -> None:
        self._max_entries = max_entries
        self._entries: OrderedDict[tuple[Any, ...], tuple[int, str]] = OrderedDict()
        self._evicted_total = 0

    def get(self, key: tuple[Any, ...]) -> tuple[int, str] | None:
        value = self._entries.get(key)
        if value is None:
            return None
        # Treat recent reads as active to avoid evicting hot keys first.
        self._entries.move_to_end(key)
        return value

    def set(self, key: tuple[Any, ...], value: tuple[int, str]) -> None:
        self._entries[key] = value
        self._entries.move_to_end(key)
        while len(self._entries) > self._max_entries:
            self._entries.popitem(last=False)
            self._evicted_total += 1

    @property
    def size(self) -> int:
        return len(self._entries)

    @property
    def evicted_total(self) -> int:
        return self._evicted_total


class _StateKeyLRU:
    """Bound active state keys to avoid unbounded handler state growth."""

    def __init__(self, *, max_entries: int = 10000) -> None:
        self._max_entries = max(1, int(max_entries))
        self._entries: OrderedDict[tuple[Any, ...], None] = OrderedDict()
        self._evicted_total = 0

    def touch(self, key: tuple[Any, ...]) -> tuple[Any, ...] | None:
        self._entries[key] = None
        self._entries.move_to_end(key)
        if len(self._entries) <= self._max_entries:
            return None
        evicted_key, _ = self._entries.popitem(last=False)
        self._evicted_total += 1
        return evicted_key

    def discard(self, key: tuple[Any, ...]) -> None:
        self._entries.pop(key, None)

    @property
    def size(self) -> int:
        return len(self._entries)

    @property
    def evicted_total(self) -> int:
        return self._evicted_total


class _OperatorStats:
    def __init__(self) -> None:
        self._counters: defaultdict[str, int] = defaultdict(int)

    def inc(self, key: str, amount: int = 1) -> None:
        self._counters[key] += amount

    def snapshot(self) -> dict[str, int]:
        return dict(self._counters)


class _StatefulOperatorMixin:
    """Shared state key tracking and bounded-cache scaffolding."""

    _states: dict[tuple[Any, ...], Any]
    _active_state_keys: _StateKeyLRU
    _state_last_touch: dict[tuple[Any, ...], int]
    _max_state_idle_ops: int
    _state_cleanup_interval: int
    _op_seq: int
    _completed: _CompletionTracker
    _latest_boundary: _BoundaryTracker
    _stats: _OperatorStats

    def _init_stateful_runtime(
        self,
        *,
        max_active_states: int,
        max_state_idle_ops: int,
        state_cleanup_interval: int,
    ) -> None:
        self._active_state_keys = _StateKeyLRU(max_entries=max_active_states)
        self._state_last_touch = {}
        self._max_state_idle_ops = max(1, int(max_state_idle_ops))
        self._state_cleanup_interval = max(1, int(state_cleanup_interval))
        self._op_seq = 0
        self._completed = _CompletionTracker()
        self._latest_boundary = _BoundaryTracker()
        self._stats = _OperatorStats()

    def _base_stats_snapshot(self) -> dict[str, int]:
        snapshot = self._stats.snapshot()
        snapshot["active_state_entries"] = len(self._states)
        snapshot["active_state_lru_entries"] = self._active_state_keys.size
        snapshot["active_state_lru_evicted_total"] = self._active_state_keys.evicted_total
        snapshot["state_touch_entries"] = len(self._state_last_touch)
        snapshot["boundary_cache_entries"] = self._latest_boundary.size
        snapshot["boundary_cache_evicted_total"] = self._latest_boundary.evicted_total
        snapshot["completion_cache_entries"] = self._completed.size
        snapshot["completion_cache_evicted_total"] = self._completed.evicted_total
        return snapshot

    def _track_state_key(self, logical_key: tuple[Any, ...]) -> None:
        self._state_last_touch[logical_key] = self._op_seq
        evicted_key = self._active_state_keys.touch(logical_key)
        if evicted_key is None or evicted_key == logical_key:
            return
        self._state_last_touch.pop(evicted_key, None)
        if self._states.pop(evicted_key, None) is not None:
            self._stats.inc("state_evicted_overflow")

    def _discard_state_key(self, logical_key: tuple[Any, ...]) -> None:
        self._active_state_keys.discard(logical_key)
        self._state_last_touch.pop(logical_key, None)

    def _tick_and_maybe_cleanup(self) -> None:
        self._op_seq += 1
        if (self._op_seq % self._state_cleanup_interval) != 0:
            return
        cutoff = self._op_seq - self._max_state_idle_ops
        if cutoff <= 0 or not self._state_last_touch:
            return
        stale_keys = [key for key, last in self._state_last_touch.items() if last <= cutoff]
        for key in stale_keys:
            self._state_last_touch.pop(key, None)
            self._active_state_keys.discard(key)
            if self._states.pop(key, None) is not None:
                self._stats.inc("state_evicted_idle")
