from __future__ import annotations

from functools import lru_cache
from typing import TYPE_CHECKING, Any, cast

import pyarrow as pa
import pyarrow.compute as pc

from narwhals._compliant import EagerSeriesNamespace
from narwhals._utils import Implementation, Version, isinstance_or_issubclass
from narwhals.exceptions import ColumnNotFoundError

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator, Mapping, Sequence
    from typing import Literal, TypeAlias

    from typing_extensions import TypeIs

    from narwhals._arrow.series import ArrowSeries
    from narwhals._arrow.typing import (  # type: ignore[attr-defined]
        ArrayAny,
        ArrayOrScalar,
        ArrayOrScalarT1,
        ArrayOrScalarT2,
        ChunkedArrayAny,
        DictionaryArrayAny,
        Incomplete,
        NativeIntervalUnit,
        NullPlacement,
        Order,
        PromoteOptions,
        ScalarAny,
    )
    from narwhals._duration import IntervalUnit
    from narwhals.dtypes import DType
    from narwhals.typing import IntoDType, PythonLiteral

    # NOTE: stubs don't allow for `ChunkedArray[StructArray]`
    # Intended to represent the `.chunks` property storing `list[pa.StructArray]`
    ChunkedArrayStructArray: TypeAlias = ChunkedArrayAny

    def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ...
    def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ...
    def is_list(t: Any) -> TypeIs[pa.ListType[Any]]: ...
    def is_large_list(t: Any) -> TypeIs[pa.LargeListType[Any]]: ...
    def is_fixed_size_list(t: Any) -> TypeIs[pa.FixedSizeListType[Any, Any]]: ...
    def is_dictionary(t: Any) -> TypeIs[pa.DictionaryType[Any, Any, Any]]: ...
    def extract_regex(
        strings: ChunkedArrayAny,
        /,
        pattern: str,
        *,
        options: Any = None,
        memory_pool: Any = None,
    ) -> ChunkedArrayStructArray: ...

else:
    from pyarrow.compute import extract_regex
    from pyarrow.types import (
        is_dictionary,
        is_duration,
        is_fixed_size_list,
        is_large_list,
        is_list,
        is_timestamp,
    )

BACKEND_VERSION = Implementation.PYARROW._backend_version()
"""Static backend version for `pyarrow`."""

UNITS_DICT: Mapping[IntervalUnit, NativeIntervalUnit] = {
    "y": "year",
    "q": "quarter",
    "mo": "month",
    "d": "day",
    "h": "hour",
    "m": "minute",
    "s": "second",
    "ms": "millisecond",
    "us": "microsecond",
    "ns": "nanosecond",
}

lit = pa.scalar
"""Alias for `pyarrow.scalar`."""


def extract_py_scalar(value: Any, /) -> Any:
    from narwhals._arrow.series import maybe_extract_py_scalar

    return maybe_extract_py_scalar(value, return_py_scalar=True)


def is_array_or_scalar(obj: Any) -> TypeIs[ArrayOrScalar]:
    """Return True for any base `pyarrow` container."""
    return isinstance(obj, (pa.ChunkedArray, pa.Array, pa.Scalar))


def arange(start: int, end: int, step: int) -> ArrayAny:
    if BACKEND_VERSION < (21,):
        import numpy as np  # ignore-banned-import

        return pa.array(np.arange(start, end, step))
    # NOTE: Added in https://github.com/apache/arrow/pull/46778
    return pa.arange(start, end, step)  # type: ignore[attr-defined]


