from __future__ import annotations

from collections.abc import Iterable, Mapping
from dataclasses import dataclass, field
import hashlib
import inspect
import json
import time
from typing import Any, Callable

from flownet.v1.runtime.flowengine.cursor_models import FlowProgramRef


@dataclass(frozen=True)
class OperatorEvaluation:
    outputs: tuple[Any, ...]
    output_tags: tuple[dict[str, str], ...] | None = None
    route_node_ids: tuple[str, ...] | None = None
    process_program_ref: FlowProgramRef | None = None


class OperatorExecutionError(RuntimeError):
    """Raised when operator runtime semantics cannot be evaluated."""


class UnsupportedOperatorTypeError(OperatorExecutionError):
    """Raised when operator type is not supported by the v1 runtime slice."""


class ProcessTargetResolutionError(OperatorExecutionError):
    """Raised when process operator cannot resolve a child FlowProgramRef."""


class StatefulProcessContractError(OperatorExecutionError):
    """Raised when sys/stateful-process violates the minimal runtime contract."""


class ShuffleOperatorContractError(OperatorExecutionError):
    """Raised when sys/shuffle violates the minimal runtime contract."""


class ReducerContractError(OperatorExecutionError):
    """Raised when sys/reducer violates the frozen v1 reducer contract."""


class JoinContractError(ReducerContractError):
    """Raised when join/join-reducer violates the frozen v1 join contract."""


class MergeReducerContractError(ReducerContractError):
    """Raised when sys/merge-reducer violates the frozen v1 merge contract."""


@dataclass
class _StateEntry:
    value: Any
    updated_at: float


@dataclass
class _ReducerState:
    stage_id: str
    mode: str
    num_partitions: int
    groupby: Mapping[str, Any]
    acc: Any
    done_partitions: set[int] = field(default_factory=set)


@dataclass
class _JoinReducerState:
    stage_id: str
    num_partitions: int
    left: dict[int, list[Any]] = field(default_factory=dict)
    right: dict[int, list[Any]] = field(default_factory=dict)
    done_left: set[int] = field(default_factory=set)
    done_right: set[int] = field(default_factory=set)


@dataclass
class _MergeReducerState:
    stage_id: str
    mode: str
    expected_partition_ids: tuple[int, ...]
    acc: Any
    partial_partitions: set[int] = field(default_factory=set)
    done_partitions: set[int] = field(default_factory=set)


@dataclass(frozen=True)
class StatefulStateHandle:
    stage_id: str
    namespace: str
    key: str
    ttl_seconds: float | None
    _runtime: "StatefulProcessRuntime" = field(repr=False, compare=False)
    _slot: tuple[int, str, str] = field(repr=False, compare=False)

    def get(self, default: Any = None) -> Any:
        return self._runtime.get(self._slot, default=default, ttl_seconds=self.ttl_seconds)

    def set(self, value: Any) -> Any:
        self._runtime.set(self._slot, value, ttl_seconds=self.ttl_seconds)
        return value

    def clear(self) -> bool:
        return self._runtime.delete(self._slot)


class StatefulProcessRuntime:
    """
    Minimal state store for v1 `sys/stateful-process`.

    Contract freeze (Lane-C, phase-2 first slice):
    1. key source priority: `state_spec.key_field` > `state_spec.key_index` > tag `state_partition_key` > `0`.
    2. handler receives a mutable `StatefulStateHandle` and can `get/set/clear` state per resolved key.
    3. unsupported advanced families (`reducer/join/merge`) fail fast with explicit error type.
    """

    def __init__(self, *, time_fn: Callable[[], float] | None = None) -> None:
        self._time_fn = time_fn or time.monotonic
        self._state: dict[tuple[int, str, str], _StateEntry] = {}

    def resolve_state_handle(
        self,
        *,
        transformation: Any,
        payload: Any,
        tags: Mapping[str, str],
    ) -> StatefulStateHandle:
        process_meta = _extract_stateful_process_meta(transformation)
        stage_id = _extract_state_stage_id(process_meta, transformation)
        namespace = _extract_state_namespace(process_meta, stage_id=stage_id)
        ttl_seconds = _extract_state_ttl_seconds(process_meta)
        key_source = _resolve_state_key_source(
            payload=payload,
            tags=tags,
            state_spec=process_meta.get("state_spec"),
        )
        key = _stable_key_fragment(key_source)
        slot = (id(transformation), namespace, key)
        self._expire_if_needed(slot, ttl_seconds=ttl_seconds)
        return StatefulStateHandle(
            stage_id=stage_id,
            namespace=namespace,
            key=key,
            ttl_seconds=ttl_seconds,
            _runtime=self,
            _slot=slot,
        )

    def get(
        self,
        slot: tuple[int, str, str],
        *,
        default: Any,
        ttl_seconds: float | None,
    ) -> Any:
        self._expire_if_needed(slot, ttl_seconds=ttl_seconds)
        entry = self._state.get(slot)
        if entry is None:
            return default
        entry.updated_at = self._time_fn()
        return entry.value

    def set(
        self,
        slot: tuple[int, str, str],
        value: Any,
        *,
        ttl_seconds: float | None,
    ) -> None:
        del ttl_seconds  # retained on handle; expiry checks are driven by caller contract
        self._state[slot] = _StateEntry(
            value=value,
            updated_at=self._time_fn(),
        )

    def delete(self, slot: tuple[int, str, str]) -> bool:
        return self._state.pop(slot, None) is not None

    def _expire_if_needed(
        self,
        slot: tuple[int, str, str],
        *,
        ttl_seconds: float | None,
    ) -> None:
        if ttl_seconds is None:
            return
        entry = self._state.get(slot)
        if entry is None:
            return
        if (self._time_fn() - entry.updated_at) >= ttl_seconds:
            self._state.pop(slot, None)


class ReducerRuntime:
    """
    Minimal in-memory runtime for v1 `sys/reducer`.

    Contract freeze (Lane-C2):
    1. state key = `(operator identity, event_group_id, reducer.stage_id)`.
    2. frame contract:
       - data: raw payload OR `{control="data", partition_id?, value}`
       - flush: `{control="flush", partition_id?}`
       - done: `{control="done", partition_id?}`
    3. terminal output is emitted once all partitions are done.
       `flush` emits snapshot without finalizing state.
    """

    def __init__(self) -> None:
        self._states: dict[tuple[int, str, str], _ReducerState] = {}

    def evaluate(
        self,
        *,
        transformation: Any,
        payload: Any,
        tags: Mapping[str, str],
        event_group_id: str | None,
    ) -> tuple[Any, ...]:
        if not isinstance(event_group_id, str) or not event_group_id.strip():
            raise ReducerContractError(
                "sys/reducer requires a non-empty event_group_id for state isolation.",
            )
        reducer_meta = _extract_reducer_meta(transformation)
        stage_id = _extract_reducer_stage_id(reducer_meta, transformation)
        mode = _extract_reducer_mode(reducer_meta)
        num_partitions = _extract_reducer_num_partitions(reducer_meta)
        groupby = _extract_reducer_groupby(reducer_meta)

        state_key = (id(transformation), event_group_id, stage_id)
        state = self._states.get(state_key)
        if state is None:
            state = _ReducerState(
                stage_id=stage_id,
                mode=mode,
                num_partitions=num_partitions,
                groupby=groupby,
                acc=_new_reducer_accumulator(mode),
            )
            self._states[state_key] = state

        control, partition_id, value = _parse_reducer_frame(
            payload=payload,
            tags=tags,
            num_partitions=state.num_partitions,
        )
        if control == "data":
            if partition_id in state.done_partitions:
                # Deterministic late-data drop after partition done.
                return ()
            _accumulate_reducer_value(state=state, value=value)
            return ()

        if control == "flush":
            return (_materialize_reducer_snapshot(state),)

        if partition_id in state.done_partitions:
            return ()
        state.done_partitions.add(partition_id)
        if len(state.done_partitions) < state.num_partitions:
            return ()

        output = _materialize_reducer_snapshot(state)
        self._states.pop(state_key, None)
        return (output,)


