#pragma once

#include <irrt/debug.hpp>
#include <irrt/exception.hpp>
#include <irrt/int_types.hpp>
#include <irrt/math_util.hpp>

namespace
{
namespace range
{
template <typename T> T len(T start, T stop, T step)
{
    // Reference: https://github.com/python/cpython/blob/9dbd12375561a393eaec4b21ee4ac568a407cdb0/Objects/rangeobject.c#L933
    if (step > 0 && start < stop)
        return 1 + (stop - 1 - start) / step;
    else if (step < 0 && start > stop)
        return 1 + (start - 1 - stop) / (-step);
    else
        return 0;
}
} // namespace range

/**
 * @brief A Python range.
 */
template <typename T> struct Range
{
    T start;
    T stop;
    T step;

    /**
     * @brief Calculate the `len()` of this range.
     */
    template <typename SizeT> T len()
    {
        debug_assert(SizeT, step != 0);
        return range::len(start, stop, step);
    }
};

namespace slice
{
/**
 * @brief Resolve a slice index under a given length like Python indexing.
 *
 * In Python, if you have a `list` of length 100, `list[-1]` resolves to
 * `list[99]`, so `resolve_index_in_length_clamped(100, -1)` returns `99`.
 *
 * If `length` is 0, 0 is returned for any value of `index`.
 *
 * If `index` is out of bounds, clamps the returned value between `0` and
 * `length - 1` (inclusive).
 *
 */
template <typename T> T resolve_index_in_length_clamped(T length, T index)
{
    if (index < 0)
    {
        return max<T>(length + index, 0);
    }
    else
    {
        return min<T>(length, index);
    }
}

/**
 * @brief Like `resolve_index_in_length_clamped`, but returns `-1` if `index` is out of bounds.
 */
template <typename T> T resolve_index_in_length(T length, T index)
{
    T resolved = index < 0 ? length + index : index;
    if (0 <= resolved && resolved < length)
    {
        return resolved;
    }
    else
    {
        return -1;
    }
}

/**
 * @brief Resolve a slice as a range.
 *
 * In Python, this would be `range(*slice(start, stop, step).indices(length))`.
 */
template <typename T>
void indices(bool start_defined, T start, bool stop_defined, T stop, bool step_defined, T step, T length,
             T *range_start, T *range_stop, T *range_step)
{
    // Reference: https://github.com/python/cpython/blob/main/Objects/sliceobject.c#L388
    *range_step = step_defined ? step : 1;
    bool step_is_negative = *range_step < 0;

    T lower, upper;
    if (step_is_negative)
    {
        lower = -1;
        upper = length - 1;
    }
    else
    {
        lower = 0;
        upper = length;
    }

    if (start_defined)
    {
        *range_start = start < 0 ? max(lower, start + length) : min(upper, start);
    }
    else
    {
        *range_start = step_is_negative ? upper : lower;
    }

    if (stop_defined)
    {
        *range_stop = stop < 0 ? max(lower, stop + length) : min(upper, stop);
    }
    else
    {
        *range_stop = step_is_negative ? lower : upper;
    }
}
} // namespace slice

/**
 * @brief A Python-like slice with **unresolved** indices.
 */
template <typename T> struct Slice
{
    bool start_defined;
    T start;

    bool stop_defined;
    T stop;

    bool step_defined;
    T step;

    Slice()
    {
        this->reset();
    }

    void reset()
    {
        this->start_defined = false;
        this->stop_defined = false;
        this->step_defined = false;
    }

    void set_start(T start)
    {
        this->start_defined = true;
        this->start = start;
    }

    void set_stop(T stop)
    {
        this->stop_defined = true;
        this->stop = stop;
    }

    void set_step(T step)
    {
        this->step_defined = true;
        this->step = step;
    }

    /**
     * @brief Resolve this slice as a range.
     *
     * In Python, this would be `range(*slice(start, stop, step).indices(length))`.
     */
    template <typename SizeT> Range<T> indices(T length)
    {
        // Reference: https://github.com/python/cpython/blob/main/Objects/sliceobject.c#L388
        debug_assert(SizeT, length >= 0);

        Range<T> result;
        slice::indices(start_defined, start, stop_defined, stop, step_defined, step, length, &result.start,
                       &result.stop, &result.step);
        return result;
    }

    /**
     * @brief Like `.indices()` but with assertions.
     */
    template <typename SizeT> Range<T> indices_checked(T length)
    {
        // TODO: Switch to `SizeT length`

        if (length < 0)
        {
            raise_exception(SizeT, EXN_VALUE_ERROR, "length should not be negative, got {0}", length, NO_PARAM,
                            NO_PARAM);
        }

        if (this->step_defined && this->step == 0)
        {
            raise_exception(SizeT, EXN_VALUE_ERROR, "slice step cannot be zero", NO_PARAM, NO_PARAM, NO_PARAM);
        }

        return this->indices<SizeT>(length);
    }
};
} // namespace

extern "C"
{
    int32_t __nac3_range_len_i32(int32_t start, int32_t stop, int32_t step)
    {
        range::len<int32_t>(start, stop, step);
    }

    int32_t __nac3_range_len_i3264(int32_t start, int32_t stop, int32_t step)
    {
        range::len<int64_t>(start, stop, step);
    }
}