def build_list_array(arrays: list[pa.Array[Any]]) -> pa.ChunkedArray[Any]:
    # This works by using concat_arrays to vertically stack the arrays.
    # Then we use take to grab the data by index and to horizontally
    # pack the arrays.
    n = len(arrays[0])
    num_cols = len(arrays)
    i = arange(0, n, 1)
    j = arange(0, num_cols, 1)
    k = arange(0, n * num_cols, 1)
    row_idx = cast("pa.Int64Array", pc.divide(k, lit(num_cols)))
    col_idx = cast("pa.Int64Array", pc.subtract(k, pc.multiply(row_idx, lit(num_cols))))
    i_rep = pc.take(i, row_idx)
    j_tile = pc.take(j, col_idx)
    indices = cast("pa.Int64Array", pc.add(pc.multiply(j_tile, lit(n)), i_rep))
    interleaved = pc.take(pa.concat_arrays(arrays), indices)
    offsets = cast("pa.Int64Array", arange(0, n * num_cols + 1, num_cols))
    # Typing suggests only allow Int32Array is valid for offsets, but Int64Array is fine.
    return pa.chunked_array([pa.ListArray.from_arrays(offsets, interleaved)])  # type: ignore[call-overload]


def chunked_array(
    arr: ArrayOrScalar | list[Iterable[Any]], dtype: pa.DataType | None = None, /
) -> ChunkedArrayAny:
    if isinstance(arr, pa.ChunkedArray):
        return arr
    if isinstance(arr, list):
        return pa.chunked_array(arr, dtype)
    return pa.chunked_array([arr], dtype)


def nulls_like(n: int, series: ArrowSeries) -> ArrayAny:
    """Create a strongly-typed Array instance with all elements null.

    Uses the type of `series`, without upseting `mypy`.
    """
    return pa.nulls(n, series.native.type)


def repeat(
    value: PythonLiteral | ScalarAny, n: int, /, dtype: pa.DataType | None = None
) -> ArrayAny:
    """Create an Array instance whose slots are the given scalar.

    *Optionally*, casting to `dtype` **before** repeating `n` times.
    """
    lit_: Incomplete = lit
    return pa.repeat(lit_(value, type=dtype), n)


def zeros(n: int, /) -> pa.Int64Array:
    return pa.repeat(0, n)


def native_to_narwhals_dtype(dtype: pa.DataType, version: Version) -> DType:
    if isinstance(dtype, pa.ExtensionType):
        return version.dtypes.Unknown()
    return native_non_extension_to_narwhals_dtype(dtype, version)


@lru_cache(maxsize=16)
def native_non_extension_to_narwhals_dtype(dtype: pa.DataType, version: Version) -> DType:  # noqa: C901, PLR0912
    dtypes = version.dtypes
    if pa.types.is_int64(dtype):
        return dtypes.Int64()
    if pa.types.is_int32(dtype):
        return dtypes.Int32()
    if pa.types.is_int16(dtype):
        return dtypes.Int16()
    if pa.types.is_int8(dtype):
        return dtypes.Int8()
    if pa.types.is_uint64(dtype):
        return dtypes.UInt64()
    if pa.types.is_uint32(dtype):
        return dtypes.UInt32()
    if pa.types.is_uint16(dtype):
        return dtypes.UInt16()
    if pa.types.is_uint8(dtype):
        return dtypes.UInt8()
    if pa.types.is_boolean(dtype):
        return dtypes.Boolean()
    if pa.types.is_float64(dtype):
        return dtypes.Float64()
    if pa.types.is_float32(dtype):
        return dtypes.Float32()
    if pa.types.is_float16(dtype):
        return dtypes.Float16()
    # bug in coverage? it shows `31->exit` (where `31` is currently the line number of
    # the next line), even though both when the if condition is true and false are covered
    if (  # pragma: no cover
        pa.types.is_string(dtype)
        or pa.types.is_large_string(dtype)
        or getattr(pa.types, "is_string_view", lambda _: False)(dtype)
    ):
        return dtypes.String()
    if pa.types.is_date32(dtype):
        return dtypes.Date()
    if is_timestamp(dtype):
        return dtypes.Datetime(time_unit=dtype.unit, time_zone=dtype.tz)
    if is_duration(dtype):
        return dtypes.Duration(time_unit=dtype.unit)
    if pa.types.is_dictionary(dtype):
        return dtypes.Categorical()
    if pa.types.is_struct(dtype):
        return dtypes.Struct(
            [
                dtypes.Field(
                    dtype.field(i).name,
                    native_to_narwhals_dtype(dtype.field(i).type, version),
                )
                for i in range(dtype.num_fields)
            ]
        )
    if is_list(dtype) or is_large_list(dtype):
        return dtypes.List(native_to_narwhals_dtype(dtype.value_type, version))
    if is_fixed_size_list(dtype):
        return dtypes.Array(
            native_to_narwhals_dtype(dtype.value_type, version), dtype.list_size
        )
    if pa.types.is_decimal128(dtype):
        return dtypes.Decimal(precision=dtype.precision, scale=dtype.scale)
    if pa.types.is_time32(dtype) or pa.types.is_time64(dtype):
        return dtypes.Time()
    if pa.types.is_binary(dtype):
        return dtypes.Binary()
    return dtypes.Unknown()  # pragma: no cover


