#pragma once

#include "record.h"
#include "record_data.h"

#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/log/errorlog.h>

#include <util/generic/vector.h>
#include <util/generic/scope.h>

namespace NReport {

    struct TRequestCtx {
        TInstant Start;
    };


    class TIoCounter final : public IIoInput, public IHttpOutput {
    public:
        TIoCounter(
            const TRecordDataVec& counters,
            const TConnDescr& descr,
            const TRequestCtx& ctx,
            IIoInput& in,
            IHttpOutput& out,
            bool hideLegacySignals
        ) noexcept
            : Counters_(counters)
            , Descr_(descr)
            , Ctx_(ctx)
            , Input_(in)
            , Output_(out)
            , HideLegacySignals_(hideLegacySignals)
        {
            if (Descr_.Request) {
                UpdateInputSpeedAndSizes(Descr_.Request->EncodedSize());
            }
        }

        void MarkResponseSent() noexcept {
            ResponseSent_ = true;
        }

        ~TIoCounter() {
            RegisterResponseEnd();
        }

    private:
        TInstant GetCurrentTime() const noexcept {
            return TInstant::Now();
        }

        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            const TInstant start = GetCurrentTime();
            Y_DEFER {
                const TDuration duration = GetCurrentTime() - start;
                ClientTime_ += duration;
            };

            Y_PROPAGATE_ERROR(Input_.Recv(lst, deadline));

            UpdateInputSpeedAndSizes(lst.size());

            return {};
        }

        TError DoSendHead(TResponse&& response, const bool forceClose, TInstant deadline) override {
            const TInstant start = GetCurrentTime();
            Y_DEFER {
                const TDuration duration = GetCurrentTime() - start;
                ClientTime_ += duration;
            };

            StatusCode_ = response.ResponseLine().StatusCode;

            UpdateOutputSpeedAndSizes(response.EncodedSize());

            Y_PROPAGATE_ERROR(Output_.SendHead(std::move(response), forceClose, deadline));

            if (HTTP_SWITCHING_PROTOCOLS == StatusCode_) {
                MarkResponseSent();
                RegisterResponseEnd();
            }
            return {};
        }

        TError DoSendTrailers(THeaders&& trailers, TInstant deadline) override {
           return Output_.SendTrailers(std::move(trailers), deadline);
        }

        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            if (HTTP_SWITCHING_PROTOCOLS == StatusCode_) {
                return Output_.Send(std::move(lst), deadline);
            }

            const TInstant start = GetCurrentTime();
            Y_DEFER {
                const TDuration duration = GetCurrentTime() - start;
                ClientTime_ += duration;
            };

            const bool finished = lst.Empty();

            UpdateOutputSpeedAndSizes(lst.size());

            Y_PROPAGATE_ERROR(Output_.Send(std::move(lst), deadline));

            if (finished) {
                MarkResponseSent();
                RegisterResponseEnd();
                if (HTTP_CONTINUE == StatusCode_) {
                    Reset100Continue();
                }
            }

            return {};
        }

        void RegisterResponseEnd() noexcept {
            if (Commited_) {
                return;
            }
            Commited_ = true;

            TLocalStats stats;
            stats.TotalTime = Now() - Ctx_.Start;
            stats.BackendTime = stats.TotalTime - ClientTime_;
            stats.InputSize = InputSize_;
            stats.OutputSize = OutputSize_;
            stats.StatusCode = ResponseSent_ ? StatusCode_ : 0;

            if (!HideLegacySignals_) {
                stats.DimData = NLegacyRange::GetDimData(Descr_, ResponseSent_ ? StatusCode_ : 0);
            }

            ApplyAll([&](TRecordData& data) {
                data.RegisterResponseEnd(stats);
            });
        }

        void Reset100Continue() noexcept {
            StatusCode_ = 0;
            OutputSize_ = 0;
            Commited_ = false;
            ResponseSent_ = false;
        }

    private:
        void UpdateOutputSpeedAndSizes(const size_t sz) {
            ApplyAll([&](TRecordData& data) {
                data.UpdateOutputSpeed(sz);
            });

            if (sz > 0 && OutputSize_ == 0) {
                ApplyAll([this](TRecordData& data) {
                    data.RegisterTimeToFirstByte(Now() - Ctx_.Start);
                });
            }

            OutputSize_ += sz;
        }

        void UpdateInputSpeedAndSizes(const size_t sz) {
            ApplyAll([&](TRecordData& data) {
                data.UpdateInputSpeed(sz);
            });
            InputSize_ += sz;
        }

        template <class TFunc>
        void ApplyAll(TFunc&& func) const {
            for (const auto& counter : Counters_) {
                func(*counter);
            }
        }

    private:
        const TRecordDataVec& Counters_;
        const TConnDescr& Descr_;
        const TRequestCtx& Ctx_;
        IIoInput& Input_;
        IHttpOutput& Output_;

        ui32 StatusCode_ = 0;

        TDuration ClientTime_;
        size_t OutputSize_ = 0;
        size_t InputSize_ = 0;
        bool HideLegacySignals_ = false;
        bool ResponseSent_ = false;
        bool Commited_ = false;
    };
}
