# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from __future__ import annotations

from typing import Callable

import numpy as np

from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_with_registered_key_functions
from .._cccl_interop import (
    get_value_type,
    set_cccl_iterator_state,
    to_cccl_value_state,
)
from .._serialization import BUILD_RESULTS, ITER, OP, VALUE, Serializable
from .._utils.protocols import get_data_pointer, get_dtype, validate_and_get_stream
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..determinism import Determinism
from ..op import OpAdapter, make_op_adapter
from ..typing import (
    DeviceArrayLike,
    GpuStruct,
    IteratorBase,
    IteratorT,
    Operator,
    _Struct,
)


class _Reduce(Serializable):
    __slots__ = [
        "d_in_cccl",
        "d_out_cccl",
        "h_init_cccl",
        "op_cccl",
        "build_results",
        "loaded_build_result",
        "device_reduce_fn",
    ]

    __serialization_schema__ = (
        ("d_in_cccl", ITER),
        ("d_out_cccl", ITER),
        ("op_cccl", OP),
        ("h_init_cccl", VALUE),
        ("build_results", BUILD_RESULTS(_bindings.DeviceReduceBuildResult)),
    )

    # TODO: constructor shouldn't require concrete `d_in`, `d_out`:
    def __init__(
        self,
        d_in: DeviceArrayLike | IteratorT,
        d_out: DeviceArrayLike | IteratorT,
        op: OpAdapter,
        h_init: np.ndarray | GpuStruct,
        determinism: Determinism,
        compute_capability=None,
    ):
        self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
        self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
        self.h_init_cccl = cccl.to_cccl_value(h_init)

        # Compile the op with value types
        value_type = get_value_type(h_init)
        self.op_cccl = op.compile((value_type, value_type), value_type)

        # loaded_build_result / device_reduce_fn are bound lazily on the first
        # __call__ (see _bind_device_reduce_fn); nothing is loaded here.
        self.build_results = cccl.build_for_ccs(
            _bindings.DeviceReduceBuildResult,
            self.d_in_cccl,
            self.d_out_cccl,
            self.op_cccl,
            self.h_init_cccl,
            determinism,
            compute_capability=compute_capability,
        )

    def _bind_device_reduce_fn(self) -> None:
        # Derived from the loaded build result (not serialized); bound at __call__
        # once resolve_build_result picks + loads the current device's build result.
        if (
            Determinism(self.loaded_build_result.determinism)
            is Determinism.NOT_GUARANTEED
        ):
            self.device_reduce_fn = self.loaded_build_result.compute_nondeterministic
        else:
            self.device_reduce_fn = self.loaded_build_result.compute

    def __call__(
        self,
        *,
        temp_storage,
        d_in,
        d_out,
        num_items: int,
        op: Callable | OpAdapter,
        h_init: np.ndarray | GpuStruct,
        stream=None,
    ):
        # Select (and lazily load) the current device's build result, then bind the
        # derived compute fn from it.
        self.loaded_build_result = cccl.resolve_build_result(self.build_results)
        self._bind_device_reduce_fn()

        set_cccl_iterator_state(self.d_in_cccl, d_in)
        set_cccl_iterator_state(self.d_out_cccl, d_out)

        # Update op state for stateful ops
        op_adapter = make_op_adapter(op)
        self.op_cccl.state = op_adapter.get_state()

        self.h_init_cccl.state = to_cccl_value_state(h_init)

        stream_handle = validate_and_get_stream(stream)

        if temp_storage is None:
            temp_storage_bytes = 0
            d_temp_storage = 0
        else:
            temp_storage_bytes = temp_storage.nbytes
            d_temp_storage = get_data_pointer(temp_storage)

        temp_storage_bytes = self.device_reduce_fn(
            d_temp_storage,
            temp_storage_bytes,
            self.d_in_cccl,
            self.d_out_cccl,
            num_items,
            self.op_cccl,
            self.h_init_cccl,
            stream_handle,
        )
        return temp_storage_bytes


