#include "indexer_manager.h"
#include "common_action.h"
#include "indexer.h"

#include <saas/util/queue.h>
#include <library/cpp/balloc/optional/operators.h>
#include <saas/rtyserver/indexer_memory/memory_indexer.h>
#include <saas/rtyserver/indexer_core/messages.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/config/realm_config.h>
#include <saas/rtyserver/config/shards_config.h>

class TReopenTask: public IObjectInQueue {
private:
    TIndexerManager& Manager;
    TQueuedDocument Document;
    const TRTYServerConfig& Config;
    TString RealmName;
public:
    TReopenTask(TIndexerManager& manager, TQueuedDocument document, const TRTYServerConfig& config, const TString& realmName)
        : Manager(manager)
        , Document(document)
        , Config(config)
        , RealmName(realmName)
    {
    }

    void Process(void* /*ThreadSpecificResource*/) override {
        ThreadDisableBalloc();

        THolder<TReopenTask> This(this);
        bool failed = false;
        for (ui32 i = 0; i < Config.GetShardsConfig().Number; i++) {
            if (Manager.GetIsActive())
                Manager.ResetIndexers(true, i, true, false, false, RealmName);
            else
                failed = true;
        }
        if (!failed) {
            Manager.WaitCloseTasks();
        }
        while (Manager.IsExistsClosing() && Manager.GetIsActive()) {
            Sleep(TDuration::Seconds(1));
        }
        if (Document) {
            const TString status = (!Manager.GetIsActive() || failed) ? "Server stopped" : "";
            Document->SetStatus(NRTYServer::TReply::OK, status);
        }
    }

};

class TCloseIndexTask: public IObjectInQueue, public IMessageProcessor {
private:
    TAtomicSharedPtr<IIndexer> Indexer;
    THolder<TGuardIncompatibleAction> Guard;
    const std::atomic<bool>* RigidStopSignal;
public:
    TCloseIndexTask(TAtomicSharedPtr<IIndexer> indexer, ITransaction* transaction, const std::atomic<bool>* rigidStopSignal) {
        if (transaction)
            Guard.Reset(new TGuardIncompatibleAction(*transaction));
        Indexer = indexer;
        RigidStopSignal = rigidStopSignal;
        RegisterGlobalMessageProcessor(this);
    }

    ~TCloseIndexTask() {
        UnregisterGlobalMessageProcessor(this);
    }

    bool Process(IMessage* message) override {
        TMessageCollectServerInfo* messCollect = dynamic_cast<TMessageCollectServerInfo*>(message);
        if (messCollect) {
            messCollect->AddClosingIndexers();
            return true;
        }
        return false;
    }

    TString Name() const override {
        return "TCloseIndexTask";
    }

    void Process(void* /*ThreadSpecificResource*/) override {
        ThreadDisableBalloc();
        if (!!Indexer) {
            NOTICE_LOG << "Index " << Indexer->Directory().BaseName() << " closing..." << Endl;
            Indexer->CloseIndex(RigidStopSignal, NRTYServer::EExecutionContext::Close);
            NOTICE_LOG << "Index " << Indexer->Directory().BaseName() << " closing...OK" << Endl;
        }
    }
};

class TIndexersThreadDisk: public TIndexersThread {
public:
    TIndexerManager& Manager;
public:
    TIndexersThreadDisk(TAtomicSharedPtr<TIndexersList> indexersList, TSynchronizedDocumentsQueue& queue, TDocumentModifier& deleter,
        const TRTYServerConfig& config, const TString& serviceName, int threadID, TIndexerManager& manager
    )
        : TIndexersThread(indexersList, queue, deleter, config, serviceName, threadID)
        , Manager(manager)
    {
        IndexerMetrics = Manager.GetIndexerMetrics();
    }

    void Process(void* /*ThreadSpecificResource*/) override {
        ThreadDisableBalloc();

        THolder<TIndexersThreadDisk> suicide(this);
        while (IsActive()) {
            TQueuedDocument document(DocumentsQueue->Get());
            if (!!document) {
                IIndexationAction::TPtr action = IIndexationAction::TFactory::Construct(document->GetCommand());
                VERIFY_WITH_LOG(!!action, "Incorrect message type for indexation: %s", NRTYServer::TMessage_TMessageType_Name(document->GetCommand()).data());
                action->Execute(document, *this);
                ui32 shard = action->GetShard();
                action.Reset(nullptr);

                Manager.ResetIndexers(false, shard, true, false, true, "");
            } else {
                const NRTYServer::TShardsConfig &shards = Config.GetShardsConfig();
                for (ui32 shard : shards.LocalShards) {
                    TWriteShardIndexGuard gShard(*IndexersList.Get(), shard);
                    Manager.ResetIndexers(false, shard, false, true, false, "");
                }
            }
        } // while (IsActive())
    }
};

void TIndexerManager::StartCloseIndexer(TAtomicSharedPtr<IIndexer> indexer, ITransaction* transaction) {
    CHECK_WITH_LOG(QueueCloseIndexers.AddAndOwn(THolder(new TCloseIndexTask(indexer, transaction, GetRigidStopSignal()))));
}

void TIndexerManager::StartReopenIndexer(TQueuedDocument replier, const TString& realmName) {
    CHECK_WITH_LOG(QueueReopen.Add(new TReopenTask(*this, replier, Config.Common.Owner, realmName)));
}

