#include "snapshot.h"
#include "messages.h"

#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/realm_config.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/common/message_collect_server_info.h>
#include <saas/rtyserver/docfetcher/library/stream_basics.h>
#include <saas/rtyserver/docfetcher/library/types.h>
#include <saas/rtyserver/indexer_core/merger_interfaces.h>
#include <saas/rtyserver/synchronizer/library/sync.h>
#include <saas/rtyserver/unistat_signals/signals.h>
#include <saas/library/daemon_base/threads/rty_pool.h>
#include <saas/util/logging/replies.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/mediator/global_notifications/system_status.h>
#include <library/cpp/retry/retry.h>

constexpr ui8 AuxStreamId = 0;
constexpr ui8 PositionAuxId = 1;

namespace {
    bool AllowEmptyIndex(const NFusion::TSnapshotStreamConfig& config) {
        return TInstant::Now() < config.AllowEmptyIndexUntil;
    }
}

NFusion::TSnapshotStream::TSnapshotStream(const TDocFetcherModule& owner, const TSnapshotStreamConfig& config, const TRTYServerConfig& globalConfig, TLog& log)
    : TBaseStream(owner, config, log, "DfSnapshotStrm")
    , RTYServerConfig(globalConfig)
    , SnapshotStreamConfig(config)
    , LastFetchedTimestamp(TInstant::Zero())
{
    // TODO: implement support for OnRequest mode
    AssertCorrectConfig(SnapshotStreamConfig.Mode == ESnapshotStreamMode::OnSchedule, "only ESnapshotStreamMode::OnSchedule is supported");
    if (SnapshotStreamConfig.ConsumeMode == NRTYServer::EConsumeMode::HardReplace) {
        // Idea of 'hard_replace' is to deploy a new index version with a reasonable degrade level.
        // If there's no FixedTimstamp then all backends may clear indices simultaneously.
        AssertCorrectConfig(!!SnapshotStreamConfig.FixedTimestamp, "Consume mode 'hard_replace' should only be used with the FixedTimestamp");
    }
    AssertCorrectConfig(!SnapshotStreamConfig.UseGroups || SnapshotStreamConfig.SnapshotManager == "yt", "Support for grouped publish is not implemented in Zookeeper SnapshotManager yet");
    AssertCorrectConfig(!SnapshotStreamConfig.UseGroups || (SnapshotStreamConfig.ConsumeMode == NRTYServer::EConsumeMode::Apply
                        || SnapshotStreamConfig.ConsumeMode == NRTYServer::EConsumeMode::Append)
                        , "replace and hard_replace modes are not compatible with grouped publish yet");

    AssertCorrectConfig(SnapshotStreamConfig.UseGroups || (SnapshotStreamConfig.ConsumeMode != NRTYServer::EConsumeMode::Apply)
                        , "apply mode is compatible only with grouped publish");
}

NFusion::TSnapshotStream::~TSnapshotStream() = default;

bool NFusion::TSnapshotStream::Stopped() const {
    return IThreadLoop::Stopped();
}

IThreadFactory* NFusion::TSnapshotStream::GetThreadPool() {
    return Singleton<TRTYPools>()->Get(TRTYPools::TLoadtype::Batch);
}

void NFusion::TSnapshotStream::OnStreamStart() {
    StopEvent.Reset();
    IIndexSnapshotManager::TContext ismContext{SnapshotStreamConfig.SnapshotServer, SnapshotStreamConfig.SnapshotPath, SnapshotStreamConfig.SnapshotToken,
                                               TRetryOptions(SnapshotStreamConfig.FetchMaxAttempts), SnapshotStreamConfig.YTHosts, SnapshotStreamConfig.UseGroups};
    SnapshotManager.Reset(IIndexSnapshotManager::TFactory::Construct(SnapshotStreamConfig.SnapshotManager, ismContext));

    InitTimestamp();
    InitSearch();
}

void NFusion::TSnapshotStream::InitTimestamp() {
    LastFetchedTimestamp = TInstant::Zero();
    LastFetchedPosition.TableRecord = -1;

    auto snapshot = GetTimestampSnapshot();
    auto it = snapshot.find(GetSubStream(SnapshotStreamConfig.StreamId, AuxStreamId));
    if (it == snapshot.end()) {
        return;
    }
    NRTYServer::TTimestampValue raw = it->second.MaxValue;
    if (raw >= static_cast<NRTYServer::TTimestampValue>(Max<ui64>())) {
        return;
    }
    TInstant t = TInstant::Seconds(SafeIntegerCast<ui64>(raw));
    if (t > TInstant::Now()) {
        return;
    }
    SetLastFetchedTimestamp(t);
    it = snapshot.find(GetSubStream(SnapshotStreamConfig.StreamId, PositionAuxId));
    if (it == snapshot.end()) {
        return;
    }
    SetLastFetchedPosition(NFusion::TimestampToPosition(it->second.MaxValue));
}

