#pragma once

#if !defined(MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED)
#define MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED 1

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

#include "./boxed_object.hpp"
#include "./optional.hpp"

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

namespace mijin
{
template<typename TKey, typename TValue>
struct VectorMapIterator
{
public:
    using difference_type = std::ptrdiff_t;
private:
    Optional<std::pair<TKey&, TValue&>> ref_;
public:
    VectorMapIterator(TKey* key, TValue* value)
    {
        ref_.emplace(std::tie(*key, *value));
    }
    VectorMapIterator(const VectorMapIterator& other) MIJIN_NOEXCEPT : VectorMapIterator(&other.ref_->first, &other.ref_->second) {}
    ~VectorMapIterator() noexcept
    {
        ref_.reset();
    }

    VectorMapIterator& operator=(const VectorMapIterator& other) MIJIN_NOEXCEPT
    {
        if (this != &other)
        {
            ref_.reset();
            ref_.emplace(std::tie(other.ref_->first, other.ref_->second));
        }
        return *this;
    }
    
    bool operator==(const VectorMapIterator& other) const MIJIN_NOEXCEPT { return &ref_->first == &other.ref_->first; }
    bool operator!=(const VectorMapIterator& other) const MIJIN_NOEXCEPT { return &ref_->first != &other.ref_->first; }
    
    std::pair<TKey&, TValue&> operator*() const MIJIN_NOEXCEPT
    {
        return *ref_;
    }
    const std::pair<TKey&, TValue&>* operator->() const MIJIN_NOEXCEPT
    {
        return &*ref_;
    }

    VectorMapIterator& operator++() MIJIN_NOEXCEPT
    {
        TKey* oldKey = &ref_->first;
        TValue* oldValue = &ref_->second;
        ref_.reset();
        ref_.emplace(std::tie(*(++oldKey), *(++oldValue)));
        return *this;
    }

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

    VectorMapIterator& operator--() MIJIN_NOEXCEPT
    {
        TKey* oldKey = &ref_->first;
        TValue* oldValue = &ref_->second;
        ref_.reset();
        ref_.emplace(std::tie(*(--oldKey), *(--oldValue)));
        return *this;
    }

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

    VectorMapIterator operator+(difference_type diff) const MIJIN_NOEXCEPT
    {
        return VectorMapIterator(&ref_->first + diff, &ref_->second + diff);
    }

    VectorMapIterator operator-(difference_type diff) const MIJIN_NOEXCEPT
    {
        return VectorMapIterator(&ref_->first - diff, &ref_->second - diff);
    }
    
    difference_type operator-(const VectorMapIterator& other) const MIJIN_NOEXCEPT
    {
        return &ref_->first - &other.ref_->first;
    }
};

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 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())) {}
    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)) {}
    VectorMap(const VectorMap&) = default;
    VectorMap(VectorMap&&) = default;
    
    VectorMap& operator=(const VectorMap&) = default;
    VectorMap& operator=(VectorMap&&) = default;
    auto operator<=>(const VectorMap& other) const noexcept = default;
    
    template<typename TIndex>
    TValue& operator[](const TIndex& key)
    {
        auto it = find(key);
        if (it != end())
        {
            return it->second;
        }
        return emplace(key, TValue()).first->second;
    }
    
    template<typename TIndex>
    const TValue& operator[](const TIndex& key) const
    {
        return at(key);
    }

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

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

    [[nodiscard]]
    const TValue& at(const TKey& key) const
    {
        auto it = find(key);
        if (it == end())
        {
            throw std::out_of_range("key not found in map");
        }
        return it->second;
    }
    
    [[nodiscard]]
    bool empty() const MIJIN_NOEXCEPT { return keys_.empty(); }
    [[nodiscard]]
    size_type size() const MIJIN_NOEXCEPT { return keys_.size(); }
    [[nodiscard]]
    size_type max_size() const MIJIN_NOEXCEPT { return std::min(keys_.max_size(), values_.max_size()); }
    void reserve(std::size_t size)
    {
        keys_.reserve(size);
        values_.reserve(size);
    }
    
    void clear()
    {
        keys_.clear();
        values_.clear();
    }
    
    template<typename... TArgs>
    std::pair<iterator, bool> emplace(TArgs&&... args)
    {
        std::pair<TKey, TValue> 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);
    }
    
    iterator erase(iterator pos) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &pos->first - &keys_[0];
        return eraseImpl(idx);
    }
    
    iterator erase(const_iterator pos) MIJIN_NOEXCEPT
    {
        const std::ptrdiff_t idx = &pos->first - &keys_[0];
        return eraseImpl(idx);
    }

    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);
    }

    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);
    }

    template<typename TSearch>
    [[nodiscard]]
    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>
    [[nodiscard]]
    const_iterator find(const TSearch& key) const MIJIN_NOEXCEPT
    {
        for (std::size_t idx = 0; idx < keys_.size(); ++idx)
        {
            if (keys_[idx] == key)
            {
                return const_iterator(&keys_[idx], &values_[idx]);
            }
        }
        return end();
    }

    [[nodiscard]]
    bool contains(const TKey& key) const MIJIN_NOEXCEPT
    {
        return std::ranges::contains(keys_, key);
    }
private:
    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
    }
};
} // namespace mijin

#endif // MIJIN_CONTAINER_VECTOR_MAP_HPP_INCLUDED