#pragma once

#include "universal_container/hierarchical_size_index.hpp"
#include "universal_container/ring_block.hpp"

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

namespace uc {

struct TieredConfig {
    std::size_t leaf_capacity = 512;
    std::size_t directory_fanout = 64;
    std::size_t directory_levels = 4;

    friend bool operator==(const TieredConfig&, const TieredConfig&) = default;
};

namespace detail {

template <class T>
class TieredStorage {
public:
    using value_type = T;
    using leaf_id_type = std::uint32_t;
    static constexpr std::size_t npos = std::numeric_limits<std::size_t>::max();

    explicit TieredStorage(TieredConfig config = {})
        : config_(normalize(config)),
          directory_(config_.directory_fanout, config_.directory_levels) {}

    TieredStorage(const TieredStorage& other)
        : config_(other.config_),
          directory_(config_.directory_fanout, config_.directory_levels),
          size_(other.size_),
          next_leaf_id_(other.next_leaf_id_),
          leaf_positions_(other.leaf_positions_) {
        leaves_.reserve(other.leaves_.size());
        for (const auto& leaf : other.leaves_) {
            leaves_.push_back(std::make_unique<Leaf>(*leaf));
        }
        rebuild_directory();
    }

    TieredStorage(TieredStorage&&) noexcept = default;

    TieredStorage& operator=(const TieredStorage& other) {
        if (this == &other) {
            return *this;
        }
        TieredStorage copy(other);
        swap(copy);
        return *this;
    }

    TieredStorage& operator=(TieredStorage&&) noexcept = default;
    ~TieredStorage() = default;

    void swap(TieredStorage& other) noexcept {
        using std::swap;
        swap(config_, other.config_);
        swap(leaves_, other.leaves_);
        swap(directory_, other.directory_);
        swap(size_, other.size_);
        swap(next_leaf_id_, other.next_leaf_id_);
        swap(leaf_positions_, other.leaf_positions_);
    }

    [[nodiscard]] std::size_t size() const noexcept { return size_; }
    [[nodiscard]] bool empty() const noexcept { return size_ == 0; }
    [[nodiscard]] const TieredConfig& config() const noexcept { return config_; }
    [[nodiscard]] std::size_t leaf_count() const noexcept { return leaves_.size(); }
    [[nodiscard]] std::size_t actual_levels() const noexcept { return directory_.actual_levels(); }

    // Changing only the directory geometry does not move elements or leaves.
    // It is substantially cheaper than rebuilding with another leaf capacity.
    void reconfigure_directory(std::size_t fanout, std::size_t maximum_levels) {
        auto requested = config_;
        requested.directory_fanout = fanout;
        requested.directory_levels = maximum_levels;
        requested = normalize(requested);
        if (requested.directory_fanout == config_.directory_fanout
            && requested.directory_levels == config_.directory_levels) {
            return;
        }

        // Build the complete replacement before changing either the active
        // directory or the public configuration.  HierarchicalSizeIndex::rebuild
        // allocates several vectors and may throw; mutating directory_ first
        // would otherwise leave locate() unusable after an allocation failure.
        std::vector<std::size_t> sizes;
        sizes.reserve(leaves_.size());
        for (const auto& leaf : leaves_) {
            sizes.push_back(leaf->values.size());
        }
        HierarchicalSizeIndex rebuilt(requested.directory_fanout,
                                      requested.directory_levels);
        rebuilt.rebuild(sizes);

        static_assert(std::is_nothrow_move_assignable_v<HierarchicalSizeIndex>);
        directory_ = std::move(rebuilt);
        config_.directory_fanout = requested.directory_fanout;
        config_.directory_levels = requested.directory_levels;
    }

    T& operator[](std::size_t index) noexcept {
        const auto location = locate_unchecked(index);
        return leaves_[location.leaf]->values[location.local];
    }

    const T& operator[](std::size_t index) const noexcept {
        const auto location = locate_unchecked(index);
        return leaves_[location.leaf]->values[location.local];
    }

    T& at(std::size_t index) {
        check_index(index);
        return (*this)[index];
    }

    const T& at(std::size_t index) const {
        check_index(index);
        return (*this)[index];
    }

    template <class Relocate>
    void push_back(T value, Relocate&& relocate) {
        if (leaves_.empty() || leaves_.back()->values.full()) {
            leaves_.push_back(make_leaf());
            rebuild_positions();
            rebuild_directory();
        }
        auto& leaf = *leaves_.back();
        leaf.values.push_back(std::move(value));
        ++size_;
        directory_.update(leaves_.size() - 1, 1);
        relocate(leaf.values[leaf.values.size() - 1], leaf.id, leaf.values.size() - 1);
    }

