#pragma once

#include "defs.h"

namespace NActors {
    // wait-free one writer multi reader hash-tree for service mapping purposes
    // on fast updates on same key - could lead to false-negatives, we don't care as such cases are broken from service-map app logic

    template <typename TKey, typename TValue, typename THash, ui64 BaseSize = 256 * 1024, ui64 ExtCount = 4, ui64 ExtBranching = 4>
    class TServiceMap : TNonCopyable {
        struct TEntry : TNonCopyable {
            ui32 CounterIn;
            ui32 CounterOut;
            TKey Key;
            TValue Value;

            TEntry()
                : CounterIn(0)
                , CounterOut(0)
                , Key()
                , Value()
            {
            }
        };

        struct TBranch : TNonCopyable {
            TEntry Entries[ExtCount];
            TBranch* Branches[ExtBranching];

            TBranch() {
                Fill(Branches, Branches + ExtBranching, (TBranch*)nullptr);
            }
        };

        ui32 Counter;
        TBranch* Line[BaseSize];

        bool ScanBranch(TBranch* branch, const TKey& key, ui64 hash, TValue& ret) {
            for (ui32 i = 0; i != ExtCount; ++i) {
                const TEntry& entry = branch->Entries[i];
                const ui32 counterIn = AtomicLoad(&entry.CounterIn);
                if (counterIn != 0 && entry.Key == key) {
                    ret = entry.Value;
                    const ui32 counterOut = AtomicLoad(&entry.CounterOut);
                    if (counterOut == counterIn)
                        return true;
                }
            }

            const ui64 hash0 = hash % ExtBranching;
            if (TBranch* next = AtomicLoad(branch->Branches + hash0))
                return ScanBranch(next, key, hash / ExtBranching, ret);

            return false;
        }

        void ScanZeroOld(TBranch* branch, const TKey& key, ui64 hash, TEntry** zeroEntry, TEntry*& oldEntry) {
            for (ui32 i = 0; i != ExtCount; ++i) {
                TEntry& entry = branch->Entries[i];
                if (entry.CounterIn == 0) {
                    if (zeroEntry && !*zeroEntry) {
                        *zeroEntry = &entry;
                        if (oldEntry != nullptr)
                            return;
                    }
                } else {
                    if (entry.Key == key) {
                        oldEntry = &entry;
                        if (!zeroEntry || *zeroEntry)
                            return;
                    }
                }
            }

            const ui64 hash0 = hash % ExtBranching;
            if (TBranch* next = branch->Branches[hash0]) {
                ScanZeroOld(next, key, hash / ExtBranching, zeroEntry, oldEntry);
            } else { // found tail, if zeroEntry requested, but not yet found - insert one
                if (zeroEntry && !*zeroEntry) {
                    TBranch* next = new TBranch();
                    *zeroEntry = next->Entries;
                    AtomicStore(branch->Branches + hash0, next);
                }
            }
        }

    public:
        TServiceMap()
            : Counter(0)
        {
            Fill(Line, Line + BaseSize, (TBranch*)nullptr);
        }

        ~TServiceMap() {
            for (ui64 i = 0; i < BaseSize; ++i) {
                delete Line[i];
            }
        }

        TValue Find(const TKey& key) {
            THash hashOp;
            const ui64 hash = hashOp(key);
            const ui64 hash0 = hash % BaseSize;

            if (TBranch* branch = AtomicLoad(Line + hash0)) {
                TValue ret;
                if (ScanBranch(branch, key, hash / BaseSize, ret))
                    return ret;
            }

            return TValue();
        }

        // returns true on update, false on insert
        TValue Update(const TKey& key, const TValue& value) {
            THash hashOp;
            const ui64 hash = hashOp(key);
            const ui64 hash0 = hash % BaseSize;

            TEntry* zeroEntry = nullptr;
            TEntry* oldEntry = nullptr;

            if (TBranch* branch = Line[hash0]) {
                ScanZeroOld(branch, key, hash / BaseSize, &zeroEntry, oldEntry);
            } else {
                TBranch* next = new TBranch();
                zeroEntry = next->Entries;
                AtomicStore(Line + hash0, next);
            }

            // now we got both entries, first - push new one
            const ui32 counter = AtomicUi32Increment(&Counter);
            AtomicStore(&zeroEntry->CounterOut, counter);
            zeroEntry->Key = key;
            zeroEntry->Value = value;
            AtomicStore(&zeroEntry->CounterIn, counter);

            if (oldEntry != nullptr) {
                const TValue ret = oldEntry->Value;
                AtomicStore<ui32>(&oldEntry->CounterOut, 0);
                AtomicStore<ui32>(&oldEntry->CounterIn, 0);
                return ret;
            } else {
                return TValue();
            }
        }

        bool Erase(const TKey& key) {
            THash hashOp;
            const ui64 hash = hashOp(key);
            const ui64 hash0 = hash % BaseSize;

            TEntry* oldEntry = 0;

            if (TBranch* branch = Line[hash0]) {
                ScanZeroOld(branch, key, hash / BaseSize, 0, oldEntry);
            }

            if (oldEntry != 0) {
                AtomicStore<ui32>(&oldEntry->CounterOut, 0);
                AtomicStore<ui32>(&oldEntry->CounterIn, 0);
                return true;
            } else {
                return false;
            }
        }
    };
}