#pragma once

#include <saas/library/searchserver/replier.h>

#include <search/request/data/reqdata.h>

#include <search/session/reqenv.h>

#include <kernel/search_daemon_iface/reqtypes.h>

#include <util/system/yassert.h>
#include <util/generic/maybe.h>

class ISearchContext;

namespace NSearchProxy {
    const TString PROXY_TYPE_META = "meta";

    class TExtendedContextBase {
    private:
        TInstant ProcessingStartTime_;
        TInstant ReportStartTime_;
        int HttpStatus_ = 0;
        bool Timeouted_ = false;
        ERequestType RequestType_ = RT_Unknown;
        TString MetaSearchType_ = PROXY_TYPE_META;

    public:
        inline TInstant ProcessingStartTime() const {
            return ProcessingStartTime_;
        }
        inline void SetProcessingStartTime(TInstant value) {
            ProcessingStartTime_ = value;
        }
        inline TInstant ReportStartTime() const {
            return ReportStartTime_;
        }
        inline void SetReportStartTime(TInstant value) {
            ReportStartTime_ = value;
        }
        inline bool Timeouted() const {
            return Timeouted_;
        }
        inline void MarkTimeouted(bool value = true) {
            Timeouted_ = value;
        }
        inline int HttpStatus() const {
            return HttpStatus_;
        }
        inline void SetHttpStatus(int status) {
            HttpStatus_ = status;
        }
        inline ERequestType RequestType() const {
            return RequestType_;
        }
        inline void SetRequestType(ERequestType type) {
            RequestType_ = type;
        }
        inline bool ForceLogging() const {
            return Timeouted() || IsServerError(HttpStatus()) || IsUserError(HttpStatus());
        }
        inline const TString& MetaSearchType() const {
            return MetaSearchType_;
        }
        inline void SetMetaSearchType(const TString& value) {
            MetaSearchType_ = value;
        }
    };
}

class TExtendedReplyContext: public NSearchProxy::TExtendedContextBase {
private:
    const IReplyContext* Context;

    ui32 DocumentsCount = 0;
    ui64 ReportByteSize = 0;
    bool AnswerIsComplete_ = true;
    ui32 FailedAttempts = 0;

public:
    TExtendedReplyContext(IReplyContext* context)
        : Context(context)
    {
        SetRequestType(RT_Search);
    }

    ui32 GetTotalDocsCount() const {
        // Because of single usage in proxy_meta(KV-storage), it is the same as GetReportDocsCount
        return DocumentsCount;
    }

    inline ui32 GetReportDocsCount() const {
        return DocumentsCount;
    }
    inline void SetReportDocsCount(ui32 value) {
        DocumentsCount = value;
    }

    inline ui64 GetReportByteSize() const {
        return ReportByteSize;
    }
    inline void SetReportByteSize(ui64 value) {
        ReportByteSize = value;
    }

    const TSearchRequestData& GetRequestData() const {
        return Context->GetRequestData();
    }

    inline TReqEnv* GetReqEnv() const {
        return nullptr;
    }

    inline bool NotFetched() const {
        return false;
    }

    inline bool AnswerIsComplete() const {
        return AnswerIsComplete_;
    }
    inline void SetAnswerIsComplete(bool answerIsComplete) {
        AnswerIsComplete_ = answerIsComplete;
    }

    inline ui32 GetFailedAttempts() const {
        return FailedAttempts;
    }
    inline void SetFailedAttempts(ui32 failedAttempts) {
        FailedAttempts = failedAttempts;
    }

    ui64 GetActiveObjects() const {
        return Context->GetActiveObjects();
    }
    TDuration GetCommonQueueDuration() const {
        return Context->GetCommonQueueDuration();
    }
};

class TExtendedSearchContext: public NSearchProxy::TExtendedContextBase {
private:
    ISearchContext** SearchContext_;

    mutable TMaybe<bool> NotFetched_;

    ui64 ReportByteSize = 0;
    ui64 ReportDocCount = 0;
public:
    TExtendedSearchContext(ISearchContext** searchContext)
        : SearchContext_(searchContext)
    {}

    ui32 GetTotalDocsCount() const {
        return SC()->Cluster()->TotalDocCount(0);
    }

    ui64 GetReportByteSize() const {
        return ReportByteSize;
    }

    void SetReportByteSize(ui64 value) {
        ReportByteSize = value;
    }

    ui64 GetReportDocsCount() const {
        return ReportDocCount;
    }

    void SetReportDocsCount(ui64 value) {
        ReportDocCount = value;
    }

    inline ui32 GetFailedAttempts() const {
        return 0;
    }

    const TSearchRequestData& GetRequestData() const {
        TReqEnv* reqEnv = static_cast<TReqEnv*>(SC()->ReqEnv());
        return *reqEnv->RequestData;
    }

    TReqEnv* GetReqEnv() const {
        return static_cast<TReqEnv*>(SC()->ReqEnv());
    }

    ISearchContext** MutableSC() {
        return SearchContext_;
    }
    ISearchContext* SC() const {
        Y_ASSERT(SearchContext_);
        return *SearchContext_;
    }
    operator ISearchContext*() const {
        return SC();
    }
    operator ISearchContext**() {
        return MutableSC();
    }
    ISearchContext* operator->() const {
        return SC();
    }
    bool NotFetched() const {
        if (NotFetched_.Defined()) {
            return NotFetched_.GetRef();
        }
        TSearcherPropsRef props = GetReqEnv()->GetProperties();
        for (auto& prop: *props) {
            if (prop.first.StartsWith("RemovedDoc_")) {
                if (prop.second.Contains("/not-fetched")) {
                    NotFetched_ = true;
                    return true;
                }
            }
        }
        NotFetched_ = false;
        return false;
    }

    bool AnswerIsComplete() const {
        return GetReqEnv()->AnswerIsComplete();
    }
};
