#pragma once

#if !defined(MIJIN_CONTAINER_STRIDE_SPAN_HPP_INCLUDED)
#define MIJIN_CONTAINER_STRIDE_SPAN_HPP_INCLUDED 1

#include <bit>
#include <cstddef>
#include <iterator>
#include <type_traits>
#include "../debug/assert.hpp"

namespace mijin
{

//
// public defines
//

//
// public constants
//

//
// public types
//

template<typename T, typename U>
[[nodiscard]] std::ptrdiff_t bytediff(T* first, U* second) noexcept
{
    return std::bit_cast<std::byte*>(second) - std::bit_cast<std::byte*>(first);
}
template<typename T>
[[nodiscard]] T* byteoffset(T* old, std::ptrdiff_t offset) noexcept
{
    return std::bit_cast<T*>(std::bit_cast<std::byte*>(old) + offset);
}

template<typename T>
class StrideSpanIterator
{
public:
    using difference_type = std::ptrdiff_t;
    using value_type = T;
private:
    T* ele_ = nullptr;
    std::size_t strideBytes_ = 0;
#if MIJIN_CHECKED_ITERATORS
    T* start_ = nullptr;
    T* end_ = nullptr;
#endif
public:
    StrideSpanIterator() = default;
    constexpr StrideSpanIterator(T* ele, std::size_t strideBytes, [[maybe_unused]] T* start, [[maybe_unused]] T* end)
        : ele_(ele), strideBytes_(strideBytes)
#if MIJIN_CHECKED_ITERATORS
        , start_(start), end_(end)
#endif
    {
        MIJIN_ASSERT(strideBytes_!= 0 && bytediff(start_, end_) % strideBytes_ == 0, "StrideSpan: stride elements don't match supplied pointers.");
    }
    StrideSpanIterator(const StrideSpanIterator&) = default;

    StrideSpanIterator& operator=(const StrideSpanIterator&) = default;
    StrideSpanIterator& operator+=(difference_type diff);
    StrideSpanIterator& operator-=(difference_type diff) {
        return *this += -diff;
    }

    constexpr auto operator<=>(const StrideSpanIterator&) const noexcept = default;

    [[nodiscard]] T& operator*() const;
    [[nodiscard]] T* operator->() const;
    StrideSpanIterator& operator++();
    StrideSpanIterator operator++(int);
    StrideSpanIterator& operator--();
    StrideSpanIterator operator--(int);

    [[nodiscard]] difference_type operator-(const StrideSpanIterator& other) const { return bytediff(other.ele_, ele_) / strideBytes_; }
    [[nodiscard]] StrideSpanIterator operator+(difference_type diff) const;
    [[nodiscard]] StrideSpanIterator operator-(difference_type diff) const { return (*this + -diff); }

    [[nodiscard]] T& operator[](difference_type diff) const { return *(*this + diff); }
};
template<typename T>
inline StrideSpanIterator<T> operator+(std::iter_difference_t<T> diff, const StrideSpanIterator<T>& iter) {
    return iter + diff;
}
static_assert(std::random_access_iterator<StrideSpanIterator<int>>);

template<typename T>
class StrideSpan
{
public:
    using element_type = T;
    using value_type = std::remove_cv_t<T>;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using pointer = T*;
    using const_pointer = const T*;
    using reference = T&;
    using const_reference = const T&;
    using iterator = StrideSpanIterator<T>;
    using const_iterator = StrideSpanIterator<const T>;
private:
    T* begin_ = nullptr;
    T* end_ = nullptr;
    std::size_t strideBytes_ = 0;
public:
    StrideSpan() = default;
    StrideSpan(const StrideSpan&) = default;
    StrideSpan(T* begin, T* end, std::size_t strideBytes) noexcept : begin_(begin), end_(end), strideBytes_(strideBytes) {
        MIJIN_ASSERT(strideBytes_ > 0, "Stride elements must be >0.");
        MIJIN_ASSERT(strideBytes_ % alignof(T) == 0, "Stride must be a multiple of element align.");
        MIJIN_ASSERT(bytediff(begin_, end_) % strideBytes_ == 0, "Difference of begin and end must be a multiple of stride bytes.");
    }

    StrideSpan& operator=(const StrideSpan&) = default;
    auto operator<=>(const StrideSpan&) const noexcept = default;

    constexpr operator bool() const noexcept { return !empty(); }
    constexpr bool operator!() const noexcept { return empty(); }
    explicit operator std::span<T>() const noexcept;

