// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once

#include <array>
#include <functional>
#include <optional>
#include <string>
#include <string_view>
#include <unordered_map>
#include <variant>
#include <vector>

namespace EFIBoot
{

extern "C"
{
#include "efivar-lite/efivar-lite.h"
#include "efivar-lite/load-option.h"
}

inline bool operator==(const efi_guid_t &first, const efi_guid_t &second)
{
    return efi_guid_cmp(&first, &second) == 0;
}

inline bool operator!=(const efi_guid_t &first, const efi_guid_t &second)
{
    return efi_guid_cmp(&first, &second) != 0;
}

using Raw_data = std::vector<uint8_t>;

template <class Type = Raw_data>
std::optional<Type> deserialize(const void *data, size_t data_size);

namespace File_path
{

{% for category in device_paths.values() %}
/*
   {{ category.name }}
   {{ category.description }}
*/
namespace {{ category.slug.upper() }}
{

{% for node in category.nodes %}
/*
   {{ node.name }}
   {{ node.description }}
*/
struct {{ node.slug.capitalize() }}
{
    {% for field in node.fields if field.type == "enum" %}
    enum class {{ field.slug.upper() }} : uint{{ field.size*8 }}_t
    {
        {% for enum in field.enum %}
        {{ enum.slug.upper() }} = {{ "0x%0x" | format(enum.value) }},
        {% endfor %}
    };

    {% endfor %}
    static const uint8_t TYPE = EFIDP_TYPE_{{ category.slug.upper() }};
    static const uint8_t SUBTYPE = EFIDP_{{ category.slug.upper() }}_{{ node.slug.upper() }};

    {% for field in node.fields %}
    {%- if field.type in ("guid", "ip4", "ip6", "mac") %}
    std::array<uint8_t, {{ field.size }}>
    {%- elif field.type == "wstring" %}
    std::u16string
    {%- elif field.type == "string" %}
    std::string
    {%- elif field.type in ("raw_data", "uri") %}
    Raw_data
    {%- elif field.type in ("int", "hex") %}
    uint{{ field.size * 8 }}_t
    {%- elif field.type == "enum" %}
    {{ field.slug.upper() }}
    {%- else %}
    {{ field.type }}
    {%- endif %} {{ field.slug }} = {};
    {% endfor %}
};

{% endfor %}
} // namespace {{ category.slug.upper() }}

{% endfor %}
struct Unknown
{
    Raw_data data = {};
    uint8_t TYPE = {};
    uint8_t SUBTYPE = {};
};

using ANY = std::variant<
{% for category in device_paths.values() %}{% for node in category.nodes %}
    {{ category.slug.upper() }}::{{ node.slug.capitalize() }},
{% endfor %}{% endfor %}
    Unknown>;

} // namespace File_path

enum class Load_option_attribute : uint32_t
{
    EMPTY = 0x00000000,

    ACTIVE = 0x00000001,
    FORCE_RECONNECT = 0x00000002,
    HIDDEN = 0x00000008,