void NFusion::TSnapshotStream::InitSearch() {
    TMessageCollectServerInfo info(/* isHumanReadable = */ false);
    bool disableSearch = false;
    if (!AllowEmptyIndex(SnapshotStreamConfig) && (!SendGlobalMessage(info) || info.GetSearchableDocs() == 0)) {
        NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": disabling search because the index is empty" << Endl;
        disableSearch = true;
    }
    if (SnapshotStreamConfig.FixedTimestamp && SnapshotStreamConfig.FixedTimestamp != LastFetchedTimestamp) {
        NOTICE_LOG << "SnapshotStream " << BaseConfig.Name
            << ": disabling search because of the index timestamp mismatch: "
            << SnapshotStreamConfig.FixedTimestamp.Seconds() << " != " << LastFetchedTimestamp.Seconds() << Endl;
        disableSearch = true;
    }
    if (disableSearch) {
        DisableSearch();
    }
}

void NFusion::TSnapshotStream::UpdateSearch() {
    TMessageCollectServerInfo info;
    if (SendGlobalMessage(info) && info.GetSearchableDocs() != 0) {
        if (!SnapshotStreamConfig.FixedTimestamp) {
            // no fixed timestamp, any non-empty index is enough
            EnableSearch();
            NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": enabled search for non-empty index" << Endl;
        } else {
            if (SnapshotStreamConfig.FixedTimestamp == LastFetchedTimestamp) {
                EnableSearch();
                NOTICE_LOG << "SnapshotStream " << BaseConfig.Name
                    << ": enabled search for fixed timestamp " << LastFetchedTimestamp.Seconds() << Endl;
            } else {
                ERROR_LOG << "SnapshotStream " << BaseConfig.Name
                    << ": index fetched, but timestamp doesn't match: "
                    << SnapshotStreamConfig.FixedTimestamp.Seconds()
                    << " != " << LastFetchedTimestamp.Seconds() << Endl;
            }
        }
    } else {
        ERROR_LOG << "SnapshotStream " << BaseConfig.Name
            << ": index fetch is successful, but index is empty" << Endl;
    }
}

void NFusion::TSnapshotStream::SetLastFetchedTimestamp(TInstant t) {
    if (LastFetchedTimestamp >= t) {
        return;
    }

    NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": LastFetchedTimestamp is updated from "
        << LastFetchedTimestamp.Seconds() << " to " << t.Seconds() << Endl;
    LastFetchedTimestamp = t;
}

void NFusion::TSnapshotStream::SetLastFetchedPosition(TMapReducePosition p) {
    if (LastFetchedPosition >= p) {
        return;
    }

    NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": LastFetchedPosition is updated from ("
        << LastFetchedPosition.TableTimestamp << ", " << LastFetchedPosition.TableRecord << ") to ("
        << p.TableTimestamp << ", " << p.TableRecord << ")" << Endl;
    LastFetchedPosition = p;
}

void NFusion::TSnapshotStream::OnStreamBeforeStop() {
}

void NFusion::TSnapshotStream::OnStreamStop() {
    SnapshotManager.Destroy();
}

TDuration NFusion::TSnapshotStream::GetDelay() const {
    if (LastFetchedTimestamp == TInstant::Zero()) {
        return GetDelayFromSubStream(
            GetTimestampSnapshot(),
            GetSubStream(SnapshotStreamConfig.StreamId, AuxStreamId)
        );
    }
    return TInstant::Now() - LastFetchedTimestamp;
}

