#pragma once

#if !defined(MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED)
#define MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED 1

#include <algorithm>
#include <functional>
#include <stdexcept>
#include <vector>

#include "../internal/common.hpp"

namespace mijin
{
template<typename TKey, typename TValue>
struct VectorMapIterator
{
public:
    using key_t = TKey;
    using value_t = TValue;
    using pair_t = std::pair<const key_t&, value_t&>;
    struct Pointer
    {
        pair_t pair;

        constexpr pair_t* operator->() { return &pair; }
    };


    using difference_type = std::ptrdiff_t;
    using value_type = std::pair<key_t&, value_t&>;
    using pointer = Pointer;
    using reference = value_type;
    using iterator_category = std::random_access_iterator_tag;
private:
    const key_t* key_;
    value_t* value_;
public:
    constexpr VectorMapIterator(const key_t* key, value_t* value) MIJIN_NOEXCEPT : key_(key), value_(value)
    {
    }
    constexpr VectorMapIterator(const VectorMapIterator& other) noexcept = default;
    constexpr VectorMapIterator& operator=(const VectorMapIterator& other) MIJIN_NOEXCEPT = default;
    
    constexpr bool operator==(const VectorMapIterator& other) const MIJIN_NOEXCEPT { return key_ == other.key_; }
    constexpr bool operator!=(const VectorMapIterator& other) const MIJIN_NOEXCEPT { return key_ != other.key_; }
    
    constexpr reference operator*() const MIJIN_NOEXCEPT
    {
        return std::tie(*key_, *value_);
    }
    constexpr pointer operator->() const MIJIN_NOEXCEPT
    {
        return {std::tie(*key_, *value_)};
    }

    constexpr VectorMapIterator& operator++() MIJIN_NOEXCEPT
    {
        ++key_;
        ++value_;
        return *this;
    }

    constexpr VectorMapIterator operator++(int) const MIJIN_NOEXCEPT
    {
        VectorMapIterator copy(*this);
        ++(*this);
        return copy;
    }

    constexpr VectorMapIterator& operator--() MIJIN_NOEXCEPT
    {
        --key_;
        --value_;
        return *this;
    }

    constexpr VectorMapIterator operator--(int) const MIJIN_NOEXCEPT
    {
        VectorMapIterator copy(*this);
        --(*this);
        return copy;
    }

    constexpr VectorMapIterator operator+(difference_type diff) const MIJIN_NOEXCEPT
    {
        return VectorMapIterator(key_ + diff, value_ + diff);
    }

    constexpr VectorMapIterator operator-(difference_type diff) const MIJIN_NOEXCEPT
    {
        return VectorMapIterator(key_ - diff, value_ - diff);
    }
    
    constexpr difference_type operator-(const VectorMapIterator& other) const MIJIN_NOEXCEPT
    {
        return key_ - other.key_;
    }
};

template<typename TKey, typename TValue, typename TKeyAllocator = MIJIN_DEFAULT_ALLOCATOR<TKey>, typename TValueAllocator = MIJIN_DEFAULT_ALLOCATOR<TValue>>
class VectorMap
{
public:
    using key_type = TKey;
    using mapped_type = TValue;
    using value_type = std::pair<const TKey, TValue>;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using reference = value_type&;
    using const_reference = const value_type&;
    using iterator = VectorMapIterator<const TKey, TValue>;
    using const_iterator = VectorMapIterator<const TKey, const TValue>;
private:
    std::vector<TKey, TKeyAllocator> keys_;
    std::vector<TValue, TValueAllocator> values_;
public:
    explicit constexpr VectorMap(TKeyAllocator keyAllocator = {})
        MIJIN_NOEXCEPT_IF((std::is_nothrow_move_constructible_v<TKeyAllocator> && std::is_nothrow_constructible_v<TValueAllocator, const TKeyAllocator&>))
        : keys_(std::move(keyAllocator)), values_(TValueAllocator(keys_.get_allocator())) {}
    constexpr VectorMap(TKeyAllocator keyAllocator, TValueAllocator valueAllocator)
        MIJIN_NOEXCEPT_IF(std::is_nothrow_move_constructible_v<TKeyAllocator> && std::is_nothrow_move_constructible_v<TValueAllocator>)
        : keys_(std::move(keyAllocator)), values_(std::move(valueAllocator)) {}
    constexpr VectorMap(const VectorMap&) = default;
    constexpr VectorMap(VectorMap&&) = default;
    
    constexpr VectorMap& operator=(const VectorMap&) = default;
    constexpr VectorMap& operator=(VectorMap&&) = default;
    constexpr auto operator<=>(const VectorMap& other) const noexcept = default;
    
    template<typename TIndex>
    constexpr mapped_type& operator[](const TIndex& key)
    {
        auto it = find(key);
        if (it != end())
        {
            return it->second;
        }
        return emplace(key, mapped_type()).first->second;
    }
    
    template<typename TIndex>
    constexpr const mapped_type& operator[](const TIndex& key) const
    {
        return at(key);
    }

    constexpr mapped_type& operator[](key_type&& key)
    {
        auto it = find(key);
        if (it != end())
        {
            return it->second;
        }
        return emplace(std::move(key), mapped_type()).first->second;
    }
    