class JoinReducerRuntime:
    """
    Minimal in-memory runtime for `sys/join-reducer`.

    Contract freeze (Lane-C3):
    1. state key = `(operator identity, event_group_id, join.stage_id)`.
    2. frame contract:
       - data: raw payload OR `{control="data", side, partition_id?, value}`
       - flush: `{control="flush"}`
       - done: `{control="done", side, partition_id?}`
    3. terminal output is emitted once all partitions are done on both sides.
       Snapshot/final output is deterministic partition-order cartesian pairs.
    """

    def __init__(self) -> None:
        self._states: dict[tuple[int, str, str], _JoinReducerState] = {}

    def evaluate(
        self,
        *,
        transformation: Any,
        payload: Any,
        tags: Mapping[str, str],
        event_group_id: str | None,
    ) -> tuple[Any, ...]:
        if not isinstance(event_group_id, str) or not event_group_id.strip():
            raise JoinContractError(
                "sys/join-reducer requires a non-empty event_group_id for state isolation.",
            )
        stage_id, num_partitions = _resolve_join_reducer_contract(
            transformation=transformation,
            tags=tags,
        )
        state_key = (id(transformation), event_group_id, stage_id)
        state = self._states.get(state_key)
        if state is None:
            state = _JoinReducerState(
                stage_id=stage_id,
                num_partitions=num_partitions,
            )
            self._states[state_key] = state
        elif state.num_partitions != num_partitions:
            raise JoinContractError(
                "join.num_partitions changed while state is active for the same stage/event_group.",
            )

        control, side, partition_id, value = _parse_join_reducer_frame(
            payload=payload,
            tags=tags,
            num_partitions=num_partitions,
        )
        if control == "flush":
            return (_materialize_join_reducer_snapshot(state),)

        bucket_done = state.done_left if side == "left" else state.done_right
        if control == "data":
            if partition_id in bucket_done:
                # Deterministic late-data drop after side/partition done.
                return ()
            bucket = state.left if side == "left" else state.right
            bucket.setdefault(partition_id, []).append(value)
            return ()

        if partition_id in bucket_done:
            return ()
        bucket_done.add(partition_id)
        if len(state.done_left) < num_partitions or len(state.done_right) < num_partitions:
            return ()

        output = _materialize_join_reducer_snapshot(state)
        self._states.pop(state_key, None)
        return (output,)


class MergeReducerRuntime:
    """
    Minimal in-memory runtime for `sys/merge-reducer`.

    Contract freeze (Lane-C3):
    1. state key = `(operator identity, event_group_id, reducer.stage_id)`.
    2. frame contract:
       - partial: raw payload OR `{control="partial", partition_id?, value}`
       - flush: `{control="flush"}`
       - done: `{control="done", partition_id?}`
    3. terminal output is emitted once all expected partitions have both
       `partial` and `done` markers.
    """

    def __init__(self) -> None:
        self._states: dict[tuple[int, str, str], _MergeReducerState] = {}

    def evaluate(
        self,
        *,
        transformation: Any,
        payload: Any,
        tags: Mapping[str, str],
        event_group_id: str | None,
    ) -> tuple[Any, ...]:
        if not isinstance(event_group_id, str) or not event_group_id.strip():
            raise MergeReducerContractError(
                "sys/merge-reducer requires a non-empty event_group_id for state isolation.",
            )
        stage_id, mode, expected_partition_ids = _resolve_merge_reducer_contract(
            transformation=transformation,
            tags=tags,
        )
        state_key = (id(transformation), event_group_id, stage_id)
        state = self._states.get(state_key)
        if state is None:
            state = _MergeReducerState(
                stage_id=stage_id,
                mode=mode,
                expected_partition_ids=expected_partition_ids,
                acc=_new_reducer_accumulator(mode),
            )
            self._states[state_key] = state
        elif state.mode != mode or state.expected_partition_ids != expected_partition_ids:
            raise MergeReducerContractError(
                "merge-reducer mode/expected partitions changed while state is active.",
            )

        control, partition_id, value = _parse_merge_reducer_frame(
            payload=payload,
            tags=tags,
            expected_partition_ids=expected_partition_ids,
        )
        if partition_id not in expected_partition_ids:
            return ()

        if control == "flush":
            return (_materialize_merge_reducer_snapshot(state),)

        if control == "partial":
            if partition_id in state.partial_partitions:
                return ()
            state.partial_partitions.add(partition_id)
            state.acc = _merge_partial_into_reducer_accumulator(
                mode=mode,
                acc=state.acc,
                value=value,
            )
            return ()

        if partition_id in state.done_partitions:
            return ()
        state.done_partitions.add(partition_id)
        if not _merge_reducer_is_complete(state):
            return ()

        output = _materialize_merge_reducer_snapshot(state)
        self._states.pop(state_key, None)
        return (output,)


_ADVANCED_UNSUPPORTED_OPERATOR_TYPES = frozenset()