NFusion::IThreadLoop::TTimeToSleep NFusion::TSnapshotStream::DoStreamIteration() {
    const NUtil::TInterval<ui64> localShard = { SnapshotStreamConfig.ShardMin, SnapshotStreamConfig.ShardMax };
    /* TODO: with OnRequest mode keep track of requested snapshot timestamp
        - in OnSchedule mode it's fine to download a week-old snapshot, they may be updated rarely
        - in OnRequest mode we need to ignore a correct week-old snapshot if other stream asks for a day-old snapshot
        So far we simply fetch anything newer than the current index.
    */
    if (SnapshotStreamConfig.FixedTimestamp && SnapshotStreamConfig.FixedTimestamp == LastFetchedTimestamp) {
        NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": FixedTimestamp " << LastFetchedTimestamp.Seconds()
                   << " is OK, sleeping until docfetcher restart" << Endl;
        return TDuration::Max();
    }

    INFO_LOG << "SnapshotStream " << BaseConfig.Name << ": acquiring list of resources..." << Endl;
    EGroupFetchStatus fetchStatus;
    if (SnapshotStreamConfig.UseGroups) {  // bin deltas
        TInstant ts = LastFetchedTimestamp - TDuration::Seconds(1);
        auto [resources, hasPriorResources] = SnapshotManager->GetGroupedByTSResourcesForShard(
            localShard, ts, TInstant::Max(), Interruptable()
        );

        NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": acquired list of resources, size: " << resources.size() << Endl;

        const bool empty = resources.empty();
        /*
         * In order to track missed snapshots we can check whether we always get the resource
         * that we read in the previous iteration. Buf if the resource was locked at the moment,
         * when the ferryman wanted to delete it, it won't happen,
         * so we need to check if we have previous resources.
         */
        if (empty || !hasPriorResources) {
            TSaasRTYServerSignals::DoUnistatRecordMissedDeltas();

            if (empty) {
                return SnapshotStreamConfig.IterationsPause;
            }
        }
        fetchStatus = FetchLatestAvailableResourceGroup(resources);
    } else {
        IIndexSnapshotManager::TShardResources resources;
        if (SnapshotStreamConfig.FixedTimestamp) {
            resources = SnapshotManager->GetResourceForShardFixed(localShard, SnapshotStreamConfig.FixedTimestamp, Interruptable());
        } else {
            resources = SnapshotManager->GetResourcesForShard(localShard, LastFetchedTimestamp, TInstant::Max(), Interruptable());
        }
        NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": acquired list of resources, size: " << resources.size() << Endl;
        if (resources.empty()) {
            return SnapshotStreamConfig.IterationsPause;
        }
        fetchStatus = FetchLatestAvailableResource(resources);
    }

    if (fetchStatus == EGroupFetchStatus::Failed) {
        ERROR_LOG << "SnapshotStream " << BaseConfig.Name << ": could not fetch resources" << Endl;
        return SnapshotStreamConfig.IterationsPause;
    }
    NOTICE_LOG << "SnapshotStream " << BaseConfig.Name
        << ": successfully fetched resource, checking search status" << Endl;

    UpdateSearch();

    if (SnapshotStreamConfig.UseGroups && fetchStatus != EGroupFetchStatus::NothingToFetch) {
        return SnapshotStreamConfig.PauseAfterPartialConsume;
    }
    return SnapshotStreamConfig.IterationsPause;
}

NFusion::EGroupFetchStatus NFusion::TSnapshotStream::FetchLatestAvailableResource(const IIndexSnapshotManager::TShardResources& resources) {
    for (const auto& resource : resources) {
        if (FetchResource(resource) == EGroupFetchStatus::Completed) {
            SetLastFetchedTimestamp(TInstant::Seconds(resource.GetTimestamp()));
            if (resource.HasTimestampEx()) {
                SetLastFetchedPosition(TMapReducePosition(resource.GetTimestampEx(), resource.GetTimestamp()));
            }
            return EGroupFetchStatus::Completed;
        }
        Sleep(SnapshotStreamConfig.FetchRetryPause);
    }
    return EGroupFetchStatus::Failed;
}

NFusion::EGroupFetchStatus NFusion::TSnapshotStream::FetchResource(const IIndexSnapshotManager::TShardResource& resource) {
    NOTICE_LOG << "SnapshotStream " << BaseConfig.Name << ": trying to fetch resource "
        << "name=" << resource.GetName().Quote() << ", timestamp=" << resource.GetTimestamp() << Endl;

    const auto& fetchConfig = SnapshotStreamConfig.ResourceFetchConfig ? *SnapshotStreamConfig.ResourceFetchConfig : RTYServerConfig.GlobalResourceFetchConfig;
    auto indexProcessorCb = [=](IIndexWithComponents& index) -> bool {
        index.GetTimestamp().Set(GetSubStream(SnapshotStreamConfig.StreamId, AuxStreamId), resource.GetTimestamp());
        if (resource.HasTimestampEx()) {
            SetLastFetchedPosition(TMapReducePosition(resource.GetTimestampEx(), resource.GetTimestamp()));
        }
        index.GetTimestamp().Flush();
        return true;
    };
    auto fetchFunc = [&]() -> bool {
        return NRTYServer::FetchAndConsumeIndices(NRTYServer::TRemoteResource{resource}, SnapshotStreamConfig.ConsumeMode, RTYServerConfig.GetRealmListConfig().GetMainRealmConfig().RealmDirectory, fetchConfig, nullptr, indexProcessorCb);
    };
    if (!DoWithRetryOnRetCode(fetchFunc, TRetryOptions{SnapshotStreamConfig.FetchMaxAttempts, SnapshotStreamConfig.FetchRetryPause})) {
        WARNING_LOG << "SnapshotStream " << BaseConfig.Name << ": Could not fetch resource " << resource.GetName() << Endl;
        return EGroupFetchStatus::Failed;
    }
    return EGroupFetchStatus::Completed;
}

