#include "main.h"

#include <library/cpp/getopt/opt.h>
#include <util/folder/path.h>

using NLastGetopt::TOpt;
using NLastGetopt::TOpts;
using NLastGetopt::TOptsParseResult;

void RunMain(int argc, const char **argv) {
    NYT::Initialize(argc, argv);

    TYtOptions options;
    bool updateOnly;
    TUploadOptions uploadOptions;

    TOpts opts = TOpts::Default();

    opts.AddLongOption('c', "cluster", "cluster name")
        .StoreResult(&options.ytCluster)
        .RequiredArgument()
        .DefaultValue("hahn");
    opts.AddLongOption('i', "input", "input table")
        .StoreResult(&options.ytInputTable)
        .RequiredArgument();
    opts.AddLongOption('o', "output", "output table")
        .StoreResult(&options.ytOutputTable)
        .RequiredArgument();
    opts.AddLongOption('s', "secrets", "secrets file")
        .StoreResult(&options.ytSecretsFile)
        .RequiredArgument();
    opts.AddLongOption('h', "header", "custom header for packets table")
        .StoreResult(&options.ytHeaderFile)
        .RequiredArgument();;
    opts.AddLongOption("tshark", "path to tshark binary at YT cluster")
        .StoreResult(&options.tsharkPath)
        .RequiredArgument()
        .DefaultValue("//home/balancer/util/tshark");
    opts.AddLongOption("libs", "path to libs package at YT cluster")
        .StoreResult(&options.libsPath)
        .RequiredArgument()
        .DefaultValue("//home/balancer/util/libraries.tar.gz");
    opts.AddLongOption("lua-script", "path to lua script at YT cluster")
        .StoreResult(&options.luaPath)
        .RequiredArgument()
        .DefaultValue("//home/balancer/util/tshark.lua");
    opts.AddLongOption("http-headers", "path to custom HTTP headers at YT cluster")
        .StoreResult(&options.httpHeadersPath)
        .RequiredArgument()
        .DefaultValue("//home/balancer/util/custom_http_header_fields");
    opts.AddLongOption("http2-headers", "path to custom HTTP2 headers at YT cluster")
        .StoreResult(&options.http2HeadersPath)
        .RequiredArgument()
        .DefaultValue("//home/balancer/util/custom_http2_header_fields");
    opts.AddLongOption("memory-limit", "YT redice memory limit in MB")
        .StoreResult(&options.memoryLimit)
        .RequiredArgument()
        .DefaultValue(800);

    opts.AddLongOption("update-only", "only update resources")
        .SetFlag(&updateOnly).NoArgument();

    opts.AddLongOption("tshark-local", "update tshark binary")
        .StoreResult(&uploadOptions.tsharkLocalPath)
        .RequiredArgument();
    opts.AddLongOption("libs-local", "update libraries package")
        .StoreResult(&uploadOptions.libsLocalPath)
        .RequiredArgument();
    opts.AddLongOption("lua-local", "update lua script")
        .StoreResult(&uploadOptions.luaLocalPath)
        .RequiredArgument();
    opts.AddLongOption("http-headers-local", "update custom headers file")
        .StoreResult(&uploadOptions.httpHeadersLocalPath)
        .RequiredArgument();
    opts.AddLongOption("http2-headers-local", "update custom headers file for http2")
        .StoreResult(&uploadOptions.http2HeadersLocalPath)
        .RequiredArgument();
    opts.AddLongOption("secrets-local", "update secrets file")
        .StoreResult(&uploadOptions.secretsLocalPath)
        .RequiredArgument();

    TOptsParseResult res(&opts, argc, argv);

    if (options.ytCluster.empty()) {
        ythrow yexception() << "Please, specify cluster name.\nTry '"
             << argv[0] << "--help' for more information.\n";
    }

    if (!updateOnly) {
        if (options.ytInputTable.empty() || options.ytOutputTable.empty() || options.ytSecretsFile.empty()) {
            ythrow yexception() << "Please, specify input table with packets, output table and secrets file.\nTry '"
                 << argv[0] << "--help' for more information.\n";
        }

        if (options.ytHeaderFile.empty()) {
            options.ytHeaderFile = options.ytInputTable + "-header";
        }

        RunMainMap(TStringBuf(argv[0]), options);
    } else {
        if (uploadOptions.tsharkLocalPath.empty() && uploadOptions.libsLocalPath.empty()
            && uploadOptions.luaLocalPath.empty() && uploadOptions.httpHeadersLocalPath.empty()
            && uploadOptions.http2HeadersLocalPath.empty() && uploadOptions.secretsLocalPath.empty()) {

            ythrow yexception() <<
                "Please, specify path to tshark binary/library and/or path to lua script and/or path to headers files and/or secrets file.\nTry '"sv
                 << argv[0] << "--help' for more information.\n";
        }

        RunMainUpload(TStringBuf(argv[0]), options, uploadOptions);
    }
}