    reference operator[](size_type index);
    const_reference operator[](size_type index) const;

    [[nodiscard]] constexpr size_type size() const noexcept { return strideBytes_ == 0 ? 0 : bytediff(begin_, end_) / strideBytes_; }
    [[nodiscard]] constexpr bool empty() const noexcept { return begin_ == end_; }

    [[nodiscard]] iterator begin() { return StrideSpanIterator<T>(begin_, strideBytes_, begin_, end_); }
    [[nodiscard]] const_iterator begin() const { return StrideSpanIterator<const T>(begin_, strideBytes_, begin_, end_); }
    [[nodiscard]] const_iterator cbegin() const { return StrideSpanIterator<const T>(begin_, strideBytes_, begin_, end_); }
    [[nodiscard]] iterator end() { return StrideSpanIterator<T>(end_, strideBytes_, begin_, end_); }
    [[nodiscard]] const_iterator end() const { return StrideSpanIterator<const T>(end_, strideBytes_, begin_, end_); }
    [[nodiscard]] const_iterator cend() const { return StrideSpanIterator<const T>(end_, strideBytes_, begin_, end_); }
};

//
// public functions
//

template<typename T>
StrideSpan<T>::operator std::span<T>() const noexcept
{
    MIJIN_ASSERT(strideBytes_ == sizeof(T), "Can't convert to regular span if stride isn't exactly size of type.");
    return {&*begin_, &*end_};
}

template<typename T>
StrideSpanIterator<T>& StrideSpanIterator<T>::operator+=(difference_type diff)
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT(byteoffset(ele_, diff * strideBytes_) <= end_ && byteoffset(ele_, diff * strideBytes_) >= start_,
        "Moving iterator out-of-bounds.");
#endif
    ele_ = byteoffset(ele_, diff * strideBytes_);
    return *this;
}

template<typename T>
T& StrideSpanIterator<T>::operator*() const
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT(ele_ >= start_ && ele_ < end_, "Attempting to dereference invalid iterator.");
#endif
    return *ele_;
}

template<typename T>
T* StrideSpanIterator<T>::operator->() const
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT(ele_ >= start_ && ele_ < end_, "Attempting to dereference invalid iterator.");
#endif
    return ele_;
}
template<typename T>
StrideSpanIterator<T>& StrideSpanIterator<T>::operator++()
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT(ele_ < end_, "Attempting to move iterator past end.");
#endif
    ele_ = byteoffset(ele_, strideBytes_);
    return *this;
}

template<typename T>
StrideSpanIterator<T> StrideSpanIterator<T>::operator++(int)
{
    StrideSpanIterator copy(*this);
    ++(*this);
    return copy;
}

template<typename T>
StrideSpanIterator<T>& StrideSpanIterator<T>::operator--()
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT(ele_ > start_, "Attempting to move iterator before start.");
#endif
    ele_ = byteoffset(ele_, -strideBytes_);
    return *this;
}

template<typename T>
StrideSpanIterator<T> StrideSpanIterator<T>::operator--(int)
{
    StrideSpanIterator copy(*this);
    --(*this);
    return copy;
}

template<typename T>
StrideSpanIterator<T> StrideSpanIterator<T>::operator+(difference_type diff) const
{
#if MIJIN_CHECKED_ITERATORS
    MIJIN_ASSERT((byteoffset(ele_ + strideBytes_ * diff) <= end_) && (byteoffset(ele_, strideBytes_ * diff) >= start_),
                 "Creating out-of-bounds iterator.");
#endif
    return StrideSpanIterator(byteoffset(ele_, diff * strideBytes_), strideBytes_,
#if MIJIN_CHECKED_ITERATORS
        start_, end_
#else
        nullptr, nullptr
#endif
    );
}

template<typename T>
auto StrideSpan<T>::operator[](size_type index) -> reference
{
    MIJIN_ASSERT(index < size(), "Attempting to access StrideSpan out of bounds.");
    return *byteoffset(begin_, index * strideBytes_);
}

template<typename T>
auto StrideSpan<T>::operator[](size_type index) const -> const_reference
{
    MIJIN_ASSERT(index < size(), "Attempting to access StrideSpan out of bounds.");
    return *byteoffset(begin_, index * strideBytes_);
}
} // namespace mijin

#endif // !defined(MIJIN_CONTAINER_STRIDE_SPAN_HPP_INCLUDED)