/**
 *  @brief Argument-sorting of string sequences.
 *  @file python/stringzilla/sort.c
 *  @author Ash Vardanian
 */
#include "stringzilla.h"

/** @brief Dispatches to the byte-wise or Unicode case-folded argsort backend. */
sz_status_t Strs_run_argsort_(sz_bool_t uncased, sz_sequence_t const *sequence, sz_sorted_idx_t *order, sz_size_t top,
                              sz_bool_t reverse) {
    return uncased ? sz_sequence_argsort_uncased(sequence, NULL, order, top, reverse)
                   : sz_sequence_argsort(sequence, NULL, order, top, reverse);
}

char const doc_argsort[] =                                                                         //
    "argsort(*, reverse=False, uncased=False, top=None, out=None) -> tuple[int, ...] | buffer\n"   //
    "\n"                                                                                           //
    "Return the stable permutation of indices that sorts the Strs.\n"                              //
    "\n"                                                                                           //
    "Args:\n"                                                                                      //
    "  reverse (bool, optional): Sort in descending order. Defaults to False.\n"                   //
    "  uncased (bool, optional): Order by Unicode case-folding. Defaults to False.\n"              //
    "  top (int, optional): Keep only the `top` leading indices. Defaults to None (all).\n"        //
    "  out (buffer, optional): Writable, C-contiguous buffer of pointer-width unsigned integers "  //
    "(e.g. numpy.uintp) to receive the indices with zero allocation. "                             //
    "Defaults to None.\n"                                                                          //
    "Returns:\n"                                                                                   //
    "  tuple[int, ...]: The sorting permutation, or `out` itself when an `out` buffer is given.\n" //
    "Example:\n"                                                                                   //
    "  >>> sz.Strs(['banana', 'apple', 'cherry']).argsort()\n"                                     //
    "  (1, 0, 2)";

/**
 *  @brief Returns the tuple of indices permuting a `Strs` object into sorted order.
 *         With `top=k`, returns just the `k` leading indices (top-K).
 */