dtypes = Version.MAIN.dtypes
NW_TO_PA_DTYPES: Mapping[type[DType], pa.DataType] = {
    dtypes.Float64: pa.float64(),
    dtypes.Float32: pa.float32(),
    dtypes.Float16: pa.float16(),
    dtypes.Binary: pa.binary(),
    dtypes.String: pa.string(),
    dtypes.Boolean: pa.bool_(),
    dtypes.Categorical: pa.dictionary(pa.uint32(), pa.string()),
    dtypes.Date: pa.date32(),
    dtypes.Time: pa.time64("ns"),
    dtypes.Int8: pa.int8(),
    dtypes.Int16: pa.int16(),
    dtypes.Int32: pa.int32(),
    dtypes.Int64: pa.int64(),
    dtypes.UInt8: pa.uint8(),
    dtypes.UInt16: pa.uint16(),
    dtypes.UInt32: pa.uint32(),
    dtypes.UInt64: pa.uint64(),
}
UNSUPPORTED_DTYPES = (dtypes.Object, dtypes.Enum)


def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> pa.DataType:
    dtypes = version.dtypes
    base_type = dtype.base_type()
    if pa_type := NW_TO_PA_DTYPES.get(base_type):
        return pa_type
    if isinstance_or_issubclass(dtype, dtypes.Datetime):
        unit = dtype.time_unit
        return pa.timestamp(unit, tz) if (tz := dtype.time_zone) else pa.timestamp(unit)
    if isinstance_or_issubclass(dtype, dtypes.Duration):
        return pa.duration(dtype.time_unit)
    if isinstance_or_issubclass(dtype, dtypes.List):
        return pa.list_(value_type=narwhals_to_native_dtype(dtype.inner, version=version))
    if isinstance_or_issubclass(dtype, dtypes.Struct):
        return pa.struct(
            [
                (field.name, narwhals_to_native_dtype(field.dtype, version=version))
                for field in dtype.fields
            ]
        )
    if isinstance_or_issubclass(dtype, dtypes.Array):
        inner = narwhals_to_native_dtype(dtype.inner, version=version)
        list_size = dtype.size
        return pa.list_(inner, list_size=list_size)
    if isinstance_or_issubclass(dtype, dtypes.Decimal):
        return pa.decimal128(precision=dtype.precision, scale=dtype.scale)
    if issubclass(base_type, UNSUPPORTED_DTYPES):
        msg = f"Converting to {base_type.__name__} dtype is not supported for PyArrow."
        raise NotImplementedError(msg)
    msg = f"Unknown dtype: {dtype}"  # pragma: no cover
    raise AssertionError(msg)