    CATEGORY_MASK = 0x00001F00,
    CATEGORY_BOOT = 0x00000000,
    CATEGORY_APP = 0x00000100,
};

inline Load_option_attribute operator|(Load_option_attribute a, Load_option_attribute b)
{
    return static_cast<Load_option_attribute>(static_cast<std::underlying_type_t<Load_option_attribute>>(a) | static_cast<std::underlying_type_t<Load_option_attribute>>(b));
}

inline Load_option_attribute operator&(Load_option_attribute a, Load_option_attribute b)
{
    return static_cast<Load_option_attribute>(static_cast<std::underlying_type_t<Load_option_attribute>>(a) & static_cast<std::underlying_type_t<Load_option_attribute>>(b));
}

struct Load_option
{
    std::u16string description = u"";
    std::vector<File_path::ANY> device_path = {};
    Raw_data optional_data = {};
    Load_option_attribute attributes = Load_option_attribute::EMPTY;
};

using Progress_fn = std::function<void(size_t, size_t)>;

template <class Type = Raw_data>
using Variable = std::tuple<Type, uint32_t>;

std::optional<tstring> init();

template <class Type = Raw_data>
std::optional<std::vector<Type>> deserialize_list(const void *data, size_t data_size);

template <class Type = Raw_data, class Size_fn, class Advance_fn>
std::optional<std::vector<Type>> deserialize_list_ex(const void *data, size_t data_size, const Size_fn &get_element_size, const Advance_fn &get_next_element);

template <class Type = Raw_data>
size_t serialize(Raw_data &output, const Type &value);

template <class Type = Raw_data>
size_t serialize_list(Raw_data &output, const std::vector<Type> &value);

template <class Filter_fn>
std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables(const Filter_fn &filter, const Progress_fn &progress);

template <class Filter_fn>
std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables(const Filter_fn &filter);
std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables();

template <class Type = Raw_data>
std::optional<Variable<Type>> get_variable(const efi_guid_t &guid, const tstring &name);

template <class Type = Raw_data>
std::optional<Variable<std::vector<Type>>> get_list_variable(const efi_guid_t &guid, const tstring &name);

template <class Type = Raw_data, class Size_fn, class Advance_fn>
std::optional<Variable<std::vector<Type>>> get_list_variable_ex(const efi_guid_t &guid, const tstring &name, const Size_fn &get_element_size, const Advance_fn &get_next_element);

template <class Type = Raw_data>
bool set_variable(const efi_guid_t &guid, const tstring &name, const Variable<Type> &variable, mode_t mode);

template <class Type = Raw_data>
bool set_list_variable(const efi_guid_t &guid, const tstring &name, const Variable<std::vector<Type>> &variable, mode_t mode);

inline std::optional<tstring> init()
{
    if(!efi_variables_supported())
        return {_T("UEFI variables not supported on this machine.")};

    return std::nullopt;
}

template <class Type>
inline std::optional<Type> deserialize(const void *data, size_t data_size)
{
    if(data_size != sizeof(Type))
        return std::nullopt;

    return {*static_cast<const Type *>(data)};
}

template <class Type>
inline size_t serialize(Raw_data &output, const Type &value)
{
    size_t pos = output.size();
    output.resize(pos + sizeof(value));
    memcpy(&output[pos], &value, sizeof(value));
    return sizeof(value);
}

template <>
inline std::optional<Raw_data> deserialize<Raw_data>(const void *data, size_t data_size)
{
    return {Raw_data{static_cast<const uint8_t *>(data), static_cast<const uint8_t *>(advance_bytes(data, data_size))}};
}

template <>
inline size_t serialize(Raw_data &output, const Raw_data &data)
{
    output.insert(std::end(output), std::begin(data), std::end(data));
    return data.size();
}

template <>
inline std::optional<std::string> deserialize<std::string>(const void *data, size_t data_size)
{
    if(data_size % sizeof(std::string::value_type))
        return std::nullopt;

    return {std::string{static_cast<const std::string::value_type *>(data), data_size / sizeof(std::string::value_type)}};
}

template <>
inline size_t serialize(Raw_data &output, const std::string &value)
{
    size_t bytes = (value.size() + 1) * sizeof(std::string::value_type);
    size_t pos = output.size();
    output.resize(pos + bytes);
    memcpy(&output[pos], value.c_str(), bytes);
    return bytes;
}

template <>
inline std::optional<std::wstring> deserialize<std::wstring>(const void *data, size_t data_size)
{
    if(data_size % sizeof(std::wstring::value_type))
        return std::nullopt;

    return {std::wstring{static_cast<const std::wstring::value_type *>(data), data_size / sizeof(std::wstring::value_type)}};
}

template <>
inline size_t serialize(Raw_data &output, const std::wstring &value)
{
    size_t bytes = (value.size() + 1) * sizeof(std::wstring::value_type);
    size_t pos = output.size();
    output.resize(pos + bytes);
    memcpy(&output[pos], value.c_str(), bytes);
    return bytes;
}

template <>
inline std::optional<std::u16string> deserialize(const void *data, size_t data_size)
{
    if(data_size % sizeof(std::u16string::value_type))
        return std::nullopt;

    return {std::u16string{static_cast<const std::u16string::value_type *>(data), data_size / sizeof(std::u16string::value_type)}};
}

template <>
inline size_t serialize(Raw_data &output, const std::u16string &value)
{
    size_t bytes = (value.size() + 1) * sizeof(std::u16string::value_type);
    size_t pos = output.size();
    output.resize(pos + bytes);
    memcpy(&output[pos], value.c_str(), bytes);
    return bytes;
}

template <class Type, class Size_fn, class Advance_fn>
inline std::optional<std::vector<Type>> deserialize_list_ex(const void *data, size_t data_size, const Size_fn &get_element_size, const Advance_fn &get_next_element)
{
    std::vector<Type> values;
    const void *data_end = advance_bytes(data, data_size);
    while(data && data < data_end)
    {
        auto element_size = get_element_size(data);
        auto value = deserialize<Type>(data, element_size);
        if(!value)
            return std::nullopt;

        values.push_back(*value);
        data = get_next_element(data, data_size);
        auto bytes_left = static_cast<const uint8_t *>(data_end) - static_cast<const uint8_t *>(data);
        data_size = static_cast<size_t>(bytes_left);
    }

    if(data != data_end)
        return std::nullopt;

    return {values};
}

template <class Type>
inline std::optional<std::vector<Type>> deserialize_list(const void *data, size_t data_size)
{
    return deserialize_list_ex<Type>(
        data, data_size,
        [](const void *)
        {
            return sizeof(Type);
        },
        [](const void *ptr, size_t)
        {
            return advance_bytes(ptr, sizeof(const Type));
        });
}

template <class Type>
size_t serialize_list(Raw_data &output, const std::vector<Type> &value)
{
    size_t bytes = 0;
    for(const auto &item: value)
        bytes += serialize(output, item);

    return bytes;
}

{% for category in device_paths.values() %}
// {{ category.name }}

{% for node in category.nodes %}
template <>
inline std::optional<File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}> deserialize(const void *data, size_t data_size)
{
    auto dp = static_cast<const efidp_{{ category.slug }}_{{ node.slug }} *>(data);
    if(dp->header.length != data_size)
        return std::nullopt;

    if(dp->header.type != File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}::TYPE)
        return std::nullopt;

    if(dp->header.subtype != File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}::SUBTYPE)
        return std::nullopt;

    File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }} value{};
    {% for field in node.fields %}
    {% if field.type == "enum" %}
    value.{{ field.slug }} = static_cast<File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}::{{ field.slug.upper() }}>(dp->{{ field.slug }});
    {% elif "string" in field.type %}
    value.{{ field.slug }} = reinterpret_cast<const decltype(value.{{ field.slug }})::value_type *>(dp->{{ field.slug }});
    {% elif field.size == "n" %}
    {
        static_assert(sizeof(decltype(value.{{ field.slug }})::value_type) == sizeof(dp->{{ field.slug }}[0]));
        auto length = static_cast<size_t>(static_cast<decltype(&dp->{{ field.slug }}[0])>(advance_bytes(data, data_size)) - &dp->{{ field.slug }}[0]);
        value.{{ field.slug }}.resize(length);
        memcpy(value.{{ field.slug }}.data(), dp->{{ field.slug }}, length * sizeof(dp->{{ field.slug }}[0]));
    }
    {% elif field.size > 8 or field.type in ("ip4", "ip6", "mac") %}
    std::copy(std::begin(dp->{{ field.slug }}), std::end(dp->{{ field.slug }}), std::begin(value.{{ field.slug }}));
    {% else %}
    value.{{ field.slug }} = dp->{{ field.slug }};
    {% endif %}
    {% endfor %}
    return {value};
}