IIndexer* TIndexerManager::GetDiskIndexer(const TString& indexDirName, const TString& tempDirName, int shard) {
    DEBUG_LOG << "TIndexerManager::GetIndexer: " << indexDirName << Endl;
    return new TIndexer(indexDirName, tempDirName, Config.GetDirectory(), Config, IndexStorage, shard);
}

void TIndexerManager::OpenIndexers() {
    TGuard<TMutex> g(Mutex);
    VERIFY_WITH_LOG(!IsActive, "Indexers already started");
    IsActive = true;

    const ui32 shards = Config.Common.Owner.GetShardsConfig().Number;
    Indexers.Reset(new TIndexersList(Config.Common.Owner, IndexStorage, shards, IndexPoolSize));
    DEBUG_LOG << "Opening indexers" << Endl;
    for (ui32 shard = 0; shard < shards; ++shard) {
        Indexers->Get(shard)->SpawnAllIndexers(shard);
    }
    IndexHandler.Start(IndexPoolSize);
    for (ui32 i = 0; i < IndexPoolSize; ++i) {
        IndexHandler.SafeAdd(CreateIndexersThread(i));
    }
    QueueCloseIndexers.Start(Config.Common.Owner.GetRealmListConfig().CloseThreads);
    QueueReopen.Start(1);
}

void TIndexerManager::ReopenIndexers() {
    TGuard<TMutex> g(Mutex);
    if (IsActive) {
        CloseIndexers(false);
        OpenIndexers();
    }
}

void TIndexerManager::CloseIndexers(bool rigidStop) {
    ThreadDisableBalloc();
    TGuard<TMutex> g(Mutex);
    RigidStop = rigidStop;
    if (IsActive) {
        IsActive = false;
        NOTICE_LOG << "Closing indexers" << Endl;
        Indexers->Stop(rigidStop);
        NOTICE_LOG << "Closing indexers...OK" << Endl;

        NOTICE_LOG << "Closing reopen threads" << Endl;
        QueueReopen.Stop();
        NOTICE_LOG << "Closing reopen threads...OK" << Endl;

        NOTICE_LOG << "Joining indexers threads" << Endl;
        IndexHandler.Stop();
        IndexHandler.Start(IndexPoolSize);
        for (ui32 indexer = 0; indexer < Indexers->Size(); indexer++) {
            TVector<TAtomicSharedPtr<IIndexer> > indexersForClose = Indexers->Get(indexer)->GetIndexers();
            for (TAtomicSharedPtr<IIndexer> indexerForClose : indexersForClose) {
                StartCloseIndexer(indexerForClose, &CloseTransaction);
            }
        }
        IndexHandler.Stop();
        NOTICE_LOG << "Joining indexers threads...OK" << Endl;

        NOTICE_LOG << "Waiting in close tasks..." << Endl;
        WaitCloseTasks();
        NOTICE_LOG << "Waiting in close tasks...OK" << Endl;

        NOTICE_LOG << "Closing indexes" << Endl;
        QueueCloseIndexers.Stop();
        NOTICE_LOG << "Closing indexes...OK" << Endl;
    }
}

TIndexersThread* TIndexerManager::CreateIndexersThread(ui32 threadId) {
    return new TIndexersThreadDisk(Indexers, DocumentsQueue, Deleter, Config.Common.Owner, ServiceName, threadId, *this);
}

bool TIndexerManager::Process(IMessage* message) {
    TMessageCollectServerInfo* messCollect = dynamic_cast<TMessageCollectServerInfo*>(message);
    if (messCollect) {
        IndexRateMeterDisk.Rate(*messCollect->GetDiskIndexRps());
        IndexRateMeterMemory.Rate(*messCollect->GetMemoryIndexRps());
        return true;
    }

    TMessageReopenIndexes* messReopen = dynamic_cast<TMessageReopenIndexes*>(message);
    if (messReopen) {
        TQueuedDocument document = messReopen->GetDocument();
        TString realmName = messReopen->GetRealmName();
        if (IsActive) {
            StartReopenIndexer(document, realmName);
        } else if (document) {
            document->SetStatus(NRTYServer::TReply::OK, "Server stopped");
        }
        return true;
    }

    TMessageDocumentIndexed* messDocIndexed = dynamic_cast<TMessageDocumentIndexed*>(message);
    if (messDocIndexed) {
        if (messDocIndexed->GetType() == "memory")
            IndexRateMeterMemory.Hit();
        else
            IndexRateMeterDisk.Hit();
        return true;
    }

    if (auto messageFastUpdate = message->As<TMessageFastUpdateInvoked>()) {
        if (messageFastUpdate->Count) {
            IndexRateMeterDisk.Hit();
        }
        return true;
    }

    return false;
}

void TIndexerManager::ResetIndexers(bool force, ui32 shard, bool doLock, bool timeToLiveSec, bool memoryIndexer, const TString& realmName) {
    TVector<TAtomicSharedPtr<IIndexer> > indexersForClose;
    {
        TAutoPtr<TWriteShardIndexGuard> gShard;
        if (doLock)
            gShard = new TWriteShardIndexGuard(*Indexers.Get(), shard);
        indexersForClose = Indexers->Get(shard)->ResetIndexers(force, shard, timeToLiveSec, memoryIndexer, realmName);
    }
    for (TAtomicSharedPtr<IIndexer> indexer : indexersForClose) {
        StartCloseIndexer(indexer, &CloseTransaction);
    }
}
