import awkward as ak
import awkward as ak
import numpy as np
import cupy as cp
from cuda.compute import ZipIterator, PermutationIterator
import nvtx

from _segment_algorithms import segment_sizes, select_segments, segmented_select, transform_segments


@nvtx.annotate("empty_like")
def empty_like(array, kind="empty"):
    # Use low-level API to avoid dispatch and from_buffers overhead
    if isinstance(array, ak.Array):
        layout = array.layout
    elif hasattr(array, 'layout'):
        layout = array.layout
    elif isinstance(array, ak.contents.Content):
        layout = array
    else:
        layout = ak.to_layout(array)

    # Recursively copy the layout tree, allocating empty buffers for data
    def copy_with_empty_buffers(content):
        backend = content._backend
        xp = backend.nplike

        if isinstance(content, ak.contents.NumpyArray):
            # Allocate empty data buffer
            empty_data = xp.empty(content.data.shape, dtype=content.data.dtype)
            return ak.contents.NumpyArray(
                empty_data,
                parameters=content._parameters,
                backend=backend
            )
        elif isinstance(content, ak.contents.ListOffsetArray):
            # Copy offsets to avoid sharing buffers between arrays
            offsets_array = xp.asarray(content.offsets).copy()
            # Wrap in appropriate Index type
            if isinstance(content.offsets, ak.index.Index32):
                new_offsets = ak.index.Index32(offsets_array)
            elif isinstance(content.offsets, ak.index.IndexU32):
                new_offsets = ak.index.IndexU32(offsets_array)
            else:
                new_offsets = ak.index.Index64(offsets_array)

            return ak.contents.ListOffsetArray(
                new_offsets,
                copy_with_empty_buffers(content.content),
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.ListArray):
            # Copy starts/stops to avoid sharing buffers
            starts_array = xp.asarray(content.starts).copy()
            stops_array = xp.asarray(content.stops).copy()
            # Wrap in appropriate Index types
            if isinstance(content.starts, ak.index.Index32):
                new_starts = ak.index.Index32(starts_array)
                new_stops = ak.index.Index32(stops_array)
            elif isinstance(content.starts, ak.index.IndexU32):
                new_starts = ak.index.IndexU32(starts_array)
                new_stops = ak.index.IndexU32(stops_array)
            else:
                new_starts = ak.index.Index64(starts_array)
                new_stops = ak.index.Index64(stops_array)

            return ak.contents.ListArray(
                new_starts,
                new_stops,
                copy_with_empty_buffers(content.content),
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.RecordArray):
            return ak.contents.RecordArray(
                [copy_with_empty_buffers(c) for c in content.contents],
                content.fields,
                length=content.length,
                parameters=content._parameters,
                backend=backend
            )
        elif isinstance(content, ak.contents.IndexedArray):
            # Copy index to avoid sharing buffers
            index_array = xp.asarray(content.index).copy()
            # Wrap in appropriate Index type
            if isinstance(content.index, ak.index.Index32):
                new_index = ak.index.Index32(index_array)
            elif isinstance(content.index, ak.index.IndexU32):
                new_index = ak.index.IndexU32(index_array)
            else:
                new_index = ak.index.Index64(index_array)

            return ak.contents.IndexedArray(
                new_index,
                copy_with_empty_buffers(content.content),
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.IndexedOptionArray):
            # Copy index to avoid sharing buffers
            index_array = xp.asarray(content.index).copy()
            # Wrap in appropriate Index type
            if isinstance(content.index, ak.index.Index32):
                new_index = ak.index.Index32(index_array)
            elif isinstance(content.index, ak.index.IndexU32):
                new_index = ak.index.IndexU32(index_array)
            else:
                new_index = ak.index.Index64(index_array)

            return ak.contents.IndexedOptionArray(
                new_index,
                copy_with_empty_buffers(content.content),
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.RegularArray):
            return ak.contents.RegularArray(
                copy_with_empty_buffers(content.content),
                content.size,
                content.length,
                parameters=content._parameters
            )
        else:
            # For other types, fallback to copy
            return content

    new_layout = copy_with_empty_buffers(layout)
    return ak.Array(new_layout)