template <>
inline size_t serialize(Raw_data &output, const File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }} &{{ node.slug }})
{
    size_t bytes = 0;
    uint8_t type = File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}::TYPE;
    bytes += serialize(output, type);
    uint8_t subtype = File_path::{{ category.slug.upper() }}::{{ node.slug.capitalize() }}::SUBTYPE;
    bytes += serialize(output, subtype);
    size_t pos = output.size();
    uint16_t length = 0;
    bytes += serialize(output, length);
    {% for field in node.fields %}
    bytes += serialize(output, {{ node.slug }}.{{ field.slug }});
    {% endfor %}
    length = static_cast<uint16_t>(bytes);
    memcpy(&output[pos], &length, sizeof(length));
    return bytes;
}

{% endfor %}{% endfor %}
template <>
inline std::optional<File_path::Unknown> deserialize(const void *data, size_t data_size)
{
    auto dp = static_cast<const efidp_header *>(data);
    if(dp->length != data_size)
        return std::nullopt;

    File_path::Unknown value{};
    value.TYPE = dp->type;
    value.SUBTYPE = dp->subtype;

    size_t data_length = data_size - sizeof(*dp);
    value.data.resize(data_length);
    memcpy(value.data.data(), static_cast<const uint8_t *>(advance_bytes(data, sizeof(const efidp_header))), data_length);
    return {value};
}

