#include "module.h"
#include "parser.h"
#include "record.h"
#include "record_data.h"
#include "stream.h"

#include <balancer/modules/report/lib/uuid_parser.h>

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/helpers/common_parsers.h>
#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/matcher/matcher.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/process/thread_info.h>

#include <library/cpp/monlib/metrics/labels.h>

#include <util/generic/scope.h>
#include <util/string/cast.h>
#include <util/string/escape.h>
#include <util/string/subst.h>
#include <util/thread/singleton.h>

namespace NReport {

    class TRecordsIndex : public TWithMutableInstance<TRecordsIndex> {
    public:
        template <class T>
        TRecord* Find(const T& uuid, size_t workerId) noexcept {
            Y_ASSERT(WorkerCount_.Defined() && 0 <= workerId && workerId <= *WorkerCount_);
            auto res = MapFindPtr(Records_[workerId], uuid);
            if (res != nullptr) {
                return res->Get();
            } else {
                return nullptr;
            }
        }

        TRecord* RegisterRecord(THolder<TRecord> record, size_t workerId = MasterWorkerId) {
            Y_ASSERT(WorkerCount_.Defined());
            TString uuid = record->Uuid();
            if (!uuid.empty()) {
                Y_ASSERT(workerId == MasterWorkerId);
                TRecord* existingRecord = Find(uuid, MasterWorkerId);
                if (existingRecord) {
                    Y_ENSURE_EX(existingRecord->UpdateRanges(record->Data().Ranges()),
                                TConfigParseError() << "records with same uuid " << uuid.Quote() <<
                                    " have different ranges");
                    existingRecord->MarkDuplicate();
                    return existingRecord;
                }
                Records_[MasterWorkerId][uuid] = std::move(record);
                return Records_[MasterWorkerId][uuid].Get();
            } else {
                EmptyRecords_[workerId].push_back(std::move(record));
                return EmptyRecords_[workerId].back().Get();
            }
        }

        void InitializeRecordsForWorker(size_t workerId) {
            Y_ASSERT(WorkerCount_.Defined() && 0 <= workerId && workerId <= *WorkerCount_);
            if (!Records_[workerId].empty()) {
                // Another module has already initialized records for the current worker.
                return;
            }
            for (auto& [uuid, record] : Records_[MasterWorkerId]) {
                Records_[workerId].emplace(uuid, MakeHolder<TRecord>(*record, workerId));
            }
        }

        void SetWorkerCount(size_t workerCount) {
            if (WorkerCount_.Defined()) {
                Y_VERIFY(*WorkerCount_ == workerCount, "different worker count came from an another module");
            } else {
                WorkerCount_ = workerCount;
                Records_.resize(workerCount + 1);
                EmptyRecords_.resize(workerCount + 1);
            }
        }

    private:
        static constexpr size_t MasterWorkerId = 0;

        TVector<THashMap<TString, THolder<TRecord>>> Records_;
        TVector<TVector<THolder<TRecord>>> EmptyRecords_;

        TMaybe<size_t> WorkerCount_;
    };


    TVector<TRecord*> GetRefers(
        const TVector<TString>& refersStrings,
        const TVector<TRecord*> commonRecords,
        TRecordsIndex& recordsIndex,
        size_t workerId
    ) {
        TVector<TRecord*> refers;
        for (const TString& uuid : refersStrings) {
            bool selfRefer = false;
            for (const auto& record : commonRecords) {
                if (record->Uuid() == uuid) {
                    selfRefer = true;
                    break;
                }
            }

            if (!selfRefer) {
                auto* record = recordsIndex.Find(uuid, workerId);
                Y_ENSURE_EX(record, TConfigParseError() << "record with uuid " << uuid.Quote() << " does not exist");
                Y_ENSURE_EX(record->IsUnique(), TConfigParseError() << "referring non-unique uuid is forbidden");
                refers.push_back(record);
            }
        }
        return refers;
    }

}

using namespace NConfig;
using namespace NSrvKernel;
using namespace NReport;

