import warnings
from functools import partial
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, TypeVar

if TYPE_CHECKING:
    import pyarrow.fs

from ray.data._internal.execution.execution_callback import ExecutionCallback
from ray.data._internal.execution.interfaces import PhysicalOperator
from ray.data._internal.execution.operators.aggregate_num_rows import (
    AggregateNumRows,
)
from ray.data._internal.execution.operators.hash_shuffle_v2 import (
    _SHUFFLE_MAP_RUNTIME_ENV,
    _make_hash_partition_fn,
)
from ray.data._internal.execution.operators.input_data_buffer import (
    InputDataBuffer,
)
from ray.data._internal.execution.operators.join import (
    JoinOperator,
    _make_join_reduce_fn,
)
from ray.data._internal.execution.operators.limit_operator import LimitOperator
from ray.data._internal.execution.operators.mix_operator import MixOperator
from ray.data._internal.execution.operators.output_splitter import OutputSplitter
from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (
    ShuffleMapOp,
)
from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (
    ShuffleReduceOp,
)
from ray.data._internal.execution.operators.union_operator import UnionOperator
from ray.data._internal.execution.operators.zip_operator import ZipOperator
from ray.data._internal.logical.interfaces import (
    LogicalOperator,
    LogicalPlan,
    PhysicalPlan,
)
from ray.data._internal.logical.operators import (
    AbstractAllToAll,
    AbstractFrom,
    AbstractUDFMap,
    Count,
    Download,
    Filter,
    InputData,
    Join,
    JoinType,
    Limit,
    ListFiles,
    Mix,
    Project,
    Read,
    ReadFiles,
    StreamingRepartition,
    StreamingSplit,
    Union,
    Write,
    Zip,
)
from ray.data._internal.planner.checkpoint import (
    plan_read_files_op_with_checkpoint_filter,
    plan_read_op_with_checkpoint_filter,
    plan_write_op_with_checkpoint_writer,
)
from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
from ray.data._internal.planner.plan_download_op import plan_download_op
from ray.data._internal.planner.plan_list_files_op import plan_list_files_op
from ray.data._internal.planner.plan_read_files_op import plan_read_files_op
from ray.data._internal.planner.plan_read_op import plan_read_op
from ray.data._internal.planner.plan_udf_map_op import (
    plan_filter_op,
    plan_project_op,
    plan_streaming_repartition_op,
    plan_udf_map_op,
)
from ray.data._internal.planner.plan_write_op import plan_write_op
from ray.data._internal.usage import create_usage_callback
from ray.data.checkpoint.load_checkpoint_callback import LoadCheckpointCallback
from ray.data.context import DataContext
from ray.data.datasource.file_datasink import _FileDatasink

LogicalOperatorType = TypeVar("LogicalOperatorType", bound=LogicalOperator)
PlanLogicalOpFn = Callable[
    [LogicalOperatorType, List[PhysicalOperator], DataContext], PhysicalOperator
]


def plan_input_data_op(
    logical_op: InputData,
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
) -> PhysicalOperator:
    """Get the corresponding DAG of physical operators for InputData."""
    assert len(physical_children) == 0

    return InputDataBuffer(
        data_context,
        input_data=logical_op.input_data,
    )


def plan_from_op(
    op: AbstractFrom,
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
) -> PhysicalOperator:
    assert len(physical_children) == 0
    return InputDataBuffer(data_context, op.input_data)


def plan_zip_op(_, physical_children, data_context):
    assert len(physical_children) >= 2
    return ZipOperator(data_context, *physical_children)


def plan_mix_op(logical_op, physical_children, data_context):
    assert len(physical_children) >= 1
    return MixOperator(
        data_context,
        *physical_children,
        weights=logical_op.weights,
        stopping_condition=logical_op.stopping_condition,
    )


def plan_union_op(_, physical_children, data_context):
    assert len(physical_children) >= 2
    return UnionOperator(data_context, *physical_children)


def plan_limit_op(logical_op, physical_children, data_context):
    assert len(physical_children) == 1
    return LimitOperator(logical_op.limit, physical_children[0], data_context)


def plan_count_op(logical_op, physical_children, data_context):
    assert len(physical_children) == 1
    return AggregateNumRows(
        [physical_children[0]], data_context, column_name=Count.COLUMN_NAME
    )