@cache_with_registered_key_functions
def make_reduce_into(
    *,
    d_in: DeviceArrayLike | IteratorT,
    d_out: DeviceArrayLike | IteratorT,
    op: Operator,
    h_init: np.ndarray | GpuStruct,
    **kwargs,
):
    """Computes a device-wide reduction using the specified binary ``op`` and initial value ``init``.

    Example:
        Below, ``make_reduce_into`` is used to create a reduction object that can be reused.

        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/reduction/reduce_object.py
            :language: python
            :start-after: # example-begin


    Args:
        d_in: Device array or iterator containing the input sequence of data items
        d_out: Device array (of size 1) or iterator that will store the result of the reduction
        op: Binary operator to apply.
            The signature is ``(T, T) -> T``, where ``T`` is
            the data type of the initial value ``h_init``.
        init: Numpy array storing initial value of the reduction
        compute_capability: Compute capability, or list of capabilities, to
            build for ahead of time. Accepts a packed int (e.g. ``90``), a
            ``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
            thereof. When ``None`` (the default), the current device's
            architecture is used.

    Returns:
        A callable object that can be used to perform the reduction
    """
    try:
        accum_dtype = get_dtype(h_init)
    except (AttributeError, TypeError) as e:
        raise TypeError(
            "Could not determine accumulator dtype from h_init; "
            "expected numpy array or object with .dtype"
        ) from e

    # Validate d_in and d_out if they are device arrays (iterators may not expose
    # dtype reliably here). Additionally, only require equality of dtypes for
    # struct objects; mixed scalar dtypes (e.g. int8 input with int64 output)
    # is acceptable
    if isinstance(h_init, _Struct):
        for arr, name in ((d_in, "input"), (d_out, "output")):
            if isinstance(arr, IteratorBase):
                continue

            dtype = get_dtype(arr)
            if dtype != accum_dtype:
                raise TypeError(
                    f"reduce_into dtype mismatch: {name} dtype {dtype} != "
                    f"accumulator dtype {accum_dtype}. "
                    f"Ensure {name} elements and h_init have identical dtype to "
                    "avoid truncation or misinterpretation."
                )

    op_adapter = make_op_adapter(op)
    return _Reduce(
        d_in,
        d_out,
        op_adapter,
        h_init,
        kwargs.get("determinism", Determinism.RUN_TO_RUN),
        compute_capability=kwargs.get("compute_capability"),
    )


def reduce_into(
    *,
    d_in: DeviceArrayLike | IteratorT,
    d_out: DeviceArrayLike | IteratorT,
    num_items: int,
    op: Operator,
    h_init: np.ndarray | GpuStruct,
    stream=None,
    **kwargs,
):
    """
    Performs device-wide reduction.

    This function automatically handles temporary storage allocation and execution.

    Example:
        Below, ``reduce_into`` is used to compute the sum of a sequence of integers.

        .. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/reduction/sum_reduction.py
            :language: python
            :start-after: # example-begin

    Args:
        d_in: Device array or iterator containing the input sequence of data items
        d_out: Device array or iterator to store the result of the reduction
        num_items: Number of items to reduce
        op: Binary operator to apply.
            The signature is ``(T, T) -> T``, where ``T`` is
            the data type of the initial value ``h_init``.
        h_init: Initial value for the reduction
        stream: CUDA stream for the operation (optional)
    """
    reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=op, h_init=h_init, **kwargs)
    tmp_storage_bytes = reducer(
        temp_storage=None,
        d_in=d_in,
        d_out=d_out,
        num_items=num_items,
        op=op,
        h_init=h_init,
        stream=stream,
    )
    tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
    reducer(
        temp_storage=tmp_storage,
        d_in=d_in,
        d_out=d_out,
        num_items=num_items,
        op=op,
        h_init=h_init,
        stream=stream,
    )