def extract_native(
    lhs: ArrowSeries, rhs: ArrowSeries | PythonLiteral | ScalarAny
) -> tuple[ChunkedArrayAny, ChunkedArrayAny | ScalarAny]:
    """Extract native objects in binary  operation.

    If the comparison isn't supported, return `NotImplemented` so that the
    "right-hand-side" operation (e.g. `__radd__`) can be tried.

    If one of the two sides has a `_broadcast` flag, then extract the scalar
    underneath it so that PyArrow can do its own broadcasting.
    """
    from narwhals._arrow.series import ArrowSeries

    if rhs is None:  # pragma: no cover
        return lhs.native, lit(None, type=lhs._type)

    if isinstance(rhs, ArrowSeries):
        if lhs._broadcast and not rhs._broadcast:
            return lhs.native[0], rhs.native
        if rhs._broadcast:
            return lhs.native, rhs.native[0]
        return lhs.native, rhs.native

    if isinstance(rhs, list):
        msg = "Expected Series or scalar, got list."
        raise TypeError(msg)

    return lhs.native, rhs if isinstance(rhs, pa.Scalar) else lit(rhs)


def floordiv_compat(left: ArrayOrScalar, right: ArrayOrScalar, /) -> Any:
    # The following lines are adapted from pandas' pyarrow implementation.
    # Ref: https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L124-L154
    # We modify it to add `safe_mask` so as to align with Polars' behaviour when dividing by zero.
    safe_mask = pc.not_equal(right, lit(0, type=right.type))
    safe_right = pc.if_else(safe_mask, right, lit(1, type=right.type))  # Dummy value
    non_safe_result = lit(None, type=left.type)

    if pa.types.is_integer(left.type) and pa.types.is_integer(right.type):
        divided = pc.if_else(
            safe_mask, pc.divide_checked(left, safe_right), non_safe_result
        )
        # TODO @dangotbanned: Use a `TypeVar` in guards
        # Narrowing to a `Union` isn't interacting well with the rest of the stubs
        # https://github.com/zen-xu/pyarrow-stubs/pull/215
        if pa.types.is_signed_integer(divided.type):
            div_type = cast("pa._lib.Int64Type", divided.type)
            has_remainder = pc.not_equal(pc.multiply(divided, right), left)
            has_one_negative_operand = pc.less(
                pc.bit_wise_xor(left, right), lit(0, div_type)
            )
            result = pc.if_else(
                pc.and_(has_remainder, has_one_negative_operand),
                pc.subtract(divided, lit(1, div_type)),
                divided,
            )
        else:
            result = divided  # pragma: no cover
        result = result.cast(left.type)
    else:
        divided = pc.if_else(safe_mask, pc.divide(left, safe_right), non_safe_result)
        result = pc.floor(divided)
    return result


def cast_for_truediv(
    arrow_array: ArrayOrScalarT1, pa_object: ArrayOrScalarT2
) -> tuple[ArrayOrScalarT1, ArrayOrScalarT2]:
    # Lifted from:
    # https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L108-L122
    # Ensure int / int -> float mirroring Python/Numpy behavior
    # as pc.divide_checked(int, int) -> int
    if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(pa_object.type):
        # GH: 56645.  # noqa: ERA001
        # https://github.com/apache/arrow/issues/35563
        return arrow_array.cast(pa.float64(), safe=False), pa_object.cast(
            pa.float64(), safe=False
        )

    return arrow_array, pa_object


# Regex for date, time, separator and timezone components
DATE_RE = r"(?P<date>\d{1,4}[-/.]\d{1,2}[-/.]\d{1,4}|\d{8})"
SEP_RE = r"(?P<sep>\s|T)"
TIME_RE = r"(?P<time>\d{2}:\d{2}(?::\d{2})?|\d{6}?)"  # \s*(?P<period>[AP]M)?)?
HMS_RE = r"^(?P<hms>\d{2}:\d{2}:\d{2})$"
HM_RE = r"^(?P<hm>\d{2}:\d{2})$"
HMS_RE_NO_SEP = r"^(?P<hms_no_sep>\d{6})$"
TZ_RE = r"(?P<tz>Z|[+-]\d{2}:?\d{2})"  # Matches 'Z', '+02:00', '+0200', '+02', etc.
FULL_RE = rf"{DATE_RE}{SEP_RE}?{TIME_RE}?{TZ_RE}?$"