Y_TLS(report) {
    TTls(const TVector<TRecord*>& commonRecords,
         const TVector<TVector<TRecord*>>& matchersRecords,
         const TVector<TString>& refersStrings,
         TRecordsIndex& recordsIndex,
         size_t workerId)
     {
        recordsIndex.InitializeRecordsForWorker(workerId);

        auto findOrCreateThreadLocalRecord = [&](const auto& masterRecord) {
            if (masterRecord->Uuid().empty()) {
                return recordsIndex.RegisterRecord(MakeHolder<TRecord>(*masterRecord, workerId), workerId);
            } else {
                TRecord* threadLocalRecord = recordsIndex.Find(masterRecord->Uuid(), workerId);
                Y_ASSERT(threadLocalRecord);
                return threadLocalRecord;
            }
        };

        WorkerCommonRecords.reserve(commonRecords.size());
        for (const auto& record : commonRecords) {
            WorkerCommonRecords.push_back(findOrCreateThreadLocalRecord(record));
        }

        WorkerMatchersRecords.reserve(matchersRecords.size());
        for (const auto& matcherRecords : matchersRecords) {
            TVector<TRecord*> records;
            records.reserve(matcherRecords.size());
            for (const auto& record : matcherRecords) {
                records.push_back(findOrCreateThreadLocalRecord(record));
            }
            WorkerMatchersRecords.push_back(std::move(records));
        }

        WorkerRefers = GetRefers(refersStrings, WorkerCommonRecords, recordsIndex, workerId);
    }

    TRecordDataVec SelectRecordData(const TConnDescr& descr, const TVector<TAtomicSharedPtr<NSrvKernel::IRequestMatcher>>& matchers) noexcept {
        TRecordDataVec counters;

        for (auto& record : WorkerCommonRecords) {
            counters.emplace_back(&record->Data());
        }

        if (WorkerCommonRecords.size() == 1 && WorkerCommonRecords.front()->IsUnique()) {
            // Uuid is unique, in this case matcher is stored in record
            const auto& record = WorkerCommonRecords.front();
            for (size_t matcherIndex = 0; matcherIndex < record->Matchers().size(); ++matcherIndex) {
                if (record->Matchers()[matcherIndex]->Match(descr)) {
                    counters.push_back(record->MatcherData()[matcherIndex]);
                }
            }
        } else {
            for (size_t matcherIndex = 0; matcherIndex < matchers.size(); ++matcherIndex) {
                auto& matcher = matchers[matcherIndex];
                if (matcher->Match(descr)) {
                    for (auto& counter : WorkerMatchersRecords[matcherIndex]) {
                        counters.emplace_back(&counter->Data());
                    }
                }
            }
        }

        for (auto& refer : WorkerRefers) {
            counters.emplace_back(&refer->Data());
            for (size_t matcherIndex = 0; matcherIndex < refer->Matchers().size(); ++matcherIndex) {
                if (refer->Matchers()[matcherIndex]->Match(descr)) {
                    counters.push_back(refer->MatcherData()[matcherIndex]);
                }
            }
        }

        return counters;
    }

    TVector<TRecord*> WorkerCommonRecords;
    TVector<TVector<TRecord*>> WorkerMatchersRecords;
    TVector<TRecord*> WorkerRefers;
};