template <>
inline size_t serialize(Raw_data &output, const File_path::Unknown &unknown)
{
    size_t bytes = 0;
    bytes += serialize(output, unknown.TYPE);
    bytes += serialize(output, unknown.SUBTYPE);
    size_t pos = output.size();
    uint16_t length = 0;
    bytes += serialize(output, length);
    bytes += serialize(output, unknown.data);
    length = static_cast<uint16_t>(bytes);
    memcpy(&output[pos], &length, sizeof(length));
    return bytes;
}

template <>
inline std::optional<File_path::ANY> deserialize(const void *data, size_t data_size)
{
    auto dp = static_cast<const_efidp>(data);
    if(dp->header.length != data_size)
        return std::nullopt;

#define TYPE_SUBTYPE(type, subtype) (((type) << 8) | (subtype))
#define casefp(Type)                                                    \
    case TYPE_SUBTYPE(File_path::Type::TYPE, File_path::Type::SUBTYPE): \
        return deserialize<File_path::Type>(dp, data_size)

    switch(TYPE_SUBTYPE(dp->header.type, dp->header.subtype))
    {
{% for category in device_paths.values() %}{% for node in category.nodes %}
        casefp({{ category.slug.upper() }}::{{ node.slug.capitalize() }});
{% endfor %}{% endfor %}

    default:
        return deserialize<File_path::Unknown>(dp, data_size);
    }

#undef casefp
#undef TYPE_SUBTYPE
}

template <>
inline size_t serialize(Raw_data &output, const File_path::ANY &file_path)
{
    return std::visit([&output](const auto &dp) -> size_t
        { return serialize(output, dp); },
        file_path);
}

template <>
inline std::optional<Load_option> deserialize(const void *data, size_t data_size)
{
    Load_option value{};
    auto ssize = static_cast<ssize_t>(data_size);
    auto load_option = const_cast<efi_load_option *>(static_cast<const efi_load_option *>(data));

#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunsafe-buffer-usage"
#endif
    for(size_t d = 0; load_option->description[d]; ++d)
        value.description.push_back(load_option->description[d]);
#if defined(__clang__)
#pragma clang diagnostic pop
#endif

    uint16_t device_path_size = efi_loadopt_pathlen(load_option, ssize);
    const_efidp device_path = efi_loadopt_path(load_option, ssize);

    auto file_paths = deserialize_list_ex<File_path::ANY>(
        device_path, device_path_size,
        [](const void *ptr)
        {
            auto size = efidp_node_size(static_cast<const_efidp>(ptr));
            return static_cast<size_t>(size);
        },
        [](const void *ptr, size_t bytes_left) -> const void *
        {
            auto dp = static_cast<const_efidp>(ptr);
            ssize_t size = efidp_node_size(dp);
            if(size < 0 || static_cast<size_t>(size) > bytes_left)
                return nullptr;

            return advance_bytes(ptr, static_cast<size_t>(size));
        });

    if(!file_paths || file_paths->empty())
        return std::nullopt;

    value.device_path = *file_paths;

    uint8_t *optional_data = nullptr;
    size_t optional_data_size = 0;
    if(int ret = efi_loadopt_optional_data(load_option, data_size, &optional_data, &optional_data_size); ret >= 0)
    {
        auto opt_data = deserialize<Raw_data>(optional_data, optional_data_size);
        if(!opt_data)
            return std::nullopt;

        value.optional_data = *opt_data;
    }

    value.attributes = static_cast<Load_option_attribute>(load_option->attributes);
    return {value};
}