def evaluate_operator(
    *,
    transformation: Any,
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
    state_runtime: StatefulProcessRuntime | None = None,
    reducer_runtime: ReducerRuntime | None = None,
    join_reducer_runtime: JoinReducerRuntime | None = None,
    merge_reducer_runtime: MergeReducerRuntime | None = None,
    event_group_id: str | None = None,
) -> OperatorEvaluation:
    operator_type = str(getattr(transformation, "type", "")).strip()
    if operator_type == "map":
        result = _invoke_target(transformation, payload, tags, seq, invoke_target)
        return OperatorEvaluation(outputs=(result,))

    if operator_type == "filter":
        keep = bool(_invoke_target(transformation, payload, tags, seq, invoke_target))
        return OperatorEvaluation(outputs=(payload,) if keep else ())

    if operator_type == "flatmap":
        result = _invoke_target(transformation, payload, tags, seq, invoke_target)
        return OperatorEvaluation(outputs=_normalize_iterable_outputs(result))

    if operator_type in {"route", "sys/router"}:
        route_node_ids = _extract_route_node_ids(transformation)
        return OperatorEvaluation(outputs=(payload,), route_node_ids=route_node_ids)

    if operator_type == "process":
        child_program_ref = _resolve_process_program_ref(transformation)
        return OperatorEvaluation(
            outputs=(payload,),
            process_program_ref=child_program_ref,
        )

    if operator_type == "sys/shuffle":
        return _evaluate_shuffle(
            transformation=transformation,
            payload=payload,
            tags=tags,
            seq=seq,
            invoke_target=invoke_target,
        )

    if operator_type == "sys/join":
        return _evaluate_join(
            transformation=transformation,
            payload=payload,
            tags=tags,
            seq=seq,
            invoke_target=invoke_target,
        )

    if operator_type == "sys/stateful-process":
        if state_runtime is None:
            raise StatefulProcessContractError(
                "state_runtime_not_configured:operator_type=sys/stateful-process",
            )
        result = _invoke_stateful_target(
            transformation=transformation,
            payload=payload,
            tags=tags,
            seq=seq,
            invoke_target=invoke_target,
            state_runtime=state_runtime,
        )
        return OperatorEvaluation(outputs=_normalize_iterable_outputs(result))

    if operator_type == "sys/reducer":
        if reducer_runtime is None:
            raise ReducerContractError(
                "reducer_runtime_not_configured:operator_type=sys/reducer",
            )
        outputs = reducer_runtime.evaluate(
            transformation=transformation,
            payload=payload,
            tags=tags,
            event_group_id=event_group_id,
        )
        return OperatorEvaluation(outputs=outputs)

    if operator_type == "sys/join-reducer":
        if join_reducer_runtime is None:
            raise JoinContractError(
                "join_reducer_runtime_not_configured:operator_type=sys/join-reducer",
            )
        outputs = join_reducer_runtime.evaluate(
            transformation=transformation,
            payload=payload,
            tags=tags,
            event_group_id=event_group_id,
        )
        return OperatorEvaluation(outputs=outputs)

    if operator_type == "sys/merge-reducer":
        if merge_reducer_runtime is None:
            raise MergeReducerContractError(
                "merge_reducer_runtime_not_configured:operator_type=sys/merge-reducer",
            )
        outputs = merge_reducer_runtime.evaluate(
            transformation=transformation,
            payload=payload,
            tags=tags,
            event_group_id=event_group_id,
        )
        return OperatorEvaluation(outputs=outputs)

    if operator_type in _ADVANCED_UNSUPPORTED_OPERATOR_TYPES:
        raise UnsupportedOperatorTypeError(f"unsupported_operator_type:{operator_type}")

    raise UnsupportedOperatorTypeError(f"unsupported_operator_type:{operator_type}")


def _invoke_target(
    transformation: Any,
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
) -> Any:
    if invoke_target is None:
        raise OperatorExecutionError(
            f"invoke_target_not_configured:operator_type={getattr(transformation, 'type', None)}",
        )
    target = getattr(transformation, "target", None)
    return invoke_target(target, payload, tags, seq)


def _normalize_iterable_outputs(result: Any) -> tuple[Any, ...]:
    if result is None:
        return ()
    if isinstance(result, tuple):
        return result
    if isinstance(result, list):
        return tuple(result)
    if isinstance(result, Iterable) and not isinstance(result, (str, bytes, bytearray, Mapping)):
        return tuple(result)
    return (result,)


def _invoke_stateful_target(
    *,
    transformation: Any,
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
    state_runtime: StatefulProcessRuntime,
) -> Any:
    state_handle = state_runtime.resolve_state_handle(
        transformation=transformation,
        payload=payload,
        tags=tags,
    )
    target = getattr(transformation, "target", None)
    if callable(target):
        return _invoke_stateful_callable(
            target=target,
            payload=payload,
            state=state_handle,
            tags=tags,
            seq=seq,
        )
    if invoke_target is not None:
        return invoke_target(target, (payload, state_handle), tags, seq)
    raise StatefulProcessContractError(
        "sys/stateful-process target must be callable or invoke_target-capable.",
    )


def _invoke_stateful_callable(
    *,
    target: Callable[..., Any],
    payload: Any,
    state: StatefulStateHandle,
    tags: Mapping[str, str],
    seq: int | None,
) -> Any:
    try:
        signature = inspect.signature(target)
    except (TypeError, ValueError):
        return target(payload, state)

    positional_capacity = 0
    has_var_positional = False
    for parameter in signature.parameters.values():
        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
            has_var_positional = True
            break
        if parameter.kind in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        ):
            positional_capacity += 1

    if has_var_positional or positional_capacity >= 4:
        return target(payload, state, tags, seq)
    if positional_capacity == 3:
        return target(payload, state, tags)
    if positional_capacity == 2:
        return target(payload, state)
    if positional_capacity == 1:
        return target(payload)
    return target()


def _extract_route_node_ids(transformation: Any) -> tuple[str, ...]:
    operator_config = getattr(transformation, "operator_config", None)
    if operator_config is None:
        return ()
    if not isinstance(operator_config, Mapping):
        raise OperatorExecutionError("route operator_config must be a mapping.")

    routes = operator_config.get("routes")
    if routes is None:
        return ()
    if not isinstance(routes, list):
        raise OperatorExecutionError("route operator_config.routes must be a list.")

    route_node_ids: list[str] = []
    for index, route in enumerate(routes):
        if not isinstance(route, Mapping):
            raise OperatorExecutionError(f"route[{index}] must be a mapping.")
        trans = route.get("transformation")
        if isinstance(trans, str) and trans.strip():
            route_node_ids.append(trans.strip())
            continue
        trans_id = getattr(trans, "trans_id", None)
        if isinstance(trans_id, str) and trans_id.strip():
            route_node_ids.append(trans_id)
            continue
        raise OperatorExecutionError(f"route[{index}].transformation missing valid trans_id.")
    return tuple(route_node_ids)


def _evaluate_shuffle(
    *,
    transformation: Any,
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
) -> OperatorEvaluation:
    shuffle_meta = _extract_shuffle_meta(transformation)
    num_partitions = _extract_shuffle_num_partitions(shuffle_meta)
    stage_id = _extract_shuffle_stage_id(shuffle_meta, transformation)
    partition_key_source = _resolve_shuffle_partition_key_source(
        shuffle_meta=shuffle_meta,
        payload=payload,
        tags=tags,
        seq=seq,
        invoke_target=invoke_target,
    )
    key_fragment = _stable_key_fragment(partition_key_source)
    partition_id = int(key_fragment, 16) % num_partitions

    output_tags = dict(tags)
    output_tags["partition_id"] = str(partition_id)
    output_tags["state_partition_key"] = key_fragment
    output_tags["shuffle_stage_id"] = stage_id
    output_tags["shuffle_num_partitions"] = str(num_partitions)

    route_node_ids = _select_shuffle_route_node_ids(
        transformation=transformation,
        partition_id=partition_id,
    )
    return OperatorEvaluation(
        outputs=(payload,),
        output_tags=(output_tags,),
        route_node_ids=route_node_ids,
    )