MODULE_WITH_TLS_BASE(report, TModuleWithSubModule) {
    friend struct TParser<TModule>;
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
        , RecordsIndex_(&TRecordsIndex::Instance())
    {
        RecordsIndex_->SetWorkerCount(mp.Control->GetCountOfChildren());

        TParser<TModule> parser(this);

        Config->ForEach(&parser);

        parser.SetDefaultValues();

        switch (parser.SignalSet) {
        case ESignalSet::Minimal:
            parser.DisabledSignals.emplace("conn_fail");
            parser.DisabledSignals.emplace("backend_fail");
            parser.DisabledSignals.emplace("sc_1xx");
            parser.DisabledSignals.emplace("outgoing_1xx");
            parser.DisabledSignals.emplace("no_backends_error");
            [[fallthrough]];

        case ESignalSet::Reduced:
            parser.DisabledSignals.emplace("ka");
            parser.DisabledSignals.emplace("nka");
            parser.DisabledSignals.emplace("reused");
            parser.DisabledSignals.emplace("nreused");
            parser.DisabledSignals.emplace("backend_keepalive_reused");
            parser.DisabledSignals.emplace("client_error");
            parser.DisabledSignals.emplace("input_speed");
            parser.DisabledSignals.emplace("output_speed");
            parser.DisabledSignals.emplace("client_fail_time");
            parser.DisabledSignals.emplace("other_fail");
            parser.DisabledSignals.emplace("hedged_attempts");
            parser.DisabledSignals.emplace("hedged_succ");
            [[fallthrough]];

        case ESignalSet::Default:
            parser.DisabledSignals.emplace("backend_short_read_answer");
            parser.DisabledSignals.emplace("backend_write_error");
            parser.DisabledSignals.emplace("input_headers_size");
            parser.DisabledSignals.emplace("backend_fail_input_headers_size");
            [[fallthrough]];

        case ESignalSet::Full:
            break;

        default:
            Y_UNREACHABLE();
        }

        for (const auto& signal : parser.EnabledSignals) {
            parser.DisabledSignals.erase(signal);
        }

        for (const TString& uuid : parser.Uuids) {
            TRecord* record = RecordsIndex_->RegisterRecord(MakeHolder<TRecord>(
                TRecordRanges(
                    parser.LegacyTimeRanges,
                    parser.BackendTimeRanges,
                    parser.ClientFailTimeRanges,
                    parser.FirstByteTimeRanges,
                    parser.InputSizeRanges,
                    parser.OutputSizeRanges,
                    parser.InputHeadersSizeRanges,
                    parser.InputHeadersSizeRanges,
                    parser.DimFilter,
                    parser.DisableSslness,
                    parser.DisableRobotness
                ),
                parser.OutgoingCodes,
                uuid,
                parser.Labels,
                mp.Control->SharedStatsManager(),
                parser.DisabledSignals
            ));
            CommonRecords_.push_back(record);
        }

        if (parser.Uuids.size() > 1) {
            for (const auto& record : CommonRecords_) {
                record->MarkDuplicate();
            }
        }

        MatchersRecords_.reserve(parser.Label2MatcherMap.size());
        for (auto& matcher : parser.Label2MatcherMap) {
            MatchersRecords_.emplace_back();
            MatchersRecords_.back().reserve(parser.Uuids.size());

            for (const TString& uuid : parser.Uuids) {
                TRecord* record = RecordsIndex_->RegisterRecord(MakeHolder<TRecord>(
                    TRecordRanges {
                        parser.LegacyTimeRanges,
                        parser.BackendTimeRanges,
                        parser.ClientFailTimeRanges,
                        parser.FirstByteTimeRanges,
                        parser.InputSizeRanges,
                        parser.OutputSizeRanges,
                        parser.InputHeadersSizeRanges,
                        parser.InputHeadersSizeRanges,
                        parser.DimFilter,
                        parser.DisableSslness,
                        parser.DisableRobotness
                    },
                    parser.OutgoingCodes,
                    uuid + "_" + matcher.first,
                    parser.Labels,
                    mp.Control->SharedStatsManager(),
                    parser.DisabledSignals
                ));
                MatchersRecords_.back().emplace_back(record);
            }

            Matchers_.push_back(std::move(matcher.second));
        }

        Y_ENSURE_EX(Submodule_, TConfigParseError() << "no submodule configured");
    }

private:
    START_PARSE {
        PARSE_EVENTS;

        ON_KEY("just_storage", JustStorage_) {
            return;
        }

        if (key == "refers") {
            CheckParse(key, [&]() {
                const TString refersValue = value->AsString();
                for (const auto& i: StringSplitter(refersValue).Split(',')) {
                    RefersStrings_.push_back(CheckedUuid(ToString(i.Token())));
                }
            });
            return;
        }

        Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
        return;
    } END_PARSE

    bool DoCanWorkWithoutHTTP() const noexcept override {
        return true;
    }

    bool DoExtraAccessLog() const noexcept override {
        return true;
    }

    void DoFinalize() override {
        HideLegacySignals_ = Control->SharedStatsManager().HideLegacySignals();

        NLegacyRange::Initialize();

        if (CommonRecords_.size() == 1 && CommonRecords_.front()->IsUnique()) {
            CommonRecords_.front()->Matchers() = std::move(Matchers_);
            for (auto& matcherData : MatchersRecords_) {
                CommonRecords_.front()->MatcherData().emplace_back(&matcherData.front()->Data());
            }
        }

        GetRefers(RefersStrings_, CommonRecords_, *RecordsIndex_, /*workerId =*/ 0);
    }

    THolder<TTls> DoInitTls(IWorkerCtl* worker) override {
        return MakeHolder<TTls>(CommonRecords_, MatchersRecords_, RefersStrings_, *RecordsIndex_, worker->WorkerId());
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (JustStorage_) {
            return Submodule_->Run(descr);
        }

        TRecordDataVec counters = tls.SelectRecordData(descr, Matchers_);

        for (const auto& c : counters) {
            if (c->Uuid()) {
                descr.ExtraAccessLog << " u:" << c->Uuid();
            }
        }

        TRequestCtx ctx{
            Now()
        };

        RegisterRequest(counters);

        TIoCounter counter{
            counters,
            descr,
            ctx,
            *descr.Input,
            *descr.Output,
            HideLegacySignals_
        };

        RegisterInProgress(counters);

        const auto prevStats = descr.Properties->ConnStats;
        bool success = false;
        bool closedByClient = false;
        TMaybe<size_t> headersSize;

        Y_DEFER {
            if (success) {
                counter.MarkResponseSent();
                RegisterSuccess(counters, headersSize);
            } else {
                RegisterFail(
                    Now() - ctx.Start,
                    descr.Properties->ConnStats - prevStats,
                    counters,
                    descr.ExtraAccessLog,
                    closedByClient,
                    headersSize
                );
            }

            for (const auto& cnt : counters) {
                cnt->RegisterConnStats(descr.Properties->ConnStats - prevStats);
            }
        };

        if (descr.Request) {
            headersSize = descr.Request->EncodedSize() - descr.Request->RequestLine().EncodedSize();
            const auto& props = descr.Request->Props();

            if (props.KeepAlive || props.HTTP2) {
                RegisterKeepAlive(counters);
            } else {
                RegisterNonKeepAlive(counters);
            }

            if (props.Reused || props.HTTP2) {
                RegisterReused(counters);
            } else {
                RegisterNotReused(counters);
            }
        }

        Y_TRY(TError, error) {
            return Submodule_->Run(descr.Copy(counter, counter));
        } Y_CATCH {
            if (error.GetAs<TForceStreamClose>() || IsHTTP2CancelledByClient(descr.Request)) {
                success = true;
            } else if (IsHTTP2ClosedByClient(descr.Request)
                    || !IsHTTP2(descr.Request) && ErrorHasErrno(error, {ECANCELED}))
            {
                closedByClient = true;
            }
            return error;
        }

        success = true;
        return {};
    }

private:
    void RegisterInProgress(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterInProg();
        }
    }

    void RegisterSuccess(const TRecordDataVec& counters, const TMaybe<size_t>& headersSize) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterSucc(headersSize);
        }
    }

    void RegisterFail(
        const TDuration duration,
        const TConnStats& delta,
        const TRecordDataVec& counters,
        const TAccessLogOutput& log,
        bool clientClosed,
        const TMaybe<size_t>& headersSize
    ) const noexcept {
        EFail fail = EFail::Other;

        if (delta.ClientError || clientClosed) {
            fail = EFail::Client;
        } else if (delta.BackendError) {
            fail = EFail::Backend;
        } else if (delta.HasConnError()) {
            log << " " << EFail::Backend;
            fail = EFail::Conn;
        }

        log << " " << fail;

        for (const auto& counter: counters) {
            counter->RegisterFail(duration, fail, headersSize);
        }
    }

    void RegisterRequest(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterRequest();
        }
    }

    void RegisterKeepAlive(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterKeepAlive();
        }
    }

    void RegisterNonKeepAlive(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterNonKeepAlive();
        }
    }

    void RegisterReused(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterReused();
        }
    }

    void RegisterNotReused(const TRecordDataVec& counters) const noexcept {
        for (const auto& counter: counters) {
            counter->RegisterNotReused();
        }
    }

private:
    TVector<TAtomicSharedPtr<NSrvKernel::IRequestMatcher>> Matchers_;
    TVector<TRecord*> CommonRecords_;
    TVector<TVector<TRecord*>> MatchersRecords_;

    TVector<TString> RefersStrings_;

    TRecordsIndex* const RecordsIndex_;

    bool HideLegacySignals_ = false;
    bool JustStorage_ = false;
};

IModuleHandle* NModReport::Handle() {
    return TModule::Handle();
}