# Separate regexes for different date formats
YMD_RE = r"^(?P<year>(?:[12][0-9])?[0-9]{2})(?P<sep1>[-/.])(?P<month>0[1-9]|1[0-2])(?P<sep2>[-/.])(?P<day>0[1-9]|[12][0-9]|3[01])$"
DMY_RE = r"^(?P<day>0[1-9]|[12][0-9]|3[01])(?P<sep1>[-/.])(?P<month>0[1-9]|1[0-2])(?P<sep2>[-/.])(?P<year>(?:[12][0-9])?[0-9]{2})$"
MDY_RE = r"^(?P<month>0[1-9]|1[0-2])(?P<sep1>[-/.])(?P<day>0[1-9]|[12][0-9]|3[01])(?P<sep2>[-/.])(?P<year>(?:[12][0-9])?[0-9]{2})$"
YMD_RE_NO_SEP = r"^(?P<year>(?:[12][0-9])?[0-9]{2})(?P<month>0[1-9]|1[0-2])(?P<day>0[1-9]|[12][0-9]|3[01])$"

DATE_FORMATS = (
    (YMD_RE_NO_SEP, "%Y%m%d"),
    (YMD_RE, "%Y-%m-%d"),
    (DMY_RE, "%d-%m-%Y"),
    (MDY_RE, "%m-%d-%Y"),
)
TIME_FORMATS = ((HMS_RE, "%H:%M:%S"), (HM_RE, "%H:%M"), (HMS_RE_NO_SEP, "%H%M%S"))


def _extract_regex_concat_arrays(
    strings: ChunkedArrayAny,
    /,
    pattern: str,
    *,
    options: Any = None,
    memory_pool: Any = None,
) -> pa.StructArray:
    r = pa.concat_arrays(
        extract_regex(strings, pattern, options=options, memory_pool=memory_pool).chunks
    )
    return cast("pa.StructArray", r)


def parse_datetime_format(arr: ChunkedArrayAny) -> str:
    """Try to infer datetime format from StringArray."""
    matches = _extract_regex_concat_arrays(arr.drop_null().slice(0, 10), pattern=FULL_RE)
    if not pc.all(matches.is_valid()).as_py():
        msg = (
            "Unable to infer datetime format, provided format is not supported. "
            "Please report a bug to https://github.com/narwhals-dev/narwhals/issues"
        )
        raise NotImplementedError(msg)

    separators = matches.field("sep")
    tz = matches.field("tz")

    # separators and time zones must be unique
    if pc.count(pc.unique(separators)).as_py() > 1:
        msg = "Found multiple separator values while inferring datetime format."
        raise ValueError(msg)

    if pc.count(pc.unique(tz)).as_py() > 1:
        msg = "Found multiple timezone values while inferring datetime format."
        raise ValueError(msg)

    date_value = _parse_date_format(cast("pc.StringArray", matches.field("date")))
    time_value = parse_time_format(cast("pc.StringArray", matches.field("time")))

    sep_value = separators[0].as_py()
    tz_value = "%z" if tz[0].as_py() else ""

    return f"{date_value}{sep_value}{time_value}{tz_value}"


def _parse_date_format(arr: pc.StringArray) -> str:
    for date_rgx, date_fmt in DATE_FORMATS:
        matches = pc.extract_regex(arr, pattern=date_rgx)
        if date_fmt == "%Y%m%d" and pc.all(matches.is_valid()).as_py():
            return date_fmt
        if (
            pc.all(matches.is_valid()).as_py()
            and pc.count(pc.unique(sep1 := matches.field("sep1"))).as_py() == 1
            and pc.count(pc.unique(sep2 := matches.field("sep2"))).as_py() == 1
            and (date_sep_value := sep1[0].as_py()) == sep2[0].as_py()
        ):
            return date_fmt.replace("-", date_sep_value)

    msg = (
        "Unable to infer datetime format. "
        "Please report a bug to https://github.com/narwhals-dev/narwhals/issues"
    )
    raise ValueError(msg)