@nvtx.annotate("awkward_to_iterator")
def awkward_to_cccl_iterator(array=None, form=None, buffers=None, dtype=None, return_offsets=True):
    """
    Convert an Awkward Array to a cuda.compute iterator (zero-copy).

    This function recursively traverses the Awkward form structure and constructs
    the corresponding cuda.compute iterator:
    - NumpyArray -> CuPy array
    - RecordArray -> ZipIterator over field iterators
    - IndexedArray -> PermutationIterator with index buffer
    - ListOffsetArray -> Iterator for the flattened content

    The resulting iterator can be used with the cuda.compute library.

    Args:
        array: Awkward Array (if starting fresh)
        form: Awkward form (from ak.to_buffers)
        buffers: Buffer dict (from ak.to_buffers)
        dtype: Optional dtype to cast to (e.g., np.float32 for GPU structs)
        return_offsets: If True, extract and return offsets for list structures (default: True)

    Returns:
        Iterator or CuPy array representing the structure
        Dictionary with metadata: {"form": ..., "buffers": ..., "offsets": ..., "length": ..., "count": ...}
        - offsets: CuPy array of offsets if list structure exists, else None
        - length: Array length (number of lists)
        - count: Total number of items across all lists (avoids .get() calls)
    """
    # Initial call: extract form and buffers from array
    initial_call = form is None and buffers is None
    length = None

    if initial_call:
        if array is None:
            raise ValueError(
                "Must provide either 'array' or both 'form' and 'buffers'")

        # Fast path: use low-level API to avoid dispatch overhead
        # Access layout directly if it's an ak.Array, otherwise convert
        if isinstance(array, ak.Array):
            layout = array.layout
        elif hasattr(array, 'layout'):
            # It's a Record or similar
            layout = array.layout
        elif isinstance(array, ak.contents.Content):
            # Already a layout
            layout = array
        else:
            # Rare fallback: need to convert to layout (will use dispatch, but rare)
            layout = ak.to_layout(array)

        # Check if already on CUDA backend, if not convert using low-level method
        if layout._backend.name != "cuda":
            from awkward._backends.dispatch import regularize_backend
            cuda_backend = regularize_backend("cuda")
            layout = layout.to_backend(cuda_backend)

        # Use low-level to_buffers to avoid @high_level_function dispatch overhead
        form, length, buffers = ak._do.to_buffers(layout)

    # Helper to extract offsets from the form structure
    def extract_offsets_from_form(form_to_search):
        """Navigate through form structure to find and extract list offsets."""
        search_form = form_to_search

        # Unwrap IndexedForm/IndexedOptionForm
        if isinstance(search_form, (ak.forms.IndexedForm, ak.forms.IndexedOptionForm)):
            search_form = search_form.content

        # Unwrap RecordForm (use first field)
        if isinstance(search_form, ak.forms.RecordForm):
            search_form = search_form.contents[0]

        # Extract offsets from ListOffsetForm/ListForm
        if isinstance(search_form, (ak.forms.ListOffsetForm, ak.forms.ListForm)):
            offsets_key = f"{search_form.form_key}-offsets"
            return cp.asarray(buffers[offsets_key], dtype=np.int64)

        return None

    # Extract offsets if this is the initial call and offsets are requested
    offsets = None
    count = None
    if initial_call and return_offsets:
        offsets = extract_offsets_from_form(form)
        # Pre-compute count (total number of items in the list array)
        if offsets is not None:
            count = int(offsets[-1])

    # Helper to create metadata dict for return
    def make_metadata():
        return {
            "form": form,
            "buffers": buffers,
            "offsets": offsets,
            "length": length,
            "count": count
        }

    # Base case: NumpyArray - return the flat data buffer
    if isinstance(form, ak.forms.NumpyForm):
        data_key = f"{form.form_key}-data"
        buffer = buffers[data_key]
        if dtype is not None:
            buffer = cp.asarray(buffer, dtype=dtype)
        else:
            buffer = cp.asarray(buffer)

        if initial_call:
            return buffer, make_metadata()
        else:
            return buffer, (form, buffers)

    # RecordArray: create ZipIterator over all fields
    elif isinstance(form, ak.forms.RecordForm):
        field_iterators = []
        for field_form in form.contents:
            field_iter, _ = awkward_to_cccl_iterator(
                form=field_form, buffers=buffers, dtype=dtype, return_offsets=False
            )
            field_iterators.append(field_iter)

        result = ZipIterator(*field_iterators)
        if initial_call:
            return result, make_metadata()
        else:
            return result, (form, buffers)

    # IndexedArray: create PermutationIterator with index mapping
    elif isinstance(form, (ak.forms.IndexedForm, ak.forms.IndexedOptionForm)):
        index_key = f"{form.form_key}-index"
        index_buffer = cp.asarray(buffers[index_key])

        # Recursively get iterator for the content
        content_iter, _ = awkward_to_cccl_iterator(
            form=form.content, buffers=buffers, dtype=dtype, return_offsets=False
        )

        result = PermutationIterator(content_iter, index_buffer)
        if initial_call:
            return result, make_metadata()
        else:
            return result, (form, buffers)

    # ListOffsetArray: return iterator for the flattened content
    # (Offsets are extracted at the top level if this is an initial call)
    elif isinstance(form, (ak.forms.ListOffsetForm, ak.forms.ListForm)):
        # Recursively handle the content (which is already flattened)
        content_iter, _ = awkward_to_cccl_iterator(
            form=form.content, buffers=buffers, dtype=dtype, return_offsets=False
        )

        if initial_call:
            return content_iter, make_metadata()
        else:
            return content_iter, (form, buffers)

    else:
        raise NotImplementedError(
            f"Form type {type(form).__name__} not yet supported. "
            f"Please add support or file an issue."
        )