PyObject *Strs_argsort(Strs *self, PyObject *const *args, Py_ssize_t positional_args_count,
                       PyObject *args_names_tuple) {

    sz_bool_t reverse = sz_false_k, uncased = sz_false_k;
    sz_size_t top = 0;
    PyObject *out_obj = NULL;
    if (positional_args_count != 0) {
        PyErr_SetString(PyExc_TypeError, "argsort() takes no positional arguments");
        return NULL;
    }

    PyObject *reverse_obj = NULL, *uncased_obj = NULL, *top_obj = NULL;
    Py_ssize_t const args_names_count = args_names_tuple ? PyTuple_GET_SIZE(args_names_tuple) : 0;
    for (Py_ssize_t i = 0; i < args_names_count; ++i) {
        PyObject *key = PyTuple_GET_ITEM(args_names_tuple, i);
        PyObject *value = args[positional_args_count + i];
        if (PyUnicode_CompareWithASCIIString(key, "reverse") == 0) { reverse_obj = value; }
        else if (PyUnicode_CompareWithASCIIString(key, "uncased") == 0) { uncased_obj = value; }
        else if (PyUnicode_CompareWithASCIIString(key, "top") == 0) { top_obj = value; }
        else if (PyUnicode_CompareWithASCIIString(key, "out") == 0) { out_obj = value; }
        else {
            PyErr_Format(PyExc_TypeError, "argsort() got an unexpected keyword argument '%U'", key);
            return NULL;
        }
    }
    if (reverse_obj) {
        if (!PyBool_Check(reverse_obj)) {
            PyErr_SetString(PyExc_TypeError, "argsort(): reverse must be a bool");
            return NULL;
        }
        reverse = (sz_bool_t)(PyObject_IsTrue(reverse_obj) != 0);
    }
    if (uncased_obj) {
        if (!PyBool_Check(uncased_obj)) {
            PyErr_SetString(PyExc_TypeError, "argsort(): uncased must be a bool");
            return NULL;
        }
        uncased = (sz_bool_t)(PyObject_IsTrue(uncased_obj) != 0);
    }
    if (top_obj && top_obj != Py_None) {
        if (!PyLong_Check(top_obj)) {
            PyErr_SetString(PyExc_TypeError, "argsort(): top must be an int or None");
            return NULL;
        }
        Py_ssize_t top_value = PyLong_AsSsize_t(top_obj);
        if (top_value == -1 && PyErr_Occurred()) return NULL;
        if (top_value < 0) {
            PyErr_SetString(PyExc_ValueError, "argsort(): top must be non-negative");
            return NULL;
        }
        top = (sz_size_t)top_value;
    }

    sz_size_t const count = Strs_len(self);
    // With `top` set, only the leading `top` indices are ordered, so we expose just those.
    sz_size_t const result_count = (top != 0 && top < count) ? top : count;

    // Optionally bind a caller-provided output buffer of `sz_sorted_idx_t` for a zero-allocation path.
    Py_buffer out_view;
    sz_bool_t have_out = sz_false_k;
    if (out_obj && out_obj != Py_None) {
        // PyBUF_CONTIG = writable + C-contiguous; rejects strided/read-only targets up front.
        if (PyObject_GetBuffer(out_obj, &out_view, PyBUF_CONTIG) != 0) return NULL;
        if (out_view.itemsize != (Py_ssize_t)sizeof(sz_sorted_idx_t)) {
            PyErr_Format(PyExc_TypeError,
                         "out must be a contiguous buffer of %zu-byte unsigned integers (e.g. numpy.uintp)",
                         sizeof(sz_sorted_idx_t));
            PyBuffer_Release(&out_view);
            return NULL;
        }
        if ((sz_size_t)(out_view.len / out_view.itemsize) < result_count) {
            PyErr_Format(PyExc_ValueError, "out buffer holds %zd indices, need %zu",
                         (Py_ssize_t)(out_view.len / out_view.itemsize), result_count);
            PyBuffer_Release(&out_view);
            return NULL;
        }
        have_out = sz_true_k;
    }

    // The backend writes a full permutation of `[0, count)` even for top-k, so the `order` buffer must
    // hold `count` indices. For a full sort the result *is* that permutation, so we can sort straight into
    // the caller's buffer; for top-k we sort into a scratch allocation and copy out only the leading `top`,
    // never writing past `result_count` of the caller's buffer.
    sz_bool_t const sort_into_out = have_out && result_count == count;
    sz_sorted_idx_t *order = sort_into_out ? (sz_sorted_idx_t *)out_view.buf
                                           : (sz_sorted_idx_t *)malloc(sizeof(sz_sorted_idx_t) * count);
    if (!order && count) {
        if (have_out) PyBuffer_Release(&out_view);
        PyErr_Format(PyExc_MemoryError, "Unable to allocate memory for the sorting operation");
        return NULL;
    }

    // Call our sorting algorithm (`reverse` and `uncased` are handled natively).
    sz_sequence_t sequence;
    sz_fill(&sequence, sizeof(sequence), 0);
    sequence.count = count;
    sequence.handle = self;
    sequence.get_start = Strs_get_start_;
    sequence.get_length = Strs_get_length_;
    sz_status_t status = Strs_run_argsort_(uncased, &sequence, order, top, reverse);
    sz_unused_(status);

    // Caller buffer path: indices are already in (or now copied into) `out`; hand the buffer back.
    if (have_out) {
        if (!sort_into_out) {
            sz_copy((sz_ptr_t)out_view.buf, (sz_cptr_t)order, result_count * sizeof(sz_sorted_idx_t));
            free(order);
        }
        PyBuffer_Release(&out_view);
        Py_INCREF(out_obj);
        return out_obj;
    }

    // No caller buffer: materialize a tuple of Python ints.
    PyObject *tuple = PyTuple_New(result_count);
    if (!tuple) {
        free(order);
        return NULL;
    }
    for (sz_size_t i = 0; i < result_count; ++i) {
        PyObject *index = PyLong_FromSize_t(order[i]);
        if (!index) {
            free(order);
            Py_DECREF(tuple);
            return NULL;
        }
        PyTuple_SET_ITEM(tuple, i, index);
    }
    free(order);
    return tuple;
}