import itertools
import uuid
from typing import TYPE_CHECKING, Callable, Iterator, List, Optional, Union

from ray.data._internal.execution.interfaces import PhysicalOperator
from ray.data._internal.execution.interfaces.task_context import TaskContext
from ray.data._internal.execution.operators.map_operator import MapOperator
from ray.data._internal.execution.operators.map_transformer import (
    BlockMapTransformFn,
    MapTransformer,
)
from ray.data.block import Block, BlockAccessor
from ray.data.context import DataContext
from ray.data.datasource.datasink import Datasink
from ray.data.datasource.datasource import Datasource

if TYPE_CHECKING:
    from ray.data._internal.logical.operators import Write

WRITE_UUID_KWARG_NAME = "write_uuid"
# Key for storing pending checkpoint paths for commit phase
PENDING_CHECKPOINTS_KWARG_NAME = "_pending_checkpoints"


def generate_write_fn(
    datasink_or_legacy_datasource: Union[Datasink, Datasource], **write_args
) -> Callable[[Iterator[Block], TaskContext], Iterator[Block]]:
    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
        """Writes the blocks to the given datasink or legacy datasource.

        Outputs the original blocks to be written."""
        # Create a copy of the iterator, so we can return the original blocks.
        it1, it2 = itertools.tee(blocks, 2)
        if isinstance(datasink_or_legacy_datasource, Datasink):
            ctx.kwargs["_datasink_write_return"] = datasink_or_legacy_datasource.write(
                it1, ctx
            )
        else:
            datasink_or_legacy_datasource.write(it1, ctx, **write_args)

        return it2

    return fn


def generate_collect_write_stats_fn() -> BlockMapTransformFn:
    # If the write op succeeds, the resulting Dataset is a list of
    # one Block which contain stats/metrics about the write.
    # Otherwise, an error will be raised. The Datasource can handle
    # execution outcomes with `on_write_complete()`` and `on_write_failed()``.
    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
        """Handles stats collection for block writes."""
        block_accessors = [BlockAccessor.for_block(block) for block in blocks]
        total_num_rows = sum(ba.num_rows() for ba in block_accessors)
        total_size_bytes = sum(ba.size_bytes() for ba in block_accessors)

        # NOTE: Write tasks can return anything, so we need to wrap it in a valid block
        # type.
        import pandas as pd

        block = pd.DataFrame(
            {
                "num_rows": [total_num_rows],
                "size_bytes": [total_size_bytes],
                "write_return": [ctx.kwargs.get("_datasink_write_return", None)],
            }
        )
        return iter([block])

    return BlockMapTransformFn(
        fn,
        is_udf=False,
        disable_block_shaping=True,
    )


def plan_write_op(
    op: "Write",
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
) -> PhysicalOperator:
    collect_stats_fn = generate_collect_write_stats_fn()

    return _plan_write_op_internal(
        op,
        physical_children,
        data_context,
        post_transformations=[collect_stats_fn],
    )


def _plan_write_op_internal(
    op: "Write",
    physical_children: List[PhysicalOperator],
    data_context: DataContext,
    post_transformations: List[BlockMapTransformFn],
    pre_transformations: Optional[List[BlockMapTransformFn]] = None,
) -> PhysicalOperator:
    """Plan a write operation with optional pre and post write transformations.

    Args:
        op: The write operator.
        physical_children: The physical children operators.
        data_context: The data context.
        post_transformations: Transformations to run AFTER the write.
        pre_transformations: Transformations to run BEFORE the write.
            Useful for 2-phase commit where pending checkpoint is written first.

    Returns:
        The physical operator for the write operation.
    """
    assert len(physical_children) == 1
    input_physical_dag = physical_children[0]

    datasink = op.datasink_or_legacy_datasource
    write_fn = generate_write_fn(datasink, **op.write_args)

    # Build transform chain: pre_write -> write -> post_write
    pre_transforms = pre_transformations or []
    write_transform = BlockMapTransformFn(
        write_fn,
        is_udf=False,
        # NOTE: No need for block-shaping
        disable_block_shaping=True,
    )
    transform_fns = pre_transforms + [write_transform] + post_transformations

    map_transformer = MapTransformer(transform_fns)

    # Set up on_start callback for datasinks.
    # This allows on_write_start to receive the schema from the first input bundle,
    # enabling schema-dependent initialization (e.g., Iceberg schema evolution).
    on_start = None
    if isinstance(datasink, Datasink):
        on_start = datasink.on_write_start

    map_op = MapOperator.create(
        map_transformer,
        input_physical_dag,
        data_context,
        name="Write",
        # Add a UUID to write tasks to prevent filename collisions. This a UUID for the
        # overall write operation, not the individual write tasks.
        map_task_kwargs={WRITE_UUID_KWARG_NAME: uuid.uuid4().hex},
        ray_remote_args=op.ray_remote_args,
        min_rows_per_bundle=op.min_rows_per_bundled_input,
        compute_strategy=op.compute,
        on_start=on_start,
    )

    return map_op