#pragma once

#include "containers/sparse_set.hpp"

#include <cstddef>   // std::size_t
#include <format>    // std::format
#include <stdexcept> // std::out_of_range
#include <utility>   // std::get
#include <vector>    // values container

namespace kuiper
{

namespace containers
{

    template<typename K, typename V>
    class sparse_map {
      public:
        using sparse_set_t       = containers::sparse_set<K>;
        using values_container_t = std::vector<V>;

      public:
        // Capacity

        /// Reserve a number of slots in the underlying containers
        void reserve(std::size_t size) {
            if (size > m_capacity) {
                m_values.resize(size);
                m_capacity = size;
            }

            return;
        }

        // Access

        /// Access value specified by key (non-const)
        V& at(const K& key) {
            if (!m_sparse_set.contains(key)) {
                throw std::out_of_range(std::format("key \"{}\" does not exist in the sparse map", key));
            }

            const auto dense_idx = m_sparse_set.sparse()[key];
            return m_values[dense_idx];
        }

        /// Access value specified by key (const, read-only)
        const V& at(const K& key) const {
            if (!m_sparse_set.contains(key)) {
                throw std::out_of_range(std::format("key \"{}\" does not exist in the sparse map", key));
            }

            const auto dense_idx = m_sparse_set.sparse()[key];
            return m_values[dense_idx];
        }

        // Modifiers

        void insert(const K& key, const V& value) {
            // 1. Insert key into the underlying sparse set
            // 2. Insert value into the corresponding index in the sparse map

            if (!m_sparse_set.contains(key)) {
                // Key does not yet exist in set

                const auto insertion = m_sparse_set.insert(key);

                if (!std::get<1>(insertion))
                    return;

                const auto dense_idx = std::get<0>(insertion);
                reserve(dense_idx + 1); // no-op if there's already enough space
                m_values[dense_idx] = value;
                ++m_size;

                return;
            } else {
                // Key already exists in set

                const auto dense_idx = m_sparse_set.sparse()[key];
                m_values[dense_idx]  = value;
                ++m_size;

                return;
            }
        }

        /// Remove an element from the map
        void erase(const K& key) {
            // 1. Retrieve dense array index from set
            // 2. Swap to back
            // 3. Erase from set
            // 4. Decrement number of elements

            if (contains(key)) {
                const auto dense_idx = m_sparse_set.sparse()[key];
                m_values[dense_idx]  = m_values[m_size - 1];
                m_sparse_set.erase(key);
                --m_size;
            }

            return;
        }

        // Lookup

        /// Check if the key exists in the map
        bool contains(const K& key) const {
            return m_sparse_set.contains(key);
        }

        // Observers

        /// Access to the underlying sparse_set container
        const sparse_set_t& sparse_set() const noexcept {
            return m_sparse_set;
        }

        /// Access to the underlying values container
        const values_container_t& values() const noexcept {
            return m_values;
        }

        // Operator overloads

        /// Access/insert element specified by key
        V& operator[](const K& key) {
            const auto k = key;

            if (!contains(k)) {
                insert(k, {});
            }

            return at(k);
        }

        /// Access element specified by key (const, read-only)
        const V& operator[](const K& key) const {
            return at(key);
        }

      private:
        containers::sparse_set<K> m_sparse_set {};
        values_container_t        m_values {};
        std::size_t               m_size     = 0;
        std::size_t               m_capacity = 0;
    };

} // namespace containers

} // namespace kuiper