def parse_time_format(arr: pc.StringArray) -> str:
    for time_rgx, time_fmt in TIME_FORMATS:
        matches = pc.extract_regex(arr, pattern=time_rgx)
        if pc.all(matches.is_valid()).as_py():
            return time_fmt
    return ""


def pad_series(
    series: ArrowSeries, *, window_size: int, center: bool
) -> tuple[ArrowSeries, int]:
    """Pad series with None values on the left and/or right side, depending on the specified parameters.

    Arguments:
        series: The input ArrowSeries to be padded.
        window_size: The desired size of the window.
        center: Specifies whether to center the padding or not.

    Returns:
        A tuple containing the padded ArrowSeries and the offset value.
    """
    if not center:
        return series, 0
    offset_left = window_size // 2
    # subtract one if window_size is even
    offset_right = offset_left - (window_size % 2 == 0)
    chunks = series.native.chunks
    arrays = nulls_like(offset_left, series), *chunks, nulls_like(offset_right, series)
    return series._with_native(pa.concat_arrays(arrays)), offset_left + offset_right


def cast_to_comparable_string_types(
    *chunked_arrays: ChunkedArrayAny | ScalarAny, separator: str
) -> tuple[Iterator[ChunkedArrayAny | ScalarAny], ScalarAny]:
    # Ensure `chunked_arrays` are either all `string` or all `large_string`.
    dtype = (
        pa.string()  # (PyArrow default)
        if not any(pa.types.is_large_string(ca.type) for ca in chunked_arrays)
        else pa.large_string()
    )
    return (ca.cast(dtype) for ca in chunked_arrays), lit(separator, dtype)


if BACKEND_VERSION >= (14,):
    # https://arrow.apache.org/docs/14.0/python/generated/pyarrow.concat_tables.html
    _PROMOTE: Mapping[PromoteOptions, Mapping[str, Any]] = {
        "default": {"promote_options": "default"},
        "permissive": {"promote_options": "permissive"},
        "none": {"promote_options": "none"},
    }
else:  # pragma: no cover
    # https://arrow.apache.org/docs/13.0/python/generated/pyarrow.concat_tables.html
    _PROMOTE = {
        "default": {"promote": True},
        "permissive": {"promote": True},
        "none": {"promote": False},
    }


def concat_tables(
    tables: Iterable[pa.Table], promote_options: PromoteOptions = "none"
) -> pa.Table:
    return pa.concat_tables(tables, **_PROMOTE[promote_options])


class ArrowSeriesNamespace(EagerSeriesNamespace["ArrowSeries", "ChunkedArrayAny"]): ...


def list_agg(
    array: ChunkedArrayAny,
    func: Literal["min", "max", "mean", "approximate_median", "sum"],
) -> ChunkedArrayAny:
    lit_: Incomplete = lit
    aggregation = (
        ("values", func, pc.ScalarAggregateOptions(min_count=0))
        if func == "sum"
        else ("values", func)
    )
    agg = pa.array(
        pa.Table.from_arrays(
            [pc.list_flatten(array), pc.list_parent_indices(array)],
            names=["values", "offsets"],
        )
        .group_by("offsets")
        .aggregate([aggregation])
        .sort_by("offsets")
        .column(f"values_{func}")
    )
    non_empty_mask = pa.array(pc.not_equal(pc.list_value_length(array), lit(0)))
    if func == "sum":
        # Make sure sum of empty list is 0.
        base_array = pc.if_else(non_empty_mask.is_null(), None, 0)
    else:
        base_array = pa.repeat(lit_(None, type=agg.type), len(array))
    return pa.chunked_array(
        [
            pc.replace_with_mask(
                base_array,
                non_empty_mask.fill_null(False),  # type: ignore[arg-type]
                agg,
            )
        ]
    )