def _plan_join_shuffle_v2(
    logical_op: Join,
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
) -> PhysicalOperator:
    left_keys = list(logical_op.left_key_columns)
    right_keys = list(logical_op.right_key_columns)
    num_partitions = logical_op.num_partitions
    join_type = JoinType(logical_op.join_type)

    left_map = ShuffleMapOp(
        physical_children[0],
        data_context,
        num_partitions=num_partitions,
        partition_fn=_make_hash_partition_fn(left_keys, num_partitions),
        map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
        name=f"JoinShuffleMapLeft(keys={tuple(left_keys)}, parts={num_partitions})",
    )
    right_map = ShuffleMapOp(
        physical_children[1],
        data_context,
        num_partitions=num_partitions,
        partition_fn=_make_hash_partition_fn(right_keys, num_partitions),
        map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
        name=f"JoinShuffleMapRight(keys={tuple(right_keys)}, parts={num_partitions})",
    )

    reduce_fn = _make_join_reduce_fn(
        join_type=join_type,
        left_key_col_names=tuple(left_keys),
        right_key_col_names=tuple(right_keys),
        left_columns_suffix=logical_op.left_columns_suffix,
        right_columns_suffix=logical_op.right_columns_suffix,
        left_schema=logical_op.input_dependencies[0].infer_schema(),
        right_schema=logical_op.input_dependencies[1].infer_schema(),
    )
    return ShuffleReduceOp(
        [left_map, right_map],
        data_context,
        num_partitions=num_partitions,
        reduce_fn=reduce_fn,
        disallow_block_splitting=False,
        reduce_ray_remote_args=logical_op.aggregator_ray_remote_args,
        name=f"JoinShuffleReduce(num_partitions={num_partitions})",
    )


def plan_join_op(
    logical_op: Join,
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
) -> PhysicalOperator:
    assert len(physical_children) == 2
    if data_context.use_hash_shuffle_v2:
        return _plan_join_shuffle_v2(logical_op, physical_children, data_context)
    return JoinOperator(
        data_context=data_context,
        left_input_op=physical_children[0],
        right_input_op=physical_children[1],
        join_type=logical_op.join_type,
        left_key_columns=logical_op.left_key_columns,
        right_key_columns=logical_op.right_key_columns,
        left_columns_suffix=logical_op.left_columns_suffix,
        right_columns_suffix=logical_op.right_columns_suffix,
        num_partitions=logical_op.num_outputs,
        partition_size_hint=logical_op.partition_size_hint,
        aggregator_ray_remote_args_override=logical_op.aggregator_ray_remote_args,
    )


def plan_streaming_split_op(
    logical_op: StreamingSplit,
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
):
    assert len(physical_children) == 1
    return OutputSplitter(
        physical_children[0],
        n=logical_op.num_splits,
        equal=logical_op.equal,
        data_context=data_context,
        locality_hints=logical_op.locality_hints,
    )


