#include "srcrwr_helper.h"

#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/net/addr.h>
#include <balancer/kernel/srcrwr/srcrwr_addrs.h>
#include <library/cpp/cgiparam/cgiparam.h>
#include <library/cpp/regex/pire/regexp.h>

#include <optional>


namespace {

const TString SrcrwrParam("l7rwr");
const TString SrcrwrHeader("X-Yandex-Internal-Srcrwr");
constexpr char SrcrwrDelimiter = ':';
constexpr char HostDelimiter = ';';
constexpr char HeaderDelim = ',';

}

namespace NSrvKernel::NSrcrwr {

bool THelper::WriteHeader(TRequest& request, const TString& id) const
{
    auto headerValue = request.Headers().GetFirstValue(SrcrwrHeader);
    if (!headerValue) { // first srcrwr record in header
        request.Headers().Add(SrcrwrHeader, id);
        return true;
    }

    for (const auto& i : StringSplitter(headerValue).Split(HeaderDelim)) {
        if (id == i.Token()) {
            return false; // this srcrwr ID was handled already
        }
    }

    request.Headers().Delete(SrcrwrHeader);
    request.Headers().Add(SrcrwrHeader, id + HeaderDelim + headerValue.data());

    return true;
}

bool THelper::GetAddrs(const TRequest& request, const TString& id, TVector<TAddr>& addrs) const noexcept
{
    try {
        return GetSrcAddresses(request, id, addrs);
    } catch (...) {
    }

    return false;
}

TString THelper::GetParamFromCgi(const TRequest& request, const TString& id) const
{
    if (request.RequestLine().CGI.Empty()) {
        return {};
    }

    const TCgiParameters cgiParams{request.RequestLine().CGI.AsStringBuf()};
    for (const auto& val: cgiParams.Range(SrcrwrParam)) { // search for "l7rwr=srcrwr_id:IP:port,hostname:port,..." params
        if (CheckId(id, val)) {
            return val; // only one srcrwr_id should be set in cgi
        }
    }
    TString param = cgiParams.Get(TString("?") + SrcrwrParam); // search for "?l7rwr=srcrwr_id:IP:port,hostname:port" param
    return CheckId(id, param) ? param : TString{};
}

bool THelper::GetSrcAddresses(const TRequest& request, const TString& id, TVector<TAddr>& addrs) const
{
    Y_ASSERT(addrs.empty());

    auto srcrwrParamFromCgi = GetParamFromCgi(request, id); // srcrwr_id:IP:port,hostname:port,...
    if (srcrwrParamFromCgi.Empty()) {
        return false;
    }

    size_t nextPos = 0;
    auto firstPos = srcrwrParamFromCgi.find(SrcrwrDelimiter);
    do {
        nextPos = srcrwrParamFromCgi.find(HostDelimiter, nextPos + 1);
        const auto lastPos  = srcrwrParamFromCgi.rfind(SrcrwrDelimiter, nextPos);
        if (firstPos == TString::npos || lastPos == TString::npos || firstPos + 1 >= lastPos) {
            return false;
        }

        auto host = srcrwrParamFromCgi.substr(firstPos + 1, lastPos - firstPos - 1);
        const auto hostLength = host.length();
        if (hostLength > 2 && host[0] == '[' && host[hostLength - 1] == ']') {
            host = host.substr(1, hostLength - 2);
        }
        ui16 port = FromString<TIpPort>(srcrwrParamFromCgi.substr(lastPos + 1, nextPos - lastPos - 1));

        auto errorOrResult = TSockAddr::FromIpPort(host, port);
        if (errorOrResult && host.find(SrcrwrDelimiter) != TString::npos) {
            return false;
        }
        addrs.emplace_back(std::move(host), "", port);
        firstPos = nextPos;
    } while(nextPos != TString::npos);

    return true;
}

bool THelper::CheckId(const TString& id, const TString& cgiParam) const
{
    const auto delimPos = cgiParam.find(SrcrwrDelimiter);

    if (delimPos == TString::npos || delimPos < 1) {
        return {};
    }

    return cgiParam.substr(0, delimPos) == id;
}

}