def sortable(array: ChunkedArrayAny, /) -> ChunkedArrayAny:
    """Return an array which orders like `array`, but with a type `pyarrow` can sort.

    `pyarrow` has no ordering kernels for `dictionary` types
    (https://github.com/apache/arrow/issues/29887), so replace the dictionary indices
    with the rank of the value they point at. Ranking the (typically small) dictionary
    is much cheaper than decoding every row, and gives the lexicographic order which
    `polars` and `pandas` use for (unordered) categoricals.
    """
    if not is_dictionary(array.type):
        return array
    # NOTE: stubs type the chunks as `Array`, which is missing `DictionaryArray` members
    chunks = cast("list[DictionaryArrayAny]", array.unify_dictionaries().chunks)
    if not chunks:
        return chunked_array([[]], pa.uint64())
    dictionary = chunks[0].dictionary
    ranks = pc.rank(dictionary, sort_keys="ascending", tiebreaker="min")
    if dictionary.null_count:
        # NOTE: a null *value* must stay null, `rank` would give it a (largest) rank.
        ranks = pc.if_else(pc.is_null(dictionary), lit(None, ranks.type), ranks)
    return pa.chunked_array([ranks.take(chunk.indices) for chunk in chunks])


def sortable_table(table: pa.Table, keys: Iterable[str], /) -> pa.Table:
    """Replace every `dictionary`-typed column named in `keys` with `sortable` ranks."""
    for name in keys:
        index = table.schema.get_field_index(name)
        if index == -1:  # pragma: no cover
            msg = f"Column '{name}' not found in table."
            raise ColumnNotFoundError(msg)
        if is_dictionary(table.field(index).type):
            table = table.set_column(index, name, sortable(table.column(index)))
    return table


def sort_indices(
    table: pa.Table, sorting: Sequence[tuple[str, Order]], null_placement: NullPlacement
) -> pa.UInt64Array:
    """`pc.sort_indices`, hiding how `null_placement` must be passed per version."""
    if BACKEND_VERSION < (25,):  # pragma: no cover
        return pc.sort_indices(table, sort_keys=sorting, null_placement=null_placement)
    # `null_placement` in `SortOptions` is deprecated since 25.0.0;
    # it must be specified per sort key instead.
    keyed = [(name, order, null_placement) for name, order in sorting]
    return pc.sort_indices(table, sort_keys=keyed)  # type: ignore[arg-type]


def list_sort(
    array: ChunkedArrayAny, *, descending: bool, nulls_last: bool
) -> ChunkedArrayAny:
    sort_direction: Literal["ascending", "descending"] = (
        "descending" if descending else "ascending"
    )
    nulls_position: Literal["at_start", "at_end"] = "at_end" if nulls_last else "at_start"
    idx, v = "idx", "values"
    is_not_sorted = pc.greater(pc.list_value_length(array), lit(0))
    indexed = pa.Table.from_arrays(
        [arange(start=0, end=len(array), step=1), array], names=[idx, v]
    )
    not_sorted_part = indexed.filter(is_not_sorted)
    pass_through = indexed.filter(pc.fill_null(pc.invert(is_not_sorted), lit(True)))  # pyright: ignore[reportArgumentType]
    exploded = pa.Table.from_arrays(
        [pc.list_flatten(array), pc.list_parent_indices(array)], names=[v, idx]
    )
    sorted_indices = sort_indices(
        exploded, [(idx, "ascending"), (v, sort_direction)], nulls_position
    )

    offsets = not_sorted_part.column(v).combine_chunks().offsets  # type: ignore[attr-defined]
    sorted_imploded = pa.ListArray.from_arrays(
        offsets, pa.array(exploded.take(sorted_indices).column(v))
    )
    imploded_by_idx = pa.Table.from_arrays(
        [not_sorted_part.column(idx), sorted_imploded], names=[idx, v]
    )
    return pa.concat_tables([imploded_by_idx, pass_through]).sort_by(idx).column(v)