class Planner:
    """The planner to convert optimized logical to physical operators.

    Note that planner is only doing operators conversion. Physical optimization work is
    done by physical optimizer.
    """

    _DEFAULT_PLAN_FNS = {
        Read: plan_read_op,
        ReadFiles: plan_read_files_op,
        ListFiles: plan_list_files_op,
        InputData: plan_input_data_op,
        Write: plan_write_op,
        AbstractFrom: plan_from_op,
        Filter: plan_filter_op,
        AbstractUDFMap: plan_udf_map_op,
        AbstractAllToAll: plan_all_to_all_op,
        Mix: plan_mix_op,
        Union: plan_union_op,
        Zip: plan_zip_op,
        Limit: plan_limit_op,
        Count: plan_count_op,
        Project: plan_project_op,
        StreamingRepartition: plan_streaming_repartition_op,
        Join: plan_join_op,
        StreamingSplit: plan_streaming_split_op,
        Download: plan_download_op,
    }
    # Operators that support checkpoint filtering. Subclasses can override.
    _CHECKPOINT_FILTER_OPS = (Read, ReadFiles)

    def __init__(self):
        self._supports_checkpointing = False
        self._plan_fns_for_checkpointing = {}

    def plan(
        self, logical_plan: LogicalPlan
    ) -> Tuple[PhysicalPlan, List["ExecutionCallback"]]:
        """Convert logical to physical operators recursively in post-order."""
        checkpoint_config = logical_plan.context.checkpoint_config

        callbacks = [cls() for cls in logical_plan.context.execution_callback_classes]
        callbacks.append(create_usage_callback(logical_plan))

        if checkpoint_config is not None and self._check_supports_checkpointing(
            logical_plan
        ):
            self._supports_checkpointing = True
            data_file_dir, data_file_fs = self._get_data_file_info(logical_plan)

            checkpoint_callback = self._create_checkpoint_callback(
                checkpoint_config,
            )

            callbacks.append(checkpoint_callback)

            # Dynamically set the plan functions for checkpointing because they
            # need to a reference to the checkpoint ref.
            self._plan_fns_for_checkpointing = self._get_plan_fns_for_checkpointing(
                data_file_dir, data_file_fs
            )

        elif checkpoint_config is not None:
            assert not self._check_supports_checkpointing(logical_plan)
            warnings.warn(
                "You've enabled checkpointing, but the logical plan doesn't support "
                "checkpointing. Checkpointing will be disabled."
            )
        physical_dag, op_map = self._plan_recursively(
            logical_plan.dag, logical_plan.context
        )
        physical_plan = PhysicalPlan(physical_dag, op_map, logical_plan.context)
        return physical_plan, callbacks

    def get_plan_fn(self, logical_op: LogicalOperator) -> PlanLogicalOpFn:
        if self._supports_checkpointing:
            assert self._plan_fns_for_checkpointing
            plan_fn = find_plan_fn(logical_op, self._plan_fns_for_checkpointing)
            if plan_fn is not None:
                return plan_fn

        plan_fn = find_plan_fn(logical_op, self._DEFAULT_PLAN_FNS)
        if plan_fn is not None:
            return plan_fn

        raise ValueError(
            f"Found unknown logical operator during planning: {logical_op}"
        )

    def _plan_recursively(
        self, logical_op: LogicalOperator, data_context: DataContext
    ) -> Tuple[PhysicalOperator, Dict[LogicalOperator, PhysicalOperator]]:
        """Plan a logical operator and its input dependencies recursively.

        Args:
            logical_op: The logical operator to plan.
            data_context: The data context.

        Returns:
            A tuple of the physical operator corresponding to the logical operator, and
            a mapping from physical to logical operators.
        """
        op_map: Dict[PhysicalOperator, LogicalOperator] = {}

        # Plan the input dependencies first.
        physical_children = []
        for child in logical_op.input_dependencies:
            physical_child, child_op_map = self._plan_recursively(child, data_context)
            physical_children.append(physical_child)
            op_map.update(child_op_map)

        plan_fn = self.get_plan_fn(logical_op)
        # We will call `set_logical_operators()` in the following for-loop,
        # no need to do it here.
        physical_op = plan_fn(logical_op, physical_children, data_context)

        # Traverse up the DAG, and set the mapping from physical to logical operators.
        # At this point, all physical operators without logical operators set
        # must have been created by the current logical operator.
        queue = [physical_op]
        while queue:
            curr_physical_op = queue.pop()
            if curr_physical_op._logical_operators:
                continue

            curr_physical_op.set_logical_operators(logical_op)
            # Add this operator to the op_map so optimizer can find it
            op_map[curr_physical_op] = logical_op
            queue.extend(curr_physical_op.input_dependencies)

        # Also add the final operator (in case the loop didn't catch it)
        op_map[physical_op] = logical_op
        return physical_op, op_map

    def _create_checkpoint_callback(
        self,
        checkpoint_config,
    ) -> LoadCheckpointCallback:
        """Factory method to create the LoadCheckpointCallback.

        Subclasses can override this to use a different callback implementation.
        """
        return LoadCheckpointCallback(
            checkpoint_config,
        )

    @staticmethod
    def _get_data_file_info(logical_plan: LogicalPlan):
        """Extract the data file directory and filesystem from the Write op's datasink.

        Returns (path, filesystem) for file-based datasinks, or (None, None)
        for non-file datasinks.
        """
        last_op = logical_plan.dag
        if isinstance(last_op, Write):
            datasink = last_op.datasink_or_legacy_datasource
            if isinstance(datasink, _FileDatasink):
                return datasink.unresolved_path, datasink.filesystem
        return None, None

    def _get_plan_fns_for_checkpointing(
        self,
        data_file_dir: Optional[str] = None,
        data_file_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
    ) -> Dict[Type[LogicalOperator], PlanLogicalOpFn]:
        plan_fns = {
            Read: partial(
                plan_read_op_with_checkpoint_filter, data_file_dir, data_file_filesystem
            ),
            ReadFiles: partial(
                plan_read_files_op_with_checkpoint_filter,
                data_file_dir,
                data_file_filesystem,
            ),
            Write: plan_write_op_with_checkpoint_writer,
        }
        return plan_fns

    def _check_supports_checkpointing(self, logical_plan: LogicalPlan) -> bool:
        """Check if the logical plan supports checkpointing.

        Subclasses can override _CHECKPOINT_FILTER_OPS to support more operators.
        """
        if not isinstance(logical_plan.dag, (Write, StreamingSplit)):
            return False

        def _all_paths_contain_checkpoint_filter(op: LogicalOperator) -> bool:
            if isinstance(op, self._CHECKPOINT_FILTER_OPS):
                return True
            return all(
                _all_paths_contain_checkpoint_filter(input_dep)
                for input_dep in op.input_dependencies
            )

        return _all_paths_contain_checkpoint_filter(logical_plan.dag)


def find_plan_fn(
    logical_op: LogicalOperator, plan_fns: Dict[Type[LogicalOperator], PlanLogicalOpFn]
) -> Optional[PlanLogicalOpFn]:
    """Find the plan function for a logical operator.

    This function goes through the plan functions in order and returns the first one
    that is an instance of the logical operator type.

    Args:
        logical_op: The logical operator to find the plan function for.
        plan_fns: The dictionary of plan functions.

    Returns:
        The plan function for the logical operator, or None if no plan function is
        found.
    """
    # TODO: This implementation doesn't account for type hierarchies conflicts or
    # multiple inheritance.
    for op_type, plan_fn in plan_fns.items():
        if isinstance(logical_op, op_type):
            return plan_fn
    return None