#pragma once

#include "value.hpp"
#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
#include <unordered_map>
#include <variant>
#include <vector>

namespace sliger::heap {

struct Array {
    std::vector<Value> values;
};

struct Struct {
    std::unordered_map<std::string, Value> fields;
};

enum class AllocType {
    Value,
    Array,
    Struct,
};

// clang-format off
template <AllocType type> struct AllocTypeType {};
template <> struct AllocTypeType<AllocType::Value> { using Type = Value; };
template <> struct AllocTypeType<AllocType::Array> { using Type = Array; };
template <> struct AllocTypeType<AllocType::Struct> { using Type = Struct; };
// clang-format on

struct AllocItem {
    template <AllocType AT> inline auto as() & -> AllocTypeType<AT>::Type&
    {
        return std::get<typename AllocTypeType<AT>::Type>(this->value);
    }
    template <AllocType AT>
    inline auto as() const& -> const AllocTypeType<AT>::Type&
    {
        return std::get<typename AllocTypeType<AT>::Type>(this->value);
    }
    template <AllocType AT> inline auto as() && -> AllocTypeType<AT>::Type&&
    {
        return std::move(
            std::get<typename AllocTypeType<AT>::Type>(this->value));
    }
    template <AllocType AT>
    inline auto as() const&& -> const AllocTypeType<AT>::Type&&
    {
        return std::move(
            std::get<typename AllocTypeType<AT>::Type>(this->value));
    }

    AllocType type;
    std::variant<Value, Array, Struct> value;
};

enum class ErrType {
    InvalidPtrAccess,
    CannotAllocate,
};

template <typename T> struct Res {
    Res(T val)
        : val(std::forward<T>(val))
    {
    }
    Res(ErrType err)
        : err(err)
    {
    }

    std::optional<T> val;
    std::optional<ErrType> err;
};

class Heap {
public:
    inline auto can_allocate() const -> bool
    {
        return this->sel->size() < this->max_size;
    }

    inline void collect(const std::vector<uint32_t>& ptr_stack_values)
    {
        this->other->reserve(this->max_size);
        for (auto ptr : ptr_stack_values) {
            move_item_to_other(ptr);
        }
        this->sel->clear();
        std::swap(this->sel, this->other);
        if (this->sel->size() + 1 >= this->max_size) {
            this->max_size *= 2;
            this->sel->reserve(this->max_size);
        } else if (this->sel->size() * 2 < this->max_size) {
            this->max_size /= 2;
            this->sel->shrink_to_fit();
            this->sel->reserve(this->max_size);
        }
    }

    inline auto at(uint32_t ptr) -> Res<AllocItem*>
    {
        if (ptr >= this->sel->size()) {
            return ErrType::InvalidPtrAccess;
        }
        return &this->sel->at(ptr);
    }

    template <AllocType type> auto alloc() -> Res<uint32_t>
    {
        if (not can_allocate()) {
            return ErrType::CannotAllocate;
        }
        auto ptr = static_cast<uint32_t>(this->sel->size());
        this->sel->push_back(typename AllocTypeType<type>::Type {});
        return ptr;
    }

private:
    inline void move_item_to_other(uint32_t ptr)
    {
        auto res = at(ptr);
        if (!res.val.has_value())
            return;
        auto val = res.val.value();
        switch (val->type) {
            case AllocType::Value: {
                auto& v = val->as<AllocType::Value>();
                if (v.type() == ValueType::Ptr) {
                    move_item_to_other(v.template as<ValueType::Ptr>().value);
                }
                break;
            }
            case AllocType::Array: {
                auto& vs = val->as<AllocType::Array>();
                for (auto& v : vs.values) {
                    if (v.type() == ValueType::Ptr) {
                        move_item_to_other(
                            v.template as<ValueType::Ptr>().value);
                    }
                }
                break;
            }
            case AllocType::Struct: {
                auto& vs = val->as<AllocType::Struct>();
                for (auto& [key, v] : vs.fields) {
                    if (v.type() == ValueType::Ptr) {
                        move_item_to_other(
                            v.template as<ValueType::Ptr>().value);
                    }
                }
                break;
            }
        }
    }

    size_t max_size = 4;

    std::vector<AllocItem> heap_1;
    std::vector<AllocItem> heap_2;

    std::vector<AllocItem>* sel = &heap_1;
    std::vector<AllocItem>* other = &heap_2;
};

}