def _evaluate_join(
    *,
    transformation: Any,
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
) -> OperatorEvaluation:
    join_meta = _extract_join_meta(transformation)
    side = _extract_join_side(transformation)

    shuffle_meta = _extract_shuffle_meta(transformation)
    shuffle_stage_id = _extract_shuffle_stage_id(shuffle_meta, transformation)
    shuffle_num_partitions = _extract_shuffle_num_partitions(shuffle_meta)

    join_stage_id = _extract_join_stage_id(
        join_meta,
        fallback_stage_id=shuffle_stage_id,
    )
    join_num_partitions = _extract_join_num_partitions(
        join_meta,
        fallback_num_partitions=shuffle_num_partitions,
    )
    if join_num_partitions != shuffle_num_partitions:
        raise JoinContractError("join.num_partitions must match shuffle.num_partitions.")

    control = _extract_join_router_control(payload)
    if control == "done":
        partition_source = _extract_mapping_value(
            payload,
            field_name="partition_id",
        )
        if partition_source is None:
            partition_source = tags.get("partition_id")
        if partition_source is None:
            raise JoinContractError("join done frame requires partition_id in payload or tags.")
        partition_id = _parse_reducer_partition_id(
            partition_source,
            num_partitions=join_num_partitions,
        )
        key_source = _extract_mapping_value(payload, field_name="state_partition_key")
        if key_source is None:
            key_source = tags.get("state_partition_key", partition_id)
        key_fragment = _stable_key_fragment(key_source)
    else:
        key_source = _resolve_shuffle_partition_key_source(
            shuffle_meta=shuffle_meta,
            payload=payload,
            tags=tags,
            seq=seq,
            invoke_target=invoke_target,
        )
        key_fragment = _stable_key_fragment(key_source)
        partition_id = int(key_fragment, 16) % join_num_partitions

    output_tags = dict(tags)
    output_tags["partition_id"] = str(partition_id)
    output_tags["state_partition_key"] = key_fragment
    output_tags["shuffle_stage_id"] = shuffle_stage_id
    output_tags["shuffle_num_partitions"] = str(join_num_partitions)
    output_tags["join_stage_id"] = join_stage_id
    output_tags["join_num_partitions"] = str(join_num_partitions)
    output_tags["join_side"] = side

    route_node_ids = _select_shuffle_route_node_ids(
        transformation=transformation,
        partition_id=partition_id,
    )
    return OperatorEvaluation(
        outputs=(payload,),
        output_tags=(output_tags,),
        route_node_ids=route_node_ids,
    )


def _extract_join_meta(transformation: Any) -> Mapping[str, Any]:
    operator_config = getattr(transformation, "operator_config", None)
    if not isinstance(operator_config, Mapping):
        raise JoinContractError("join operator_config must be a mapping.")
    join_meta = operator_config.get("join")
    if not isinstance(join_meta, Mapping):
        raise JoinContractError("join operator_config.join must be a mapping.")
    return join_meta


def _extract_join_side(transformation: Any) -> str:
    operator_config = getattr(transformation, "operator_config", None)
    if not isinstance(operator_config, Mapping):
        raise JoinContractError("join operator_config must be a mapping.")
    side = str(operator_config.get("side", "")).strip().lower()
    if side not in {"left", "right"}:
        raise JoinContractError("join operator_config.side must be one of: left, right.")
    return side


def _extract_join_stage_id(
    join_meta: Mapping[str, Any],
    *,
    fallback_stage_id: str,
) -> str:
    stage_id = join_meta.get("stage_id", fallback_stage_id)
    if stage_id is None:
        stage_id = fallback_stage_id
    return _join_as_non_empty_str(stage_id, field_name="join.stage_id")


def _extract_join_num_partitions(
    join_meta: Mapping[str, Any],
    *,
    fallback_num_partitions: int,
) -> int:
    raw = join_meta.get("num_partitions", fallback_num_partitions)
    try:
        num_partitions = int(raw)
    except (TypeError, ValueError) as exc:
        raise JoinContractError("join.num_partitions must be an integer >= 1.") from exc
    if num_partitions < 1:
        raise JoinContractError("join.num_partitions must be an integer >= 1.")
    return num_partitions


def _extract_join_router_control(payload: Any) -> str:
    control = _extract_mapping_value(payload, field_name="control")
    if control is None:
        return "data"
    normalized = str(control or "").strip().lower()
    if not normalized:
        return "data"
    if normalized not in {"data", "done"}:
        raise JoinContractError("join.control must be one of: data, done.")
    return normalized


def _extract_mapping_value(payload: Any, *, field_name: str) -> Any:
    if not isinstance(payload, Mapping):
        return None
    return payload.get(field_name)


def _extract_shuffle_meta(transformation: Any) -> Mapping[str, Any]:
    operator_config = getattr(transformation, "operator_config", None)
    if not isinstance(operator_config, Mapping):
        raise ShuffleOperatorContractError(
            "shuffle operator_config must be a mapping.",
        )
    shuffle_meta = operator_config.get("shuffle")
    if not isinstance(shuffle_meta, Mapping):
        raise ShuffleOperatorContractError(
            "shuffle operator_config.shuffle must be a mapping.",
        )
    return shuffle_meta


def _extract_shuffle_num_partitions(shuffle_meta: Mapping[str, Any]) -> int:
    raw = shuffle_meta.get("num_partitions", 1)
    try:
        num_partitions = int(raw)
    except (TypeError, ValueError) as exc:
        raise ShuffleOperatorContractError(
            "shuffle.num_partitions must be an integer > 0.",
        ) from exc
    if num_partitions <= 0:
        raise ShuffleOperatorContractError(
            "shuffle.num_partitions must be an integer > 0.",
        )
    return num_partitions


def _extract_shuffle_stage_id(shuffle_meta: Mapping[str, Any], transformation: Any) -> str:
    stage_id = shuffle_meta.get("stage_id")
    if stage_id is None:
        stage_id = getattr(transformation, "trans_id", None)
    if not isinstance(stage_id, str) or not stage_id.strip():
        raise ShuffleOperatorContractError("shuffle.stage_id must be a non-empty string.")
    return stage_id.strip()


def _resolve_shuffle_partition_key_source(
    *,
    shuffle_meta: Mapping[str, Any],
    payload: Any,
    tags: Mapping[str, str],
    seq: int | None,
    invoke_target: Callable[[Any, Any, Mapping[str, str], int | None], Any] | None,
) -> Any:
    key_fn = shuffle_meta.get("key_fn")
    key_field = shuffle_meta.get("key_field")
    key_index = shuffle_meta.get("key_index", 0)

    if key_fn is not None and key_field is not None:
        raise ShuffleOperatorContractError(
            "shuffle.key_fn cannot be combined with shuffle.key_field.",
        )
    if key_fn is not None and key_index not in (None, 0):
        raise ShuffleOperatorContractError(
            "shuffle.key_fn cannot be combined with non-default shuffle.key_index.",
        )

    if key_fn is not None:
        if invoke_target is None:
            raise ShuffleOperatorContractError(
                "shuffle.key_fn requires invoke_target runtime support.",
            )
        return invoke_target(key_fn, payload, tags, seq)

    key_index_value = _extract_shuffle_key_index(key_index)
    key_field_value: str | None = None
    if key_field is not None:
        if not isinstance(key_field, str) or not key_field.strip():
            raise ShuffleOperatorContractError("shuffle.key_field must be a non-empty string.")
        key_field_value = key_field.strip()
    return _extract_shuffle_partition_key_from_payload(
        payload=payload,
        key_index=key_index_value,
        key_field=key_field_value,
    )


