#pragma once

#include <atomic>
#include <cstddef>
#include <memory>
#include <type_traits>
#include <typeinfo>

#include <userver/utils/fast_pimpl.hpp>

USERVER_NAMESPACE_BEGIN

namespace engine::impl::task_local {

using Key = std::size_t;

enum class VariableKind { kNormal, kInherited };

class DataBase {
public:
    using Deleter = void (*)(DataBase&) noexcept;

    void DeleteSelf() noexcept;

protected:
    explicit DataBase(Deleter deleter);

private:
    const Deleter deleter_;
};

class NormalDataBase : public DataBase {
protected:
    template <typename Derived>
    NormalDataBase(std::in_place_type_t<Derived>, Key /*key*/)
        : DataBase([](DataBase& base) noexcept { delete &static_cast<Derived&>(base); }) {}
};

class InheritedDataBase : public DataBase {
public:
    void AddRef() noexcept;

    Key GetKey() const noexcept;

protected:
    template <typename Derived>
    InheritedDataBase(std::in_place_type_t<Derived>, Key key)
        : DataBase([](DataBase& base) noexcept {
              // NOLINTNEXTLINE(cppcoreguidelines-pro-type-static-cast-downcast)
              if (--static_cast<InheritedDataBase&>(base).ref_counter_ == 0) {
                  delete &static_cast<Derived&>(base);
              }
          }),
          key_(key) {}

private:
    std::atomic<std::size_t> ref_counter_{1};
    const Key key_{};
};

template <VariableKind Kind>
using ConditionalDataBase = std::conditional_t<Kind == VariableKind::kInherited, InheritedDataBase, NormalDataBase>;

template <typename T, VariableKind Kind>
class DataImpl final : public ConditionalDataBase<Kind> {
    using Base = ConditionalDataBase<Kind>;

public:
    template <typename... Args>
    explicit DataImpl(Key key, Args&&... args)
        : Base(std::in_place_type<DataImpl>, key), variable_(std::forward<Args>(args)...) {}

    T& Get() noexcept { return variable_; }

private:
    T variable_;
};

[[noreturn]] void ReportVariableNotSet(const std::type_info& type);

class Storage final {
public:
    Storage();

    Storage(Storage&&) = delete;
    Storage& operator=(Storage&&) = delete;
    ~Storage();

    // Copies pointers to inherited variables from 'other'
    // 'this' must not contain any variables
    void InheritFrom(Storage& other);

    // Copies pointers to specific inherited variables from 'other'
    // 'this' must not contain the variable being copied
    // does nothing if there is nothing to copy
    void InheritNodeIfExists(Storage& other, Key key);

    // Moves other's variables into 'this', leaving 'other' in an empty state
    // 'this' must not contain any variables
    void InitializeFrom(Storage&& other) noexcept;

    // Variable accessors must be called with the same T, Kind, key.
    // Otherwise it is UB.
    template <typename T, VariableKind Kind>
    T& GetOrEmplace(Key key) {
        DataBase* const old_data = GetGeneric(key);
        if (!old_data) {
            const bool has_existing_variable = false;
            return DoEmplace<T, Kind>(key, has_existing_variable);
        }
        return static_cast<DataImpl<T, Kind>&>(*old_data).Get();
    }

    template <typename T, VariableKind Kind>
    T* GetOptional(Key key) noexcept {
        DataBase* const data = GetGeneric(key);
        if (!data) return nullptr;
        return &static_cast<DataImpl<T, Kind>&>(*data).Get();
    }

    template <typename T, VariableKind Kind>
    T& Get(Key key) {
        T* const result = GetOptional<T, Kind>(key);
        if (!result) ReportVariableNotSet(typeid(T));
        return *result;
    }

    template <typename T, VariableKind Kind, typename... Args>
    T& Emplace(Key key, Args&&... args) {
        DataBase* const old_data = GetGeneric(key);
        const bool has_existing_variable = old_data != nullptr;
        auto& result = DoEmplace<T, Kind>(key, has_existing_variable, std::forward<Args>(args)...);
        if (old_data) old_data->DeleteSelf();
        return result;
    }

    template <typename T, VariableKind Kind>
    void Erase(Key key) noexcept {
        static_assert(Kind == VariableKind::kInherited);
        EraseInherited(key);
    }

private:
    DataBase* GetGeneric(Key key) noexcept;

    void SetGeneric(Key key, NormalDataBase& node, bool has_existing_variable);

    void SetGeneric(Key key, InheritedDataBase& node, bool has_existing_variable);

    void EraseInherited(Key key) noexcept;

    void InheritNode(InheritedDataBase&);

    // Provides strong exception guarantee. Does not delete the old data, if any.
    template <typename T, VariableKind Kind, typename... Args>
    T& DoEmplace(Key key, bool has_existing_variable, Args&&... args) {
        auto new_data = std::make_unique<DataImpl<T, Kind>>(key, std::forward<Args>(args)...);
        SetGeneric(key, *new_data, has_existing_variable);
        return new_data.release()->Get();
    }

    struct Impl;
    utils::FastPimpl<Impl, 40, 8> impl_;
};

class Variable final {
public:
    Variable();

    Variable(const Variable&) = delete;
    Variable(Variable&&) = delete;
    Variable& operator=(const Variable&) = delete;
    Variable& operator=(Variable&&) = delete;

    Key GetKey() const noexcept;

private:
    const Key key_;
};

Storage& GetCurrentStorage() noexcept;

struct InternalTag final {
    explicit InternalTag() = default;
};

}  // namespace engine::impl::task_local

USERVER_NAMESPACE_END