template <>
inline size_t serialize(Raw_data &output, const Load_option &load_option)
{
    size_t size = 0;
    size += serialize(output, load_option.attributes);
    auto file_path_list_length_pos = output.size();
    uint16_t file_path_list_size = 0;
    size += serialize(output, file_path_list_size);
    size += serialize(output, load_option.description);

    file_path_list_size = static_cast<uint16_t>(serialize_list(output, load_option.device_path));
    // Always set END_ENTIRE tag at the end of device path
    if(!load_option.device_path.size() || std::visit([](const auto &file_path)
                                              { return file_path.SUBTYPE; },
                                              load_option.device_path.back())
            != File_path::END::Entire::SUBTYPE)
    {
        File_path::END::Entire end{};
        file_path_list_size = static_cast<uint16_t>(file_path_list_size + static_cast<uint16_t>(serialize(output, end))); // Older GCC complains about conversion when using += `conversion from ‘int’ to ‘uint16_t’ {aka ‘short unsigned int’} may change value`
    }
    size += file_path_list_size;
    memcpy(&output[file_path_list_length_pos], &file_path_list_size, sizeof(file_path_list_size));

    size += serialize(output, load_option.optional_data);
    return size;
}

extern Progress_fn _get_variables_progress_fn;

template <class Filter_fn>
inline std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables(const Filter_fn &filter_fn, const Progress_fn &progress_fn)
{
    std::unordered_map<tstring, efi_guid_t> variables;
    efi_guid_t *guid = nullptr;
    TCHAR *name = nullptr;
    _get_variables_progress_fn = progress_fn;
    efi_set_get_next_variable_name_progress_cb([](size_t step, size_t total) noexcept
        { try { if(_get_variables_progress_fn)_get_variables_progress_fn(step, total); } catch (...) {/* ignore */} });

    int ret = 0;
    while((ret = efi_get_next_variable_name(&guid, &name)) > 0)
    {
        if(!filter_fn(*guid, name))
            continue;

        memcpy(&variables[name], guid, sizeof(efi_guid_t));
    }

    efi_set_get_next_variable_name_progress_cb(nullptr);
    _get_variables_progress_fn = nullptr;

    if(ret < 0)
        return std::nullopt;

    return {variables};
}

template <class Filter_fn>
inline std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables(Filter_fn filter_fn)
{
    return get_variables(filter_fn, [](size_t, size_t) { /* noprogress */ });
}

inline std::optional<std::unordered_map<tstring, efi_guid_t>> get_variables()
{
    return get_variables(
        [](const efi_guid_t &, const tstring_view)
        { return true; },
        [](size_t, size_t) { /* noprogress */ });
}

template <class Type>
inline std::optional<Variable<Type>> get_variable(const efi_guid_t &guid, const tstring &name)
{
    uint8_t *data = nullptr;
    size_t data_size = 0;
    uint32_t attributes = 0;
    if(int ret = efi_get_variable(guid, name.c_str(), &data, &data_size, &attributes); ret < 0)
        return std::nullopt;

    auto value = deserialize<Type>(data, data_size);
    if(!value)
        return std::nullopt;

    return {{ "{{*value, attributes}}" }};
}