def _extract_shuffle_key_index(raw: Any) -> int:
    if raw is None:
        return 0
    try:
        key_index = int(raw)
    except (TypeError, ValueError) as exc:
        raise ShuffleOperatorContractError(
            "shuffle.key_index must be an integer >= 0.",
        ) from exc
    if key_index < 0:
        raise ShuffleOperatorContractError("shuffle.key_index must be an integer >= 0.")
    return key_index


def _extract_shuffle_partition_key_from_payload(
    *,
    payload: Any,
    key_index: int,
    key_field: str | None,
) -> Any:
    if key_field is not None:
        if not isinstance(payload, Mapping):
            raise ShuffleOperatorContractError(
                "shuffle.key_field requires mapping payload.",
            )
        if key_field not in payload:
            raise ShuffleOperatorContractError(
                f"shuffle.key_field '{key_field}' is missing in payload.",
            )
        return payload[key_field]

    if isinstance(payload, (list, tuple)):
        if key_index >= len(payload):
            raise ShuffleOperatorContractError("shuffle.key_index out of range.")
        return payload[key_index]

    if key_index != 0:
        raise ShuffleOperatorContractError("shuffle.key_index out of range.")
    return payload


def _select_shuffle_route_node_ids(
    *,
    transformation: Any,
    partition_id: int,
) -> tuple[str, ...] | None:
    downstreams = getattr(transformation, "downstreams", None)
    if downstreams is None:
        return None
    if not isinstance(downstreams, list):
        raise ShuffleOperatorContractError("transformation.downstreams must be a list.")
    if not downstreams:
        return ()

    downstream_node_ids: list[str] = []
    for index, downstream in enumerate(downstreams):
        trans_id = getattr(downstream, "trans_id", None)
        if not isinstance(trans_id, str) or not trans_id.strip():
            raise ShuffleOperatorContractError(
                f"downstreams[{index}] missing valid trans_id",
            )
        downstream_node_ids.append(trans_id)
    selected_index = partition_id % len(downstream_node_ids)
    return (downstream_node_ids[selected_index],)


def _extract_reducer_meta(transformation: Any) -> Mapping[str, Any]:
    operator_config = getattr(transformation, "operator_config", None)
    if not isinstance(operator_config, Mapping):
        raise ReducerContractError("reducer operator_config must be a mapping.")
    reducer_meta = operator_config.get("reducer")
    if not isinstance(reducer_meta, Mapping):
        raise ReducerContractError("reducer operator_config.reducer must be a mapping.")
    return reducer_meta


def _extract_reducer_stage_id(reducer_meta: Mapping[str, Any], transformation: Any) -> str:
    stage_id = reducer_meta.get("stage_id")
    if stage_id is None:
        stage_id = getattr(transformation, "trans_id", None)
    return _reducer_as_non_empty_str(stage_id, field_name="reducer.stage_id")


def _extract_reducer_mode(reducer_meta: Mapping[str, Any]) -> str:
    mode = _reducer_as_non_empty_str(
        reducer_meta.get("mode", "concat"),
        field_name="reducer.mode",
    )
    if mode not in {
        "count",
        "sum",
        "concat",
        "count_by_key",
        "sum_by_key",
        "concat_by_key",
    }:
        raise ReducerContractError(f"unsupported_reducer_mode:{mode}")
    return mode


def _extract_reducer_num_partitions(reducer_meta: Mapping[str, Any]) -> int:
    raw_value = reducer_meta.get("num_partitions", 1)
    try:
        num_partitions = int(raw_value)
    except (TypeError, ValueError) as exc:
        raise ReducerContractError("reducer.num_partitions must be an integer >= 1.") from exc
    if num_partitions < 1:
        raise ReducerContractError("reducer.num_partitions must be an integer >= 1.")
    return num_partitions


def _extract_reducer_groupby(reducer_meta: Mapping[str, Any]) -> Mapping[str, Any]:
    groupby = reducer_meta.get("groupby")
    if groupby is None:
        return {}
    if not isinstance(groupby, Mapping):
        raise ReducerContractError("reducer.groupby must be a mapping when provided.")
    return dict(groupby)


def _parse_reducer_frame(
    *,
    payload: Any,
    tags: Mapping[str, str],
    num_partitions: int,
) -> tuple[str, int, Any]:
    control = "data"
    partition_source: Any = None
    value = payload

    if isinstance(payload, Mapping) and "control" in payload:
        control = _normalize_reducer_control(payload.get("control"))
        partition_source = payload.get("partition_id")
        if control == "data":
            value = payload.get("value")
            if "value" not in payload:
                raise ReducerContractError(
                    "reducer data frame with control field requires payload.value.",
                )

    if partition_source is None:
        partition_source = tags.get("partition_id")
    if partition_source is None:
        partition_source = 0
    partition_id = _parse_reducer_partition_id(
        partition_source,
        num_partitions=num_partitions,
    )
    if control == "data":
        return control, partition_id, value
    return control, partition_id, None


def _normalize_reducer_control(value: Any) -> str:
    normalized = str(value or "").strip().lower()
    if not normalized:
        return "data"
    if normalized not in {"data", "flush", "done"}:
        raise ReducerContractError("reducer.control must be one of: data, flush, done.")
    return normalized


def _parse_reducer_partition_id(value: Any, *, num_partitions: int) -> int:
    try:
        partition_id = int(value)
    except (TypeError, ValueError) as exc:
        raise ReducerContractError("reducer.partition_id must be an integer >= 0.") from exc
    if partition_id < 0 or partition_id >= num_partitions:
        raise ReducerContractError("reducer.partition_id out of range.")
    return partition_id


def _new_reducer_accumulator(mode: str) -> Any:
    if mode in {"count_by_key", "sum_by_key", "concat_by_key"}:
        return {}
    if mode == "count":
        return 0
    if mode == "sum":
        return 0
    return []


def _accumulate_reducer_value(*, state: _ReducerState, value: Any) -> None:
    mode = state.mode
    if mode == "count":
        state.acc += 1
        return
    if mode == "sum":
        try:
            state.acc += value
        except Exception as exc:
            raise ReducerContractError("reducer.sum payload is not additive.") from exc
        return
    if mode == "concat":
        state.acc.append(value)
        return

    key = _resolve_reducer_group_key(value, groupby=state.groupby)
    if mode == "count_by_key":
        state.acc[key] = int(state.acc.get(key, 0)) + 1
        return
    if mode == "sum_by_key":
        amount = _resolve_reducer_group_sum_value(value, groupby=state.groupby)
        try:
            state.acc[key] = state.acc.get(key, 0) + amount
        except Exception as exc:
            raise ReducerContractError("reducer.sum_by_key payload value is not additive.") from exc
        return

    bucket = state.acc.setdefault(key, [])
    bucket.append(value)


def _resolve_reducer_group_key(value: Any, *, groupby: Mapping[str, Any]) -> Any:
    key_field = groupby.get("key_field")
    key_index = _parse_groupby_index(
        groupby.get("key_index", 0),
        field_name="reducer.groupby.key_index",
    )

    if isinstance(value, Mapping):
        if key_field is not None:
            field = _reducer_as_non_empty_str(
                key_field,
                field_name="reducer.groupby.key_field",
            )
            if field not in value:
                raise ReducerContractError(f"reducer.groupby.key_field '{field}' is missing in payload.")
            return value[field]
        if "key" in value:
            return value["key"]

    if isinstance(value, (list, tuple)):
        if key_index >= len(value):
            raise ReducerContractError("reducer.groupby.key_index out of range.")
        return value[key_index]

    raise ReducerContractError(
        "by-key reducer payload must be mapping(key_field/key) or sequence(key_index).",
    )