@nvtx.annotate("reconstruct_with_offsets")
def reconstruct_with_offsets(list_array, new_offsets):
    """
    Given a list array and new offsets representing for example
    a filtered view, reconstruct the list array with the new offsets.
    """

    if isinstance(list_array, ak.Array):
        layout = list_array.layout
    elif hasattr(list_array, 'layout'):
        layout = list_array.layout
    elif isinstance(list_array, ak.contents.Content):
        layout = list_array
    else:
        layout = ak.to_layout(list_array)

    # Wrap new_offsets in an Index if it's not already
    if not isinstance(new_offsets, ak.index.Index):
        # Determine the appropriate Index type based on dtype
        if hasattr(new_offsets, 'dtype'):
            dtype = new_offsets.dtype
        else:
            dtype = np.int64

        if dtype == np.int32:
            new_offsets = ak.index.Index32(new_offsets)
        elif dtype == np.uint32:
            new_offsets = ak.index.IndexU32(new_offsets)
        else:
            new_offsets = ak.index.Index64(new_offsets)

    # Find the top-level list and reconstruct with new offsets
    def reconstruct_list(content, new_offsets):
        if isinstance(content, ak.contents.ListOffsetArray):
            # Slice content to match new offsets
            num_data = int(new_offsets.data[-1])
            sliced_content = content.content[:num_data]
            return ak.contents.ListOffsetArray(
                new_offsets,
                sliced_content,
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.IndexedArray):
            # Recurse through indexed wrapper
            new_content = reconstruct_list(content.content, new_offsets)
            return ak.contents.IndexedArray(
                content.index,
                new_content,
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.IndexedOptionArray):
            # Recurse through indexed option wrapper
            new_content = reconstruct_list(content.content, new_offsets)
            return ak.contents.IndexedOptionArray(
                content.index,
                new_content,
                parameters=content._parameters
            )
        elif isinstance(content, ak.contents.RecordArray):
            # For records, reconstruct each field
            new_contents = [reconstruct_list(
                c, new_offsets) for c in content.contents]
            # Length should match the number of lists (offsets length - 1)
            new_length = len(new_offsets.data) - 1 if isinstance(new_offsets,
                                                                 ak.index.Index) else len(new_offsets) - 1
            return ak.contents.RecordArray(
                new_contents,
                content.fields,
                length=new_length,
                parameters=content._parameters,
                backend=content._backend
            )
        else:
            # Shouldn't reach here for typical list arrays
            return content

    new_layout = reconstruct_list(layout, new_offsets)
    return ak.Array(new_layout)


@nvtx.annotate("filter_lists")
def filter_lists(array, cond):
    it, meta = awkward_to_cccl_iterator(array)
    in_segments = meta["offsets"]
    out_array = empty_like(array)
    it_out, meta_out = awkward_to_cccl_iterator(out_array)
    out_segments = meta_out["offsets"]
    num_items = meta["count"]
    segmented_select(
        it,
        in_segments,
        it_out,
        out_segments,
        cond,
        num_items
    )
    return reconstruct_with_offsets(out_array, out_segments)


@nvtx.annotate("select_lists")
def select_lists(array, mask):
    data_in, meta = awkward_to_cccl_iterator(array)
    offsets_in = meta["offsets"]
    offsets_out = meta["offsets"]
    num_lists = meta["length"]
    num_elements = meta["count"]
    out_array = empty_like(array)
    data_out, meta = awkward_to_cccl_iterator(out_array)
    d_num_selected_out = cp.empty(2, np.int32)
    select_segments(
        data_in,
        offsets_in,
        mask,
        data_out,
        offsets_out,
        d_num_selected_out,
        num_elements,
        num_lists)
    num_elements_kept, num_lists_kept = d_num_selected_out
    offsets_out = offsets_out[:num_lists_kept+1]
    return reconstruct_with_offsets(out_array, offsets_out)


@nvtx.annotate("list_sizes")
def list_sizes(array):
    _, meta = awkward_to_cccl_iterator(array)
    return segment_sizes(meta["offsets"])


@nvtx.annotate("transform_lists")
def transform_lists(array, out_array, list_size, op):
    data_in, meta = awkward_to_cccl_iterator(array)
    data_out, _ = awkward_to_cccl_iterator(out_array)
    num_segments = meta["length"]
    transform_segments(data_in, data_out, list_size, op, num_segments)
    return out_array