    void push_back(T value) {
        push_back(std::move(value), NoRelocate{});
    }

    template <class Relocate>
    void insert(std::size_t index, T value, Relocate&& relocate) {
        if (index > size_) {
            throw std::out_of_range("tiered insertion index out of range");
        }
        if (index == size_) {
            push_back(std::move(value), std::forward<Relocate>(relocate));
            return;
        }

        const auto location = directory_.locate(index);
        auto& leaf = *leaves_[location.leaf];
        if (!leaf.values.full()) {
            leaf.values.insert(location.local, std::move(value));
            ++size_;
            directory_.update(location.leaf, 1);
            refresh_leaf(location.leaf, relocate);
            return;
        }

        split_and_insert(location.leaf, location.local, std::move(value), relocate);
        ++size_;
        rebuild_directory();
    }

    void insert(std::size_t index, T value) {
        insert(index, std::move(value), NoRelocate{});
    }

    template <class Relocate>
    T erase(std::size_t index, Relocate&& relocate) {
        check_index(index);
        const auto location = directory_.locate(index);
        auto removed = leaves_[location.leaf]->values.erase(location.local);
        --size_;

        if (leaves_[location.leaf]->values.empty() && leaves_.size() > 1) {
            retire_leaf(location.leaf);
            rebuild_positions();
            rebuild_directory();
            return removed;
        }

        if (try_merge(location.leaf, relocate)) {
            return removed;
        }

        directory_.update(location.leaf, -1);
        refresh_leaf(location.leaf, relocate);
        return removed;
    }

    T erase(std::size_t index) {
        return erase(index, NoRelocate{});
    }

    void clear() noexcept {
        leaves_.clear();
        leaf_positions_.clear();
        size_ = 0;
        next_leaf_id_ = 0;
        directory_.rebuild({});
    }

    template <class Function>
    void for_each(Function&& function) {
        for (auto& leaf : leaves_) {
            for (std::size_t i = 0; i < leaf->values.size(); ++i) {
                function(leaf->values[i]);
            }
        }
    }

    template <class Function>
    void for_each(Function&& function) const {
        for (const auto& leaf : leaves_) {
            for (std::size_t i = 0; i < leaf->values.size(); ++i) {
                function(leaf->values[i]);
            }
        }
    }

    template <class Function>
    void for_each_with_location(Function&& function) const {
        for (const auto& leaf : leaves_) {
            for (std::size_t i = 0; i < leaf->values.size(); ++i) {
                function(leaf->values[i], leaf->id, i);
            }
        }
    }

    [[nodiscard]] std::vector<T> to_vector_copy() const {
        std::vector<T> result;
        result.reserve(size_);
        for_each([&](const T& value) { result.push_back(value); });
        return result;
    }

    [[nodiscard]] std::vector<T> to_vector_move() {
        std::vector<T> result;
        result.reserve(size_);
        for_each([&](T& value) { result.push_back(std::move(value)); });
        return result;
    }

    template <class Relocate>
    static TieredStorage from_vector(std::vector<T>&& source,
                                     TieredConfig config,
                                     Relocate&& relocate) {
        TieredStorage result(config);
        if (source.empty()) {
            return result;
        }
        const auto target_occupancy = bulk_target_occupancy(
            result.config_.leaf_capacity);
        const auto count = (source.size() + target_occupancy - 1)
                         / target_occupancy;
        result.leaves_.reserve(count);
        for (auto& value : source) {
            if (result.leaves_.empty()
                || result.leaves_.back()->values.size() == target_occupancy) {
                result.leaves_.push_back(result.make_leaf());
            }
            result.leaves_.back()->values.push_back(std::move(value));
            ++result.size_;
        }
        result.rebuild_positions();
        result.rebuild_directory();
        for (std::size_t leaf = 0; leaf < result.leaves_.size(); ++leaf) {
            result.refresh_leaf(leaf, relocate);
        }
        return result;
    }

    static TieredStorage from_vector(std::vector<T>&& source, TieredConfig config) {
        return from_vector(std::move(source), config, NoRelocate{});
    }