def _resolve_reducer_group_sum_value(value: Any, *, groupby: Mapping[str, Any]) -> Any:
    value_field = groupby.get("value_field")
    value_index_raw = groupby.get("value_index")
    key_index = _parse_groupby_index(
        groupby.get("key_index", 0),
        field_name="reducer.groupby.key_index",
    )

    if isinstance(value, Mapping):
        if value_field is not None:
            field = _reducer_as_non_empty_str(
                value_field,
                field_name="reducer.groupby.value_field",
            )
            if field not in value:
                raise ReducerContractError(f"reducer.groupby.value_field '{field}' is missing in payload.")
            return value[field]
        if "value" in value:
            return value["value"]

    if isinstance(value, (list, tuple)):
        if value_index_raw is not None:
            value_index = _parse_groupby_index(
                value_index_raw,
                field_name="reducer.groupby.value_index",
            )
            if value_index >= len(value):
                raise ReducerContractError("reducer.groupby.value_index out of range.")
            return value[value_index]
        if len(value) >= 2:
            value_index = 1 if key_index == 0 else 0
            if value_index < len(value):
                return value[value_index]

    raise ReducerContractError(
        "sum_by_key requires value_field/value_index or pair-like sequence payload.",
    )


def _parse_groupby_index(value: Any, *, field_name: str) -> int:
    try:
        index = int(value)
    except (TypeError, ValueError) as exc:
        raise ReducerContractError(f"{field_name} must be an integer >= 0.") from exc
    if index < 0:
        raise ReducerContractError(f"{field_name} must be an integer >= 0.")
    return index


def _materialize_reducer_snapshot(state: _ReducerState) -> Any:
    if state.mode == "count":
        return int(state.acc)
    if state.mode == "sum":
        return state.acc
    if state.mode == "concat":
        return tuple(state.acc)

    items: list[tuple[Any, Any]] = []
    for key, aggregate in state.acc.items():
        if state.mode == "concat_by_key":
            items.append((key, tuple(aggregate)))
        else:
            items.append((key, aggregate))
    items.sort(key=lambda kv: repr(kv[0]))
    return tuple(items)