NFusion::EGroupFetchStatus NFusion::TSnapshotStream::FetchResourceGroup(TInstant groupTimestamp, const IIndexSnapshotManager::TShardResources& resources) {
    ui32 consumedResources = 0;
    ui32 skippedResources = 0;
    for (const auto& resource : resources) {
        if (Stopped() || consumedResources == SnapshotStreamConfig.MaxResourcesPerIteration) {
            INFO_LOG << "Partial fetch of resource group (" << resource.GetTimestampEx() << "; " << resource.GetTimestamp() << ") is completed" << Endl;
            break;
        }
        if (resource.HasTimestampEx()) {
            // FIXME: probably doesn't work correctly with FixedTimestamp option
            const TMapReducePosition resourcePosition = TMapReducePosition(resource.GetTimestampEx(), resource.GetTimestamp());
            if (resourcePosition <= LastFetchedPosition) {
                DEBUG_LOG << "Skip resource (" << resource.GetTimestampEx() << "; " << resource.GetTimestamp() << ") as already consumed" << Endl;
                ++skippedResources;
                continue;
            }
        }
        if (FetchResource(resource) == EGroupFetchStatus::Failed) {
            return EGroupFetchStatus::Failed;
        }
        consumedResources++;
    }

    if (consumedResources + skippedResources == resources.size()) {
        INFO_LOG << "Resource group (" << groupTimestamp.Seconds() << ") is fully processed. Consumed resources: " << consumedResources << ", skipped: " << skippedResources << Endl;

        return skippedResources == resources.size() ? EGroupFetchStatus::NothingToFetch : EGroupFetchStatus::Completed;
    }

    return EGroupFetchStatus::Partial;
}

NFusion::EGroupFetchStatus NFusion::TSnapshotStream::FetchLatestAvailableResourceGroup(const IIndexSnapshotManager::TGroupedShardResources& resources) {
    for (const auto& [groupTimestamp, resourceGroup] : resources) {
        if (Stopped()) {
            return EGroupFetchStatus::NothingToFetch;
        }
        const EGroupFetchStatus fetchStatus = FetchResourceGroup(groupTimestamp, resourceGroup);
        if (fetchStatus == EGroupFetchStatus::Completed || fetchStatus == EGroupFetchStatus::NothingToFetch) {
            SetLastFetchedTimestamp(groupTimestamp);
            SetLastFetchedPosition(TMapReducePosition(Max<ui64>(), groupTimestamp.Seconds()));
        }
        if (fetchStatus != EGroupFetchStatus::NothingToFetch) {
            return fetchStatus;
        }
    }
    return EGroupFetchStatus::NothingToFetch;
}

NFusion::IThreadLoop::TTimeToSleep NFusion::TSnapshotStream::OnException() {
    // TODO: add metric for fetch failure
    ERROR_LOG << "An exception occurred in stream " << BaseConfig.Name << ": " << CurrentExceptionMessage() << Endl;
    return SnapshotStreamConfig.IterationsPause;
}

bool NFusion::TSnapshotStream::Process(IMessage* message_) {
    if (auto message = message_->As<TMessageGetDocfetcherStatus>()) {
        message->DocfetcherPresent = true;
        ui64 timestamp = 0;
        if (message->Status.IsMap()) {
            auto status = message->Status.GetMapSafe();
            timestamp = status["SearchableTimestamp"].GetUInteger();
        }
        if (LastFetchedTimestamp.Seconds() > timestamp) {
            message->Status.InsertValue("SearchableTimestamp", LastFetchedTimestamp.Seconds());
        }
        return true;
    }
    return TBaseStream::Process(message_);
}