template <class Type>
inline std::optional<Variable<std::vector<Type>>> get_list_variable(const efi_guid_t &guid, const tstring &name)
{
    uint8_t *data = nullptr;
    size_t data_size = 0;
    uint32_t attributes = 0;
    if(int ret = efi_get_variable(guid, name.c_str(), &data, &data_size, &attributes); ret < 0)
        return std::nullopt;

    auto value = deserialize_list<Type>(data, data_size);
    if(!value)
        return std::nullopt;

    return {{ "{{*value, attributes}}" }};
}

template <class Type, class Size_fn, class Advance_fn>
inline std::optional<Variable<std::vector<Type>>> get_list_variable_ex(const efi_guid_t &guid, const tstring &name, const Size_fn &get_element_size, const Advance_fn &get_next_element)
{
    uint8_t *data = nullptr;
    size_t data_size = 0;
    uint32_t attributes = 0;
    if(int ret = efi_get_variable(guid, name.c_str(), &data, &data_size, &attributes); ret < 0)
        return std::nullopt;

    auto value = deserialize_list_ex<Type>(data, data_size, get_element_size, get_next_element);
    if(!value)
        return std::nullopt;

    return {{ "{{*value, attributes}}" }};
}

template <class Type>
inline bool set_variable(const efi_guid_t &guid, const tstring &name, const Variable<Type> &variable, mode_t mode)
{
    auto [value, attributes] = variable;
    Raw_data bytes;
    size_t size = serialize(bytes, value);
    // Skip overwriting with exactly the same value
    if(const auto current = EFIBoot::get_variable<Raw_data>(guid, name); current)
    {
        const auto &[current_bytes, current_attributes] = *current;
        if(current_attributes == attributes && current_bytes == bytes)
            return true;
    }

    // Don't care about the error from get_variable
    efi_error_clear();
    return efi_set_variable(guid, name.c_str(), bytes.data(), size, attributes, mode) == 0;
}

template <class Type>
inline bool set_list_variable(const efi_guid_t &guid, const tstring &name, const Variable<std::vector<Type>> &variable, mode_t mode)
{
    auto [value, attributes] = variable;
    Raw_data bytes;
    size_t size = serialize_list(bytes, value);
    // Skip overwriting with exactly the same value
    if(const auto current = EFIBoot::get_variable<Raw_data>(guid, name); current)
    {
        const auto &[current_bytes, current_attributes] = *current;
        if(current_attributes == attributes && current_bytes == bytes)
            return true;
    }

    // Don't care about the error from get_variable
    efi_error_clear();
    return efi_set_variable(guid, name.c_str(), bytes.data(), size, attributes, mode) == 0;
}

inline bool del_variable(const efi_guid_t &guid, const tstring &name)
{
    return efi_del_variable(guid, name.c_str()) == 0;
}

inline tstring get_error_trace()
{
    tstring output = _T("Error trace:\n");
    unsigned int i = 0;
    while(true)
    {
        TCHAR *filename = nullptr;
        TCHAR *function = nullptr;
        int line = 0;
        TCHAR *message = nullptr;
        int error = 0;
        const int ERROR_STR_BUFFER_SIZE = 1024;
        TCHAR error_str[ERROR_STR_BUFFER_SIZE] = {};

        int rc = efi_error_get(i, &filename, &function, &line, &message, &error);
        if(rc < 0)
            output += _T("error fetching trace value\n");

        if(rc <= 0)
            break;

        ++i;
        if(_tcserror_s(error_str, ERROR_STR_BUFFER_SIZE - 1, error) != 0)
            output += _T("error translating error code to string\n");

        output += filename;
        output += _T(":");
        output += to_tstring(line);
        output += _T(" ");
        output += function;
        output += _T("(): \n    ");
        output += error_str;
        output += _T("[");
        output += to_tstring(error);
        output += _T("]: ");
        output += message;
        output += _T("\n");
    }

    if(i == 0)
        output += _T("no errors?");

    return output;
}

inline void error_clear()
{
    efi_error_clear();
}

} // namespace EFIBoot