def _reducer_as_non_empty_str(value: Any, *, field_name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ReducerContractError(f"{field_name} must be a non-empty string.")
    return value.strip()


def _resolve_join_reducer_contract(
    *,
    transformation: Any,
    tags: Mapping[str, str],
) -> tuple[str, int]:
    operator_config = getattr(transformation, "operator_config", None)
    join_meta: Mapping[str, Any] = {}
    if operator_config is not None:
        if not isinstance(operator_config, Mapping):
            raise JoinContractError("join-reducer operator_config must be a mapping when provided.")
        join_meta_raw = operator_config.get("join")
        if join_meta_raw is not None:
            if not isinstance(join_meta_raw, Mapping):
                raise JoinContractError("join-reducer operator_config.join must be a mapping when provided.")
            join_meta = dict(join_meta_raw)

    stage_id_raw = join_meta.get("stage_id")
    tag_stage_id = tags.get("join_stage_id")
    if stage_id_raw is not None and tag_stage_id and str(stage_id_raw).strip() != str(tag_stage_id).strip():
        raise JoinContractError("join.stage_id mismatch between operator config and tags.")
    stage_id = stage_id_raw if stage_id_raw is not None else tag_stage_id
    if stage_id is None:
        stage_id = getattr(transformation, "trans_id", None)
    stage_id_value = _join_as_non_empty_str(stage_id, field_name="join.stage_id")

    num_partitions_raw = join_meta.get("num_partitions")
    tag_num_partitions = tags.get("join_num_partitions")
    if num_partitions_raw is None:
        num_partitions_raw = tag_num_partitions
    elif tag_num_partitions is not None:
        try:
            if int(num_partitions_raw) != int(tag_num_partitions):
                raise JoinContractError("join.num_partitions mismatch between operator config and tags.")
        except (TypeError, ValueError) as exc:
            raise JoinContractError("join.num_partitions must be an integer >= 1.") from exc

    if num_partitions_raw is None:
        num_partitions_raw = 1
    try:
        num_partitions = int(num_partitions_raw)
    except (TypeError, ValueError) as exc:
        raise JoinContractError("join.num_partitions must be an integer >= 1.") from exc
    if num_partitions < 1:
        raise JoinContractError("join.num_partitions must be an integer >= 1.")
    return stage_id_value, num_partitions


def _parse_join_reducer_frame(
    *,
    payload: Any,
    tags: Mapping[str, str],
    num_partitions: int,
) -> tuple[str, str, int, Any]:
    control = _extract_join_reducer_control(payload)
    if control == "flush":
        return control, "left", 0, None

    side = _resolve_join_reducer_side(payload=payload, tags=tags)
    partition_source = _extract_mapping_value(payload, field_name="partition_id")
    if partition_source is None:
        partition_source = tags.get("partition_id")
    if partition_source is None:
        partition_source = 0
    partition_id = _parse_reducer_partition_id(
        partition_source,
        num_partitions=num_partitions,
    )

    if control == "done":
        return control, side, partition_id, None

    value = payload
    if isinstance(payload, Mapping) and "control" in payload:
        if "value" not in payload:
            raise JoinContractError(
                "join-reducer data frame with control field requires payload.value.",
            )
        value = payload.get("value")
    return control, side, partition_id, value


def _extract_join_reducer_control(payload: Any) -> str:
    control = _extract_mapping_value(payload, field_name="control")
    if control is None:
        return "data"
    normalized = str(control or "").strip().lower()
    if not normalized:
        return "data"
    if normalized not in {"data", "flush", "done"}:
        raise JoinContractError("join-reducer.control must be one of: data, flush, done.")
    return normalized


def _resolve_join_reducer_side(*, payload: Any, tags: Mapping[str, str]) -> str:
    side = _extract_mapping_value(payload, field_name="side")
    if side is None:
        side = tags.get("join_side") or tags.get("side")
    side_value = str(side or "").strip().lower()
    if side_value not in {"left", "right"}:
        raise JoinContractError("join-reducer side must be one of: left, right.")
    return side_value


def _materialize_join_reducer_snapshot(state: _JoinReducerState) -> tuple[tuple[Any, Any], ...]:
    pairs: list[tuple[Any, Any]] = []
    for partition_id in range(state.num_partitions):
        left_values = state.left.get(partition_id, [])
        right_values = state.right.get(partition_id, [])
        for left_value in left_values:
            for right_value in right_values:
                pairs.append((left_value, right_value))
    return tuple(pairs)


def _resolve_merge_reducer_contract(
    *,
    transformation: Any,
    tags: Mapping[str, str],
) -> tuple[str, str, tuple[int, ...]]:
    operator_config = getattr(transformation, "operator_config", None)
    reducer_meta: Mapping[str, Any] = {}
    if operator_config is not None:
        if not isinstance(operator_config, Mapping):
            raise MergeReducerContractError("merge-reducer operator_config must be a mapping when provided.")
        reducer_meta_raw = operator_config.get("reducer")
        if reducer_meta_raw is not None:
            if not isinstance(reducer_meta_raw, Mapping):
                raise MergeReducerContractError(
                    "merge-reducer operator_config.reducer must be a mapping when provided.",
                )
            reducer_meta = dict(reducer_meta_raw)

    stage_id_raw = reducer_meta.get("stage_id")
    tag_stage_id = tags.get("reducer_stage_id")
    if stage_id_raw is not None and tag_stage_id and str(stage_id_raw).strip() != str(tag_stage_id).strip():
        raise MergeReducerContractError("reducer.stage_id mismatch between operator config and tags.")
    stage_id = stage_id_raw if stage_id_raw is not None else tag_stage_id
    if stage_id is None:
        stage_id = getattr(transformation, "trans_id", None)
    stage_id_value = _merge_as_non_empty_str(stage_id, field_name="reducer.stage_id")

    mode_raw = reducer_meta.get("mode")
    tag_mode = tags.get("reducer_mode")
    if mode_raw is not None and tag_mode and str(mode_raw).strip() != str(tag_mode).strip():
        raise MergeReducerContractError("reducer.mode mismatch between operator config and tags.")
    mode = str(mode_raw if mode_raw is not None else (tag_mode or "concat")).strip()
    mode_value = _extract_reducer_mode({"mode": mode})

    num_partitions_raw = reducer_meta.get("num_partitions")
    tag_num_partitions = tags.get("reducer_num_partitions")
    if num_partitions_raw is None:
        num_partitions_raw = tag_num_partitions
    if num_partitions_raw is None:
        num_partitions_raw = 1
    try:
        num_partitions = int(num_partitions_raw)
    except (TypeError, ValueError) as exc:
        raise MergeReducerContractError("reducer.num_partitions must be an integer >= 1.") from exc
    if num_partitions < 1:
        raise MergeReducerContractError("reducer.num_partitions must be an integer >= 1.")

    expected_partition_ids = _parse_merge_expected_partition_ids(
        reducer_meta.get("merge_partitions"),
        num_partitions=num_partitions,
    )
    return stage_id_value, mode_value, expected_partition_ids


def _parse_merge_expected_partition_ids(raw: Any, *, num_partitions: int) -> tuple[int, ...]:
    if raw is None:
        return tuple(range(num_partitions))
    if not isinstance(raw, (list, tuple)):
        raise MergeReducerContractError("reducer.merge_partitions must be a list/tuple when provided.")
    parsed: set[int] = set()
    for item in raw:
        partition_id = _parse_non_negative_int_strict(
            item,
            field_name="reducer.merge_partitions[]",
            exc_type=MergeReducerContractError,
        )
        parsed.add(partition_id)
    if not parsed:
        raise MergeReducerContractError("reducer.merge_partitions must not be empty.")
    return tuple(sorted(parsed))


def _parse_merge_reducer_frame(
    *,
    payload: Any,
    tags: Mapping[str, str],
    expected_partition_ids: tuple[int, ...],
) -> tuple[str, int, Any]:
    control = _extract_merge_reducer_control(payload)
    if control == "flush":
        # flush does not depend on one partition; keep deterministic placeholder.
        return control, expected_partition_ids[0], None

    partition_source = _extract_mapping_value(payload, field_name="partition_id")
    if partition_source is None:
        partition_source = tags.get("partition_id")
    if partition_source is None:
        partition_source = 0
    partition_id = _parse_non_negative_int_strict(
        partition_source,
        field_name="merge-reducer.partition_id",
        exc_type=MergeReducerContractError,
    )

    if control == "done":
        return control, partition_id, None

    value = payload
    if isinstance(payload, Mapping) and "control" in payload:
        if "value" not in payload:
            raise MergeReducerContractError(
                "merge-reducer partial frame with control field requires payload.value.",
            )
        value = payload.get("value")
    return control, partition_id, value


def _extract_merge_reducer_control(payload: Any) -> str:
    control = _extract_mapping_value(payload, field_name="control")
    if control is None:
        return "partial"
    normalized = str(control or "").strip().lower()
    if not normalized:
        return "partial"
    if normalized not in {"partial", "flush", "done"}:
        raise MergeReducerContractError("merge-reducer.control must be one of: partial, flush, done.")
    return normalized


def _merge_partial_into_reducer_accumulator(*, mode: str, acc: Any, value: Any) -> Any:
    if mode == "count":
        try:
            return acc + value
        except Exception as exc:
            raise MergeReducerContractError("merge-reducer.count partial is not additive.") from exc
    if mode == "sum":
        try:
            return acc + value
        except Exception as exc:
            raise MergeReducerContractError("merge-reducer.sum partial is not additive.") from exc
    if mode == "concat":
        if isinstance(value, (list, tuple)):
            acc.extend(value)
        else:
            acc.append(value)
        return acc

    items = _extract_merge_partial_keyed_items(value)
    if mode == "count_by_key":
        for key, partial in items:
            try:
                acc[key] = acc.get(key, 0) + partial
            except Exception as exc:
                raise MergeReducerContractError(
                    "merge-reducer.count_by_key partial value is not additive.",
                ) from exc
        return acc
    if mode == "sum_by_key":
        for key, partial in items:
            try:
                acc[key] = acc.get(key, 0) + partial
            except Exception as exc:
                raise MergeReducerContractError(
                    "merge-reducer.sum_by_key partial value is not additive.",
                ) from exc
        return acc

    for key, partial in items:
        bucket = acc.setdefault(key, [])
        if isinstance(partial, (list, tuple)):
            bucket.extend(partial)
        else:
            bucket.append(partial)
    return acc


def _extract_merge_partial_keyed_items(value: Any) -> tuple[tuple[Any, Any], ...]:
    if isinstance(value, Mapping):
        return tuple((key, item_value) for key, item_value in value.items())
    if isinstance(value, (list, tuple)):
        items: list[tuple[Any, Any]] = []
        for index, item in enumerate(value):
            if not isinstance(item, (list, tuple)) or len(item) != 2:
                raise MergeReducerContractError(
                    f"merge-reducer by-key partial[{index}] must be a (key, value) pair.",
                )
            items.append((item[0], item[1]))
        return tuple(items)
    raise MergeReducerContractError(
        "merge-reducer by-key partial must be mapping or list/tuple of (key, value) pairs.",
    )


def _merge_reducer_is_complete(state: _MergeReducerState) -> bool:
    expected = set(state.expected_partition_ids)
    return expected.issubset(state.partial_partitions) and expected.issubset(state.done_partitions)


def _materialize_merge_reducer_snapshot(state: _MergeReducerState) -> Any:
    if state.mode == "count":
        return int(state.acc)
    if state.mode == "sum":
        return state.acc
    if state.mode == "concat":
        return tuple(state.acc)

    items: list[tuple[Any, Any]] = []
    for key, aggregate in state.acc.items():
        if state.mode == "concat_by_key":
            items.append((key, tuple(aggregate)))
        else:
            items.append((key, aggregate))
    items.sort(key=lambda kv: repr(kv[0]))
    return tuple(items)


def _join_as_non_empty_str(value: Any, *, field_name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise JoinContractError(f"{field_name} must be a non-empty string.")
    return value.strip()


def _merge_as_non_empty_str(value: Any, *, field_name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise MergeReducerContractError(f"{field_name} must be a non-empty string.")
    return value.strip()


def _parse_non_negative_int_strict(
    value: Any,
    *,
    field_name: str,
    exc_type: type[ReducerContractError],
) -> int:
    try:
        resolved = int(value)
    except (TypeError, ValueError) as exc:
        raise exc_type(f"{field_name} must be an integer >= 0.") from exc
    if resolved < 0:
        raise exc_type(f"{field_name} must be an integer >= 0.")
    return resolved


def _resolve_process_program_ref(transformation: Any) -> FlowProgramRef:
    target = getattr(transformation, "target", None)
    ref = _coerce_flow_program_ref(target)
    if ref is not None:
        return ref

    pipe_flow = getattr(transformation, "pipe_flow", None)
    ref = _coerce_flow_program_ref(pipe_flow)
    if ref is not None:
        return ref

    raise ProcessTargetResolutionError(
        "process operator target must provide canonical program_uri/program_rev metadata.",
    )


def _coerce_flow_program_ref(raw: Any) -> FlowProgramRef | None:
    if raw is None:
        return None

    program_uri = None
    program_rev = None
    if isinstance(raw, Mapping):
        program_uri = raw.get("program_uri")
        program_rev = raw.get("program_rev")
    else:
        program_uri = getattr(raw, "program_uri", None)
        program_rev = getattr(raw, "program_rev", None)

    if isinstance(program_uri, str) and program_uri.strip() and isinstance(program_rev, str) and program_rev.strip():
        return FlowProgramRef(program_uri=program_uri.strip(), program_rev=program_rev.strip())

    return None


def _extract_stateful_process_meta(transformation: Any) -> Mapping[str, Any]:
    operator_config = getattr(transformation, "operator_config", None)
    if not isinstance(operator_config, Mapping):
        raise StatefulProcessContractError(
            "stateful_process operator_config must be a mapping.",
        )
    process_meta = operator_config.get("stateful_process")
    if not isinstance(process_meta, Mapping):
        raise StatefulProcessContractError(
            "stateful_process operator_config.stateful_process must be a mapping.",
        )
    return process_meta


def _extract_state_stage_id(process_meta: Mapping[str, Any], transformation: Any) -> str:
    stage_id = process_meta.get("stage_id")
    if stage_id is None:
        stage_id = getattr(transformation, "trans_id", None)
    return _as_non_empty_str(
        stage_id,
        field_name="stateful_process.stage_id",
    )


def _extract_state_namespace(process_meta: Mapping[str, Any], *, stage_id: str) -> str:
    namespace = process_meta.get("state_namespace")
    if namespace is None:
        namespace = f"stateful_process:{stage_id}"
    return _as_non_empty_str(
        namespace,
        field_name="stateful_process.state_namespace",
    )


def _extract_state_ttl_seconds(process_meta: Mapping[str, Any]) -> float | None:
    ttl_value = process_meta.get("state_ttl_s")
    if ttl_value is None:
        return None
    return _as_positive_float(
        ttl_value,
        field_name="stateful_process.state_ttl_s",
    )


def _resolve_state_key_source(
    *,
    payload: Any,
    tags: Mapping[str, str],
    state_spec: Any,
) -> Any:
    if state_spec is not None and not isinstance(state_spec, Mapping):
        raise StatefulProcessContractError(
            "stateful_process.state_spec must be a mapping when provided.",
        )

    if isinstance(state_spec, Mapping):
        key_field = state_spec.get("key_field")
        if key_field is not None:
            key_field_name = _as_non_empty_str(
                key_field,
                field_name="stateful_process.state_spec.key_field",
            )
            if not isinstance(payload, Mapping):
                raise StatefulProcessContractError(
                    "stateful_process.state_spec.key_field requires mapping payload.",
                )
            if key_field_name not in payload:
                raise StatefulProcessContractError(
                    f"stateful_process.state_spec.key_field '{key_field_name}' is missing in payload.",
                )
            return payload[key_field_name]

        key_index = state_spec.get("key_index")
        if key_index is not None:
            if not isinstance(payload, (list, tuple)):
                raise StatefulProcessContractError(
                    "stateful_process.state_spec.key_index requires list/tuple payload.",
                )
            try:
                resolved_index = int(key_index)
            except (TypeError, ValueError) as exc:
                raise StatefulProcessContractError(
                    "stateful_process.state_spec.key_index must be an integer >= 0.",
                ) from exc
            if resolved_index < 0 or resolved_index >= len(payload):
                raise StatefulProcessContractError(
                    "stateful_process.state_spec.key_index out of range.",
                )
            return payload[resolved_index]

    fallback_key = tags.get("state_partition_key")
    if isinstance(fallback_key, str) and fallback_key.strip():
        return fallback_key.strip()
    return 0


def _stable_key_fragment(value: Any) -> str:
    normalized = _normalize_state_key_component(value)
    encoded = json.dumps(
        normalized,
        sort_keys=True,
        separators=(",", ":"),
        default=str,
    )
    return hashlib.sha1(encoded.encode("utf-8")).hexdigest()


def _normalize_state_key_component(value: Any) -> Any:
    if isinstance(value, Mapping):
        items: list[list[Any]] = []
        for key, item_value in value.items():
            items.append(
                [
                    _normalize_state_key_component(key),
                    _normalize_state_key_component(item_value),
                ]
            )
        items.sort(key=lambda item: repr(item[0]))
        return {"dict": items}

    if isinstance(value, (list, tuple)):
        return {"seq": [_normalize_state_key_component(item) for item in value]}

    if isinstance(value, set):
        items = [_normalize_state_key_component(item) for item in value]
        items.sort(key=repr)
        return {"set": items}

    if isinstance(value, (bytes, bytearray)):
        return {"bytes": bytes(value).hex()}

    try:
        hash(value)
    except TypeError:
        return {"repr": repr(value)}
    return value


def _as_non_empty_str(value: Any, *, field_name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise StatefulProcessContractError(f"{field_name} must be a non-empty string.")
    return value.strip()


def _as_positive_float(value: Any, *, field_name: str) -> float:
    try:
        resolved = float(value)
    except (TypeError, ValueError) as exc:
        raise StatefulProcessContractError(f"{field_name} must be a float > 0.") from exc
    if resolved <= 0:
        raise StatefulProcessContractError(f"{field_name} must be a float > 0.")
    return resolved


__all__ = [
    "OperatorEvaluation",
    "OperatorExecutionError",
    "UnsupportedOperatorTypeError",
    "ProcessTargetResolutionError",
    "ReducerContractError",
    "JoinContractError",
    "MergeReducerContractError",
    "ReducerRuntime",
    "JoinReducerRuntime",
    "MergeReducerRuntime",
    "ShuffleOperatorContractError",
    "StatefulProcessContractError",
    "StatefulProcessRuntime",
    "StatefulStateHandle",
    "evaluate_operator",
]
