#include "mirror.h"

#include <saas/util/queue.h>
#include <saas/util/network/address.h>
#include <library/cpp/logger/global/global.h>

#include <library/cpp/dns/cache.h>
#include <library/cpp/http/misc/httpreqdata.h>

#include <util/memory/blob.h>
#include <util/network/socket.h>
#include <util/random/random.h>

class TFlowMirror::TReq: public IObjectInQueue {
private:
    TFlowMirror& Parent;
    TString Request;
public:
    TReq(TFlowMirror& parent,
         const TString& request,
         const TBaseServerRequestData& rd,
         const TBlob& data)
         : Parent(parent)
    {
        Request += request + CrLf();
        if (Parent.GetConfig().CopyHeaders) {
            for (ui32 i = 0; i < rd.HeadersCount(); ++i) {
                const TString& header = rd.HeaderByIndex(i);
                if (header.StartsWith("Host:"))
                    continue;

                Request += header + CrLf();
            }

            Request += "Host: " + Parent.GetConfig().Host + ":" +
                ToString(Parent.GetConfig().Port) + CrLf();
        }
        Request += CrLf();
        Request.append(data.AsCharPtr(), data.Size());
    }

    const TString& GetRequestString() const {
        return Request;
    }

    void Process(void* /*ThreadSpecificResource*/) override {
        THolder<TReq> cleanup(this);
        Parent.DoProcessRequest(this);
    }

    static inline TStringBuf CrLf() {
        return "\r\n"sv;
    }
};

void TFlowMirror::DoProcessRequest(const TReq* req) {
    Y_ASSERT(req);
    Y_ASSERT(Address);
    try {
        if (StopFlag) {
            return;
        }
        TInstant start = Now();
        TSocket s(*Address, Config.Timeout);
        {
            TSocketOutput so(s);
            so << req->GetRequestString();
        }
        TInstant end = Now();
        Redirected.Hit((end - start).MilliSeconds());
    } catch (...) {
#if !defined(NDEBUG)
        Cerr << "Error while processing FlowMirror request: " << CurrentExceptionMessage() << Endl;
#endif
        Failed.Inc();
    }
}

bool TFlowMirror::SelectRequest(const TString& request) {
    for (auto&& slice : Config.Slices) {
        if (slice.CompiledRegexp.Match(request.c_str()) && RandomNumber<float>() <= slice.Fraction) {
            return true;
        }
    }
    if (RandomNumber<float>() <= Config.Fraction) {
        return true;
    }
    return false;
}

bool TFlowMirror::ProcessRequest(const TString& request, const TBaseServerRequestData& rd, const TBlob& data) {
    if (!Config.Enabled)
        return false;
    if (!SelectRequest(request)) {
        return false;
    }

    THolder<TReq> req(new TReq(*this, request, rd, data));
    if (!Queue->Add(req.Get())) {
#if !defined(NDEBUG)
        Cerr << "Cannot queue FlowMirror request" << Endl;
#endif
        Failed.Inc();
        return false;
    } else {
        Y_UNUSED(req.Release());
        return true;
    }
}

void TFlowMirror::RegisterMetrics(TMetrics& controller) {
    if (Config.Enabled) {
        Redirected.Register(controller);
        Failed.Register(controller);
    }
}

void TFlowMirror::UnregisterMetrics(TMetrics& controller) {
    if (Config.Enabled) {
        Redirected.Deregister(controller);
        Failed.Deregister(controller);
    }
}

const TFlowMirror::TConfig& TFlowMirror::GetConfig() const {
    return Config;
}

TFlowMirror::TFlowMirror(const TConfig& config)
    : Config(config)
    , Redirected("FlowMirror_Redirected")
    , Failed("FlowMirror_Failed")
    , StopFlag(false)
{
    if (!Config.Enabled) {
        return;
    }

    auto address = NUtil::Resolve(Config.Host, Config.Port);
    if (!address.Defined()) {
        ERROR_LOG << "Redirect to " << Config.Host << ":" << Config.Port << " has not been started" << Endl;
        return;
    }

    Address.Reset(new TNetworkAddress(address.GetRef()->Addr));
    {
        NUtil::TSmartMtpQueue::TOptions opts;
        opts.ThreadName = "FlowMirror";
        Queue.Reset(CreateRTYQueue(config.Threads, 0, opts));
    }

    INFO_LOG << "Redirecting " << Config.Fraction << " of stream to " << Config.Host << ":" << Config.Port << Endl;
}

TFlowMirror::~TFlowMirror() {
    StopFlag = true;
}
