#pragma once

#include <variant>
#include <functional>
#include <type_traits>

// Virtual Function Information
template <typename T>
class __Function;
template <typename T>
class __VirtualFunctionInfo {
    __VirtualFunctionInfo(T *addr_, void *parent_);
    [[nodiscard]] bool can_overwrite() const;
    T *const addr;
    void *const parent;
    friend class __Function<std::remove_pointer_t<T>>;
};

// Thunks
typedef void *(*thunk_enabler_t)(void *target, void *thunk);
void enable_all_thunks(const thunk_enabler_t &thunk_enabler);

// Function Information
template <typename Ret, typename... Args>
class __Function<Ret(Args...)> {
public:
    // Types
    using ptr_type = Ret (*)(Args...);
    using type = std::function<Ret(Args...)>;
    using overwrite_type = std::function<Ret(type, Args...)>;

    // Normal Function
    __Function(const char *name_, ptr_type func_, ptr_type thunk_);
    // Virtual Function
    __Function(const char *name_, ptr_type *func_, void *parent, ptr_type thunk_);

    // Overwrite Function
    [[nodiscard]] bool overwrite(overwrite_type target) {
        // Check If Enabled
        if (!enabled) {
            return false;
        }
        // Overwrite
        type original = get_thunk_target();
        thunk_target = [original, target](Args... args) {
            return target(original, args...);
        };
        return true;
    }

    // Getters
    [[nodiscard]] ptr_type get_backup() const {
        return backup;
    }
    [[nodiscard]] ptr_type get() const {
        if (!enabled) {
            return nullptr;
        } else if (is_virtual) {
            return *get_vtable_addr();
        } else {
            return std::get<ptr_type>(func);
        }
    }
    [[nodiscard]] ptr_type *get_vtable_addr() const {
        if (is_virtual) {
            return std::get<__VirtualFunctionInfo<ptr_type>>(func).addr;
        } else {
            return nullptr;
        }
    }
    [[nodiscard]] const char *get_name() const {
        return name;
    }
    [[nodiscard]] type get_thunk_target() const {
        if (thunk_target) {
            return thunk_target;
        } else {
            return get_backup();
        }
    }

private:
    // Current Function
    const bool is_virtual;
    std::variant<ptr_type, __VirtualFunctionInfo<ptr_type>> func;

    // State
    const bool enabled;
    const char *const name;

    // Backup Of Original Function Pointer
    const ptr_type backup;

    // Thunk
    const ptr_type thunk;
    type thunk_target;
    void enable_thunk(const thunk_enabler_t &thunk_enabler);
    friend void enable_all_thunks(const thunk_enabler_t &);
};