    [[nodiscard]]
    constexpr iterator begin() MIJIN_NOEXCEPT { return {keys_.data(), values_.data()}; }
    [[nodiscard]]
    constexpr const_iterator begin() const MIJIN_NOEXCEPT { return {keys_.data(), values_.data()}; }
    [[nodiscard]]
    constexpr const_iterator cbegin() const MIJIN_NOEXCEPT { return {keys_.data(), values_.data()}; }
    [[nodiscard]]
    constexpr iterator end() MIJIN_NOEXCEPT { return {keys_.data() + keys_.size(), values_.data() + values_.size()}; }
    [[nodiscard]]
    constexpr const_iterator end() const MIJIN_NOEXCEPT { return {keys_.data() + keys_.size(), values_.data() + values_.size()}; }
    [[nodiscard]]
    constexpr const_iterator cend() const MIJIN_NOEXCEPT { return {keys_.data() + keys_.size(), values_.data() + values_.size()}; }

    [[nodiscard]]
    constexpr mapped_type& at(const key_type& key)
    {
        auto it = find(key);
        if (it == end())
        {
            throw std::out_of_range("key not found in map");
        }
        return it->second;
    }

    [[nodiscard]]
    constexpr const mapped_type& at(const key_type& key) const
    {
        auto it = find(key);
        if (it == end())
        {
            throw std::out_of_range("key not found in map");
        }
        return it->second;
    }
    
    [[nodiscard]]
    constexpr bool empty() const MIJIN_NOEXCEPT { return keys_.empty(); }
    [[nodiscard]]
    constexpr size_type size() const MIJIN_NOEXCEPT { return keys_.size(); }
    [[nodiscard]]
    constexpr size_type max_size() const MIJIN_NOEXCEPT { return std::min(keys_.max_size(), values_.max_size()); }
    constexpr void reserve(std::size_t size)
    {
        keys_.reserve(size);
        values_.reserve(size);
    }
    
    constexpr void clear()
    {
        keys_.clear();
        values_.clear();
    }
    
    template<typename... TArgs>
    constexpr std::pair<iterator, bool> emplace(TArgs&&... args)
        MIJIN_NOEXCEPT_IF((std::is_nothrow_constructible_v<std::pair<key_type, mapped_type>, TArgs&&...>
                            && std::is_nothrow_move_constructible_v<key_type> && std::is_nothrow_move_constructible_v<mapped_type>))
    {
        std::pair<key_type, mapped_type> asPair(std::forward<TArgs>(args)...);
        auto it = find(asPair.first);
        if (it != end())
        {
            return std::make_pair(std::move(it), false);
        }
        keys_.push_back(std::move(asPair.first));
        values_.push_back(std::move(asPair.second));
        return std::make_pair(iterator(&keys_.back(), &values_.back()), true);
    }
    
    constexpr iterator erase(iterator pos) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &pos->first - &keys_[0];
        return eraseImpl(idx);
    }
    
    constexpr iterator erase(const_iterator pos) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &pos->first - &keys_[0];
        return eraseImpl(idx);
    }

    constexpr iterator erase(iterator first, iterator last) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &first->first - &keys_[0];
        const std::ptrdiff_t count = last - first;
        return eraseImpl(idx, count);
    }

    constexpr iterator erase(const_iterator first, const_iterator last) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &first->first - &keys_[0];
        const std::ptrdiff_t count = last - first;
        return eraseImpl(idx, count);
    }

    constexpr iterator eraseAt(std::size_t idx, std::size_t count = 1) MIJIN_NOEXCEPT
    {
        return eraseImpl(static_cast<std::ptrdiff_t>(idx), static_cast<std::ptrdiff_t>(count));
    }

    template<typename TSearch>
    [[nodiscard]]
    constexpr iterator find(const TSearch& key) MIJIN_NOEXCEPT
    {
        for (std::size_t idx = 0; idx < keys_.size(); ++idx)
        {
            if (keys_[idx] == key)
            {
                return iterator(&keys_[idx], &values_[idx]);
            }
        }
        return end();
    }

    template<typename TSearch, typename TProj = std::identity>
    [[nodiscard]]
    constexpr const_iterator find(const TSearch& key, TProj proj = {}) const MIJIN_NOEXCEPT
    {
        auto itKey = std::ranges::find(keys_, key, std::move(proj));
        if (itKey == keys_.end()) {
            return end();
        }
        const std::ptrdiff_t idx = std::distance(keys_.begin(), itKey);
        return const_iterator(&*itKey, &values_[idx]);
    }

    template<typename TSearch, typename TProj = std::identity>
    [[nodiscard]]
    constexpr bool contains(const TSearch& key, TProj proj = {}) const MIJIN_NOEXCEPT
    {
        return std::ranges::contains(keys_, key, std::move(proj));
    }

    [[nodiscard]]
    constexpr std::pair<const key_type&, mapped_type&> entry(std::size_t idx) MIJIN_NOEXCEPT
    {
        return std::tie(keys_.at(idx), values_.at(idx));
    }

    [[nodiscard]]
    constexpr std::pair<const key_type&, const mapped_type&> entry(std::size_t idx) const MIJIN_NOEXCEPT
    {
        return std::tie(keys_.at(idx), values_.at(idx));
    }
private:
    constexpr iterator eraseImpl(std::ptrdiff_t idx, std::ptrdiff_t count = 1) MIJIN_NOEXCEPT
    {
        auto itKey = keys_.begin() + idx;
        auto itValue = values_.begin() + idx;
        itKey = keys_.erase(itKey, itKey + count);
        itValue = values_.erase(itValue, itValue + count);
        return itKey == keys_.end() ? end() : iterator(&*itKey, &*itValue); // cannot dereference the iterators if the last element was removed
    }
};

#if 0
consteval VectorMap<int, int> testMap() {
    VectorMap<int, int> result;
    result.emplace(5, 17);
    result[0] = -10;
    return result;
};
static_assert(testMap().at(0) == -10);
#endif
} // namespace mijin

#endif // MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED