#include "memory_manager.h"

ui32 TTrigramDocsMap::RemoveDoc(ui32 docId) {
    TWriteGuard wg(Mutex);
    DEBUG_LOG << "Removing " << docId << " from rt manager" << Endl;
    auto it = HashesByDoc.find(docId);
    if (it == HashesByDoc.end()) {
        DEBUG_LOG << "Removing " << docId << " from rt manager -> 0" << Endl;
        return 0;
    }

    SUPERLONG firstPos = 0;
    TWordPosition::SetDoc(firstPos, docId);

    for (auto&& i : it->second) {
        auto itF = PositionsByHash.find(i);
        CHECK_WITH_LOG(itF != PositionsByHash.end());
        auto firstIter = itF->second.lower_bound(firstPos);
        auto lastIter = firstIter;
        while (lastIter != itF->second.end() && TWordPosition::Doc(*lastIter) == docId)
            ++lastIter;
        itF->second.erase(firstIter, lastIter);
        if (itF->second.empty())
            PositionsByHash.erase(itF);
    }
    HashesByDoc.erase(it);
    DEBUG_LOG << "Removing " << docId << " from rt manager -> OK" << Endl;
    return true;
}

bool TTrigramDocsMap::AddData(const TString& value, const TSet<SUPERLONG>& positions) {
    TWriteGuard wg(Mutex);
    for (const auto& pos : positions)
        HashesByDoc[TWordPosition::Doc(pos)].insert(value);
    PositionsByHash[value].insert(positions.begin(), positions.end());
    return true;
}

bool TTrigramDocsMap::GetData(const TString& value, TSet<SUPERLONG>& positions) const {
    TReadGuard rg(Mutex);
    auto it = PositionsByHash.find(value);
    if (it == PositionsByHash.end()) {
        return false;
    }
    positions = it->second;
    return true;
}

bool TTrigramDocsMap::Empty() const {
    TReadGuard rg(Mutex);
    return PositionsByHash.empty();
}

ui32 TTrigramIndexManagerMemory::RemoveDocids(const TVector<ui32>& docids) {
    TReadGuard rg(Mutex);
    ui32 count = 0;
    TSet<TString> toDelete;
    for (auto&& docid : docids) {
        for (auto&& i : Map) {
            if (ui32 removed = i.second->RemoveDoc(docid)) {
                count += removed;
                if (i.second->Empty())
                    toDelete.insert(i.first);
            }
        }
    }
    rg.Release();
    if (!toDelete.empty()) {
        TWriteGuard wg(Mutex);
        for (const auto& key : toDelete) {
            auto i = Map.find(key);
            if (i != Map.end() && i->second->Empty()) {
                Map.erase(i);
            }
        }
    }
    return count;
}

bool TTrigramIndexManagerMemory::AddData(const TString& keyName, const TString& value, const TSet<SUPERLONG>& positions) {
    TReadGuard rg(Mutex);
    auto it = Map.find(keyName);
    if (it == Map.end()) {
        rg.Release();
        TWriteGuard wg(Mutex);
        if (!Map.contains(keyName)) {
            Map[keyName] = new TTrigramDocsMap();
        }
        it = Map.find(keyName);
        return it->second->AddData(value, positions);
    } else {
        return it->second->AddData(value, positions);
    }
}

bool TTrigramIndexManagerMemory::GetData(const TString& keyName, const TString& value, const IDocProcessor* filter, TSet<SUPERLONG>& result) const {
    TReadGuard rg(Mutex);
    auto it = Map.find(keyName);
    if (it == Map.end()) {
        return false;
    } else {
        TSet<SUPERLONG> positions;
        it->second->GetData(value, positions);
        for (auto&& pos : positions) {
            if (!filter || !filter->DocumentBanned(TWordPosition::Doc(pos))) {
                result.insert(pos);
            }
        }
    }
    return true;
}

class TTrigramIndexManagerMemory::TMemoryIterator: public TTrigramIndexManager::TIterator {
public:
    typedef TAtomicSharedPtr<TSet<SUPERLONG>> TPosPtr;
    TMemoryIterator(TPosPtr positions, const IDocProcessor* filter)
        : Filter(filter)
        , Positions(positions)
        , PosIterator(Positions->cbegin())
    {
        FindFiltered();
    }

    bool Valid() const override {
        return PosIterator != Positions->cend();
    }

    void Next() override {
        ++PosIterator;
        FindFiltered();
    }

    SUPERLONG GetPosition() const override {
        return *PosIterator;
    }

private:
    inline void FindFiltered() {
        while (Filter && (PosIterator != Positions->cend()) && Filter->DocumentBanned(TWordPosition::Doc(*PosIterator))) {
            ++PosIterator;
        }
    }
    const IDocProcessor* Filter;
    TAtomicSharedPtr<TSet<SUPERLONG>> Positions;
    TSet<SUPERLONG>::const_iterator PosIterator;
};

TTrigramIndexManager::TIterator::TPtr TTrigramIndexManagerMemory::CreateIterator(const TString& keyName, const TString& value, const IDocProcessor* filter) const {
    TReadGuard rg(Mutex);
    auto it = Map.find(keyName);
    if (it == Map.end()) {
        return nullptr;
    }

    TMemoryIterator::TPosPtr positions(MakeAtomicShared<TSet<SUPERLONG>>());
    it->second->GetData(value, *positions);
    if (positions->empty()) {
        return nullptr;
    }
    return MakeIntrusive<TMemoryIterator>(positions, filter);
}
