#pragma once

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <limits>
#include <stdexcept>
#include <utility>
#include <vector>

namespace uc::detail {

// A compact directory over leaf sizes.  Level zero describes leaves, every
// next level groups `fanout` children, and the requested maximum number of
// levels is respected.  The top level uses binary search; descent only scans a
// single fanout-sized group.  A Fenwick side index supplies O(log L) prefix
// sums needed by stable-ID -> logical-index resolution.
class HierarchicalSizeIndex {
public:
    struct Location {
        std::size_t leaf = 0;
        std::size_t local = 0;
    };

    HierarchicalSizeIndex(std::size_t fanout = 16, std::size_t max_levels = 3)
        : fanout_(std::max<std::size_t>(2, fanout)),
          max_levels_(std::clamp<std::size_t>(max_levels, 1, 8)) {}

    void configure(std::size_t fanout, std::size_t max_levels) {
        const auto leaf_sizes = levels_.empty()
            ? std::vector<std::size_t>{} : levels_.front();
        HierarchicalSizeIndex replacement(fanout, max_levels);
        replacement.rebuild(leaf_sizes);
        swap(replacement);
    }

    void rebuild(const std::vector<std::size_t>& leaf_sizes) {
        // Build all allocation-owning state off to the side.  A failed
        // allocation must not leave a directory that no longer describes its
        // storage.
        std::vector<std::vector<std::size_t>> new_levels;
        std::vector<std::size_t> new_top_prefix;
        std::vector<std::size_t> new_fenwick(leaf_sizes.size() + 1, 0);
        std::size_t new_total = 0;
        if (leaf_sizes.empty()) {
            levels_.swap(new_levels);
            top_prefix_.swap(new_top_prefix);
            fenwick_.swap(new_fenwick);
            total_ = 0;
            return;
        }

        new_levels.push_back(leaf_sizes);
        // Keep the top directory at most fanout-sized.  Binary search handles
        // that top level, so building extra single-parent levels only adds a
        // lookup and cannot improve asymptotic complexity.
        while (new_levels.back().size() > fanout_
               && new_levels.size() < max_levels_) {
            const auto& children = new_levels.back();
            std::vector<std::size_t> parents;
            parents.reserve((children.size() + fanout_ - 1) / fanout_);
            for (std::size_t begin = 0; begin < children.size(); begin += fanout_) {
                const auto end = std::min(children.size(), begin + fanout_);
                std::size_t weight = 0;
                for (auto i = begin; i < end; ++i) {
                    weight += children[i];
                }
                parents.push_back(weight);
            }
            new_levels.push_back(std::move(parents));
        }

        for (const auto weight : new_levels.back()) {
            new_total += weight;
            new_top_prefix.push_back(new_total);
        }
        for (std::size_t i = 0; i < leaf_sizes.size(); ++i) {
            for (auto node = i + 1; node < new_fenwick.size();
                 node += node & (~node + 1)) {
                new_fenwick[node] += leaf_sizes[i];
            }
        }

        levels_.swap(new_levels);
        top_prefix_.swap(new_top_prefix);
        fenwick_.swap(new_fenwick);
        total_ = new_total;
    }

    void swap(HierarchicalSizeIndex& other) noexcept {
        using std::swap;
        swap(fanout_, other.fanout_);
        swap(max_levels_, other.max_levels_);
        swap(total_, other.total_);
        levels_.swap(other.levels_);
        top_prefix_.swap(other.top_prefix_);
        fenwick_.swap(other.fenwick_);
    }

    void update(std::size_t leaf, std::ptrdiff_t delta) {
        if (levels_.empty() || leaf >= levels_.front().size()) {
            throw std::out_of_range("directory leaf index out of range");
        }
        apply_delta(levels_[0][leaf], delta);
        std::size_t node = leaf;
        for (std::size_t level = 1; level < levels_.size(); ++level) {
            node /= fanout_;
            apply_delta(levels_[level][node], delta);
        }
        total_ = apply_delta_copy(total_, delta);
        fenwick_add_signed(leaf, delta);
        rebuild_top_prefix();
    }

    [[nodiscard]] Location locate(std::size_t logical_index) const {
        if (logical_index >= total_ || levels_.empty()) {
            throw std::out_of_range("logical index out of range");
        }

        const auto top_it = std::upper_bound(top_prefix_.begin(), top_prefix_.end(), logical_index);
        std::size_t node = static_cast<std::size_t>(top_it - top_prefix_.begin());
        std::size_t remaining = logical_index - (node == 0 ? 0 : top_prefix_[node - 1]);

        for (std::size_t level = levels_.size() - 1; level > 0; --level) {
            const auto& children = levels_[level - 1];
            const auto begin = node * fanout_;
            const auto end = std::min(children.size(), begin + fanout_);
            auto child = begin;
            for (; child < end; ++child) {
                if (remaining < children[child]) {
                    break;
                }
                remaining -= children[child];
            }
            assert(child < end);
            node = child;
        }
        return {node, remaining};
    }

    [[nodiscard]] std::size_t prefix_before(std::size_t leaf) const noexcept {
        std::size_t sum = 0;
        for (std::size_t i = leaf; i > 0; i -= i & (~i + 1)) {
            sum += fenwick_[i];
        }
        return sum;
    }

    [[nodiscard]] std::size_t total() const noexcept { return total_; }
    [[nodiscard]] std::size_t actual_levels() const noexcept { return levels_.size(); }
    [[nodiscard]] std::size_t fanout() const noexcept { return fanout_; }

    [[nodiscard]] std::size_t allocated_bytes() const noexcept {
        std::size_t bytes = fenwick_.capacity() * sizeof(std::size_t)
                          + top_prefix_.capacity() * sizeof(std::size_t);
        for (const auto& level : levels_) {
            bytes += level.capacity() * sizeof(std::size_t);
        }
        return bytes;
    }

private:
    static void apply_delta(std::size_t& value, std::ptrdiff_t delta) {
        value = apply_delta_copy(value, delta);
    }

    static std::size_t apply_delta_copy(std::size_t value, std::ptrdiff_t delta) {
        if (delta < 0) {
            const auto magnitude = static_cast<std::size_t>(-delta);
            assert(value >= magnitude);
            return value - magnitude;
        }
        return value + static_cast<std::size_t>(delta);
    }

    void rebuild_top_prefix() {
        top_prefix_.clear();
        top_prefix_.reserve(levels_.back().size());
        std::size_t prefix = 0;
        for (const auto weight : levels_.back()) {
            prefix += weight;
            top_prefix_.push_back(prefix);
        }
    }

    void fenwick_add(std::size_t leaf, std::size_t delta) noexcept {
        for (auto i = leaf + 1; i < fenwick_.size(); i += i & (~i + 1)) {
            fenwick_[i] += delta;
        }
    }

    void fenwick_add_signed(std::size_t leaf, std::ptrdiff_t delta) noexcept {
        for (auto i = leaf + 1; i < fenwick_.size(); i += i & (~i + 1)) {
            fenwick_[i] = apply_delta_copy(fenwick_[i], delta);
        }
    }

    std::size_t fanout_;
    std::size_t max_levels_;
    std::size_t total_ = 0;
    std::vector<std::vector<std::size_t>> levels_;
    std::vector<std::size_t> top_prefix_;
    std::vector<std::size_t> fenwick_;
};

} // namespace uc::detail