from __future__ import annotations

from typing import TYPE_CHECKING

from flownet.engine.operators.base import OperatorHandler
from flownet.engine.operators.router import RouterOperatorHandler
from flownet.engine.operators.shuffle import ShuffleRouterOperatorHandler
from flownet.engine.operators.stateless import FilterOperatorHandler, PipeOperatorHandler
from flownet.engine.operators.stateful_process import StatefulProcessOperatorHandler
from flownet.engine.operators.reducers import (
    JoinReducerOperatorHandler,
    MergeReducerOperatorHandler,
    ReducerOperatorHandler,
)

if TYPE_CHECKING:
    from flownet.engine.operator_registry import OperatorRegistry


def register_default_operators(registry: "OperatorRegistry") -> "OperatorRegistry":
    default_handler = registry.default_handler
    registry.register("map", default_handler)
    registry.register("flatmap", default_handler)
    registry.register("tap", default_handler)
    registry.register("filter", FilterOperatorHandler())
    registry.register("pipe", PipeOperatorHandler())
    registry.register_internal("sys/router", RouterOperatorHandler())
    registry.register_internal("sys/shuffle", ShuffleRouterOperatorHandler())
    registry.register_internal("sys/stateful-process", StatefulProcessOperatorHandler())
    registry.register_internal("sys/reducer", ReducerOperatorHandler())
    registry.register_internal("sys/join", ShuffleRouterOperatorHandler())
    registry.register_internal("sys/join-reducer", JoinReducerOperatorHandler())
    registry.register_internal("sys/merge-reducer", MergeReducerOperatorHandler())
    return registry


__all__ = ["register_default_operators"]
