#include "dolbilka_plan.h"

#include <util/string/cast.h>
#include <util/string/strip.h>
#include <util/generic/vector.h>

namespace NSaas {

TStringBuf TDolbilkaPlan::TrimQuery(TStringBuf src) {
    if (src.StartsWith("http")) {
        src = src.After(':');
        for (int i = 0; i < 3; i++) {
            src = src.After('/');
        }
    }
    if (!src.StartsWith('&')) {
        src = src.After('?');
    }
    return src;
}

namespace {
    TRequestId ToRequestId(const TStringBuf& strRequestId) {
        TRequestId requestId = 0;
        Y_ENSURE(
            TryFromString(strRequestId, requestId),
            "Cannot convert request id " << TString(strRequestId).Quote() << " to integer. Please check input files format: 'RequestId \\t Request' ");
        return requestId;
    }

    template <typename Foo>
    void ForEachLine(IInputStream& inp, Foo foo) {
        for (TString line; inp.ReadLine(line);) {
            if (line.StartsWith("# "))
                continue;

            TStringBuf trimmed = StripStringLeft(line);
            if (trimmed.size() <= 1)
                continue;

            bool shouldStop = foo(trimmed);
            if (shouldStop) {
                break;
            }
        }
    }
}

void TDolbilkaPlanReader::ReadAll(IInputStream& inp, TVector<TDolbilkaPlan::TQueryData>& plan) {
    ui64 autoId = 0;
    plan.clear();

    const size_t limit = MaxItemsToRead ? MaxItemsToRead : Max<size_t>();

    if (Mode == EPlanFormat::PlainQueries) {
        ForEachLine(inp, [&](TStringBuf trimmed) {
            trimmed = TDolbilkaPlan::TrimQuery(trimmed);
            trimmed = StripStringRight(trimmed);
            plan.emplace_back(TDolbilkaPlan::TQueryData{++autoId, TString(trimmed)});
            return plan.size() >= limit;
        });
    } else {
        Y_ENSURE(Mode == EPlanFormat::TsvWithIds);
        ForEachLine(inp, [&](TStringBuf trimmed) {
            const TStringBuf requestIdBuf = trimmed.NextTok('\t');
            const TRequestId requestId = ToRequestId(requestIdBuf);

            trimmed = TDolbilkaPlan::TrimQuery(trimmed);
            trimmed = StripStringRight(trimmed);
            if (!trimmed.empty()) {
                plan.emplace_back(TDolbilkaPlan::TQueryData{requestId, "/?" + TString(trimmed)});
            }
            return plan.size() >= limit;
        });
    }
}

}