void RunMainUpload(TStringBuf argv0, const TYtOptions &options, const TUploadOptions &uploadOptions) {
    TFileOperations operations(options.ytCluster);

    if (!uploadOptions.tsharkLocalPath.empty() && !CheckFileExists(uploadOptions.tsharkLocalPath)) {
        ythrow yexception() << "Incorrect tshark binary file " << uploadOptions.tsharkLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.tsharkLocalPath.empty() && options.tsharkPath.empty()) {
        ythrow yexception() << "Empty tshark YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.tsharkLocalPath.empty()) {
        operations.UploadFile(options.tsharkPath, uploadOptions.tsharkLocalPath, true);
    }

    if (!uploadOptions.libsLocalPath.empty() && !CheckFileExists(uploadOptions.libsLocalPath)) {
        ythrow yexception() << "Incorrect libraries package file " << uploadOptions.libsLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.libsLocalPath.empty() && options.libsPath.empty()) {
        ythrow yexception() << "Empty libraries package YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.libsLocalPath.empty()) {
        operations.UploadFile(options.libsPath, uploadOptions.libsLocalPath, false);
    }

    if (!uploadOptions.luaLocalPath.empty() && !CheckFileExists(uploadOptions.luaLocalPath)) {
        ythrow yexception() << "Incorrect lua script file " << uploadOptions.luaLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.luaLocalPath.empty() && options.luaPath.empty()) {
        ythrow yexception() << "Empty lua script YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.luaLocalPath.empty()) {
        operations.UploadFile(options.luaPath, uploadOptions.luaLocalPath, false);
    }

    if (!uploadOptions.httpHeadersLocalPath.empty() && !CheckFileExists(uploadOptions.httpHeadersLocalPath)) {
        ythrow yexception() << "Incorrect http headers file " << uploadOptions.httpHeadersLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.httpHeadersLocalPath.empty() && options.httpHeadersPath.empty()) {
        ythrow yexception() << "Empty http headers YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.httpHeadersLocalPath.empty()) {
        operations.UploadFile(options.httpHeadersPath, uploadOptions.httpHeadersLocalPath, false);
    }

    if (!uploadOptions.http2HeadersLocalPath.empty() && !CheckFileExists(uploadOptions.http2HeadersLocalPath)) {
        ythrow yexception() << "Incorrect http2 headers file " << uploadOptions.http2HeadersLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.http2HeadersLocalPath.empty() && options.http2HeadersPath.empty()) {
        ythrow yexception() << "Empty http2 headers YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.http2HeadersLocalPath.empty()) {
        operations.UploadFile(options.http2HeadersPath, uploadOptions.http2HeadersLocalPath, false);
    }

    if (!uploadOptions.secretsLocalPath.empty() && !CheckFileExists(uploadOptions.secretsLocalPath)) {
        ythrow yexception() << "Incorrect secrets file " << uploadOptions.secretsLocalPath << ".\nTry '"
             << argv0 << "--help' for more information.\n";
    }

    if (!uploadOptions.secretsLocalPath.empty() && options.ytSecretsFile.empty()) {
        ythrow yexception() << "Empty secrets YT path.\nTry '"
             << argv0 << "--help' for more information.\n";
    } else if (!uploadOptions.secretsLocalPath.empty()) {
        operations.UploadFile(options.ytSecretsFile, uploadOptions.secretsLocalPath, false);
    }
}

int main(int argc, const char** argv) noexcept {
    try {
        RunMain(argc, argv);
    } catch(const yexception& e) {
        Cerr << e.what() << "\n";
        return 1;
    } catch(...) {
        return 1;
    }
}

bool CheckFileExists(const TString& path) {
    return TFsPath(path).IsFile();
}