    static TieredStorage from_vector_copy(const std::vector<T>& source,
                                          TieredConfig config) {
        TieredStorage result(config);
        if (source.empty()) {
            return result;
        }
        const auto target_occupancy = bulk_target_occupancy(
            result.config_.leaf_capacity);
        const auto count = (source.size() + target_occupancy - 1)
                         / target_occupancy;
        result.leaves_.reserve(count);
        for (const auto& value : source) {
            if (result.leaves_.empty()
                || result.leaves_.back()->values.size() == target_occupancy) {
                result.leaves_.push_back(result.make_leaf());
            }
            result.leaves_.back()->values.push_back(value);
            ++result.size_;
        }
        result.rebuild_positions();
        result.rebuild_directory();
        return result;
    }

    static TieredStorage reconfigured_copy(const TieredStorage& source,
                                            TieredConfig config) {
        TieredStorage result(config);
        if (source.empty()) {
            return result;
        }
        const auto target_occupancy = bulk_target_occupancy(
            result.config_.leaf_capacity);
        const auto count = (source.size() + target_occupancy - 1)
                         / target_occupancy;
        result.leaves_.reserve(count);
        source.for_each([&](const T& value) {
            if (result.leaves_.empty()
                || result.leaves_.back()->values.size() == target_occupancy) {
                result.leaves_.push_back(result.make_leaf());
            }
            result.leaves_.back()->values.push_back(value);
            ++result.size_;
        });
        result.rebuild_positions();
        result.rebuild_directory();
        return result;
    }

    [[nodiscard]] std::size_t logical_index(leaf_id_type leaf_id,
                                             std::size_t local) const {
        if (leaf_id >= leaf_positions_.size()) {
            throw std::out_of_range("unknown tiered leaf id");
        }
        const auto position = leaf_positions_[leaf_id];
        if (position == npos || local >= leaves_[position]->values.size()) {
            throw std::out_of_range("stale tiered location");
        }
        return directory_.prefix_before(position) + local;
    }

    [[nodiscard]] std::size_t allocated_bytes() const noexcept {
        std::size_t bytes = leaves_.capacity() * sizeof(typename decltype(leaves_)::value_type)
                          + leaf_positions_.capacity() * sizeof(std::size_t)
                          + directory_.allocated_bytes();
        for (const auto& leaf : leaves_) {
            bytes += sizeof(Leaf) + leaf->values.allocated_bytes();
        }
        return bytes;
    }

private:
    struct NoRelocate {
        void operator()(const T&, leaf_id_type, std::size_t) const noexcept {}
    };

    struct Leaf {
        Leaf(leaf_id_type leaf_id, std::size_t capacity)
            : id(leaf_id), values(capacity) {}

        leaf_id_type id;
        RingBlock<T> values;
    };

    static TieredConfig normalize(TieredConfig config) {
        config.leaf_capacity = std::clamp<std::size_t>(config.leaf_capacity, 4, 1U << 20U);
        config.directory_fanout = std::clamp<std::size_t>(config.directory_fanout, 2, 1U << 16U);
        config.directory_levels = std::clamp<std::size_t>(config.directory_levels, 1, 8);
        return config;
    }

    // Bulk-loading completely full leaves makes the first random insertion
    // into almost every leaf pay for a split and a directory rebuild.  Keep a
    // small, deterministic reserve instead.  Seven eighths preserves compact
    // memory usage while giving each freshly built leaf enough room for the
    // short edit bursts that motivate switching away from vector storage.
    [[nodiscard]] static std::size_t
    bulk_target_occupancy(std::size_t capacity) noexcept {
        return std::max<std::size_t>(1, capacity - capacity / 8);
    }

    [[nodiscard]] std::unique_ptr<Leaf> make_leaf() {
        if (next_leaf_id_ == std::numeric_limits<leaf_id_type>::max()) {
            throw std::length_error("tiered leaf id space exhausted");
        }
        const auto id = next_leaf_id_++;
        if (leaf_positions_.size() <= id) {
            leaf_positions_.resize(static_cast<std::size_t>(id) + 1, npos);
        }
        return std::make_unique<Leaf>(id, config_.leaf_capacity);
    }

    void check_index(std::size_t index) const {
        if (index >= size_) {
            throw std::out_of_range("tiered index out of range");
        }
    }

    [[nodiscard]] typename HierarchicalSizeIndex::Location
    locate_unchecked(std::size_t index) const noexcept {
        assert(index < size_);
        try {
            return directory_.locate(index);
        } catch (...) {
            std::terminate();
        }
    }

    void rebuild_positions() {
        std::fill(leaf_positions_.begin(), leaf_positions_.end(), npos);
        for (std::size_t i = 0; i < leaves_.size(); ++i) {
            const auto id = leaves_[i]->id;
            if (leaf_positions_.size() <= id) {
                leaf_positions_.resize(static_cast<std::size_t>(id) + 1, npos);
            }
            leaf_positions_[id] = i;
        }
    }

    void rebuild_directory() {
        std::vector<std::size_t> sizes;
        sizes.reserve(leaves_.size());
        for (const auto& leaf : leaves_) {
            sizes.push_back(leaf->values.size());
        }
        directory_.rebuild(sizes);
        assert(directory_.total() == size_);
    }

    template <class Relocate>
    void refresh_leaf(std::size_t leaf_index, Relocate& relocate) {
        auto& leaf = *leaves_[leaf_index];
        for (std::size_t i = 0; i < leaf.values.size(); ++i) {
            relocate(leaf.values[i], leaf.id, i);
        }
    }

    template <class Relocate>
    void split_and_insert(std::size_t leaf_index,
                          std::size_t local,
                          T value,
                          Relocate& relocate) {
        auto& old = *leaves_[leaf_index];
        std::vector<T> combined;
        combined.reserve(old.values.size() + 1);
        for (std::size_t i = 0; i < old.values.size(); ++i) {
            combined.push_back(std::move(old.values[i]));
        }
        combined.insert(combined.begin() + static_cast<std::ptrdiff_t>(local), std::move(value));

        const auto left_id = old.id;
        auto left = std::make_unique<Leaf>(left_id, config_.leaf_capacity);
        auto right = make_leaf();
        const auto middle = combined.size() / 2;
        for (std::size_t i = 0; i < middle; ++i) {
            left->values.push_back(std::move(combined[i]));
        }
        for (std::size_t i = middle; i < combined.size(); ++i) {
            right->values.push_back(std::move(combined[i]));
        }

        leaves_[leaf_index] = std::move(left);
        leaves_.insert(leaves_.begin() + static_cast<std::ptrdiff_t>(leaf_index + 1),
                       std::move(right));
        rebuild_positions();
        refresh_leaf(leaf_index, relocate);
        refresh_leaf(leaf_index + 1, relocate);
    }

    void retire_leaf(std::size_t leaf_index) {
        const auto id = leaves_[leaf_index]->id;
        if (id < leaf_positions_.size()) {
            leaf_positions_[id] = npos;
        }
        leaves_.erase(leaves_.begin() + static_cast<std::ptrdiff_t>(leaf_index));
    }

    template <class Relocate>
    bool try_merge(std::size_t leaf_index, Relocate& relocate) {
        if (leaves_.size() < 2) {
            return false;
        }
        const auto threshold = std::max<std::size_t>(1, config_.leaf_capacity / 4);
        if (leaves_[leaf_index]->values.size() >= threshold) {
            return false;
        }

        std::size_t left_index = leaf_index;
        std::size_t right_index = leaf_index + 1;
        if (right_index >= leaves_.size()) {
            left_index = leaf_index - 1;
            right_index = leaf_index;
        }
        const auto combined_size = leaves_[left_index]->values.size()
                                 + leaves_[right_index]->values.size();
        if (combined_size > config_.leaf_capacity) {
            return false;
        }

        const auto survivor_id = leaves_[left_index]->id;
        auto merged = std::make_unique<Leaf>(survivor_id, config_.leaf_capacity);
        for (std::size_t i = 0; i < leaves_[left_index]->values.size(); ++i) {
            merged->values.push_back(std::move(leaves_[left_index]->values[i]));
        }
        for (std::size_t i = 0; i < leaves_[right_index]->values.size(); ++i) {
            merged->values.push_back(std::move(leaves_[right_index]->values[i]));
        }

        const auto retired_id = leaves_[right_index]->id;
        leaves_[left_index] = std::move(merged);
        leaves_.erase(leaves_.begin() + static_cast<std::ptrdiff_t>(right_index));
        if (retired_id < leaf_positions_.size()) {
            leaf_positions_[retired_id] = npos;
        }
        rebuild_positions();
        rebuild_directory();
        refresh_leaf(left_index, relocate);
        return true;
    }

    TieredConfig config_;
    std::vector<std::unique_ptr<Leaf>> leaves_;
    HierarchicalSizeIndex directory_;
    std::size_t size_ = 0;
    leaf_id_type next_leaf_id_ = 0;
    std::vector<std::size_t> leaf_positions_;
};

} // namespace detail
} // namespace uc