// Copyright 2021 Bloomberg Finance L.P
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <csignal>
#include <cstddef>
#include <errno.h>
#include <fcntl.h>
#include <iostream>
#include <limits>
#include <netinet/in.h>
#include <sys/socket.h>

#include <buildboxcommon_commandline.h>
#include <buildboxcommon_logging.h>

#include <rumbad_bufferedpublisher.h>
#include <rumbad_compositewriter.h>
#include <rumbad_filewriter.h>
#include <rumbad_statsmanager.h>
#include <rumbad_stdoutwriter.h>
#include <span>

using namespace rumbad;

const int MSG_SIZE = 128 * 1024;
const int BATCH_SIZE = 200;
const std::chrono::seconds DEFAULT_PUBLISH_INTERVAL(60);
bool RECEIVED_INTERRUPT = false;
int PIPE_FDS[2];

enum ReturnCode : std::uint8_t {
    RC_OK = 0,
    RC_USAGE = 100,
    RC_SOCKET_CREATION_ERROR = 101,
    RC_SOCKET_BIND_ERROR = 102,
    RC_SOCKET_BUFFER_SIZE_ERROR = 103,
    RC_PORT_ERROR = 104,
    RC_SOCKET_FLAGS_ERROR = 105,
    RC_PIPE_ERROR = 106,
    RC_SIGNAL_HANDLER_ERROR = 107
};

using ArgumentSpec = buildboxcommon::CommandLineTypes::ArgumentSpec;
using DataType = buildboxcommon::CommandLineTypes::DataType;
using TypeInfo = buildboxcommon::CommandLineTypes::TypeInfo;
using DefaultValue = buildboxcommon::CommandLineTypes::DefaultValue;
using CommandLine = buildboxcommon::CommandLine;

void receiveData(int socketFd, BufferedPublisher *publisher, char *buffer)
{
    ssize_t nBytes = 0;
    fd_set fds;

    FD_ZERO(&fds);
    FD_SET(socketFd, &fds);
    FD_SET(PIPE_FDS[0], &fds);
    int maxFd = std::max(socketFd, PIPE_FDS[0]);

    if (select(maxFd + 1, &fds, NULL, NULL, NULL) > 0) {
        // Only read from the socket if there's something to read. Otherwise
        // we got here thanks to an interrupt coming down the interrupt pipe,
        // and don't need to do anything.
        if (FD_ISSET(socketFd, &fds)) {
            memset(buffer, 0, MSG_SIZE);

            // Don't bother storing the client address, we don't need to
            // respond.
            nBytes = recvfrom(socketFd, buffer, MSG_SIZE, 0, NULL, NULL);

            if (nBytes == 0) {
                BUILDBOX_LOG_DEBUG("Received an empty message from the client")
            }
            else if (nBytes < 0) {
                BUILDBOX_LOG_ERROR("Error receiving message: rcode="
                                   << nBytes << ", errno=" << errno);
            }
            else {
                std::string strMessage(buffer, nBytes);
                publisher->publish(strMessage);
            }
        }
    }
}

void setupLogger(const char *programName)
{
    buildboxcommon::LogLevel logLevel = buildboxcommon::LogLevel::DEBUG;
    BUILDBOX_LOG_SET_LEVEL(logLevel);

    auto &logger = buildboxcommon::logging::Logger::getLoggerInstance();
    logger.initialize(programName);
}

void handleSignal(int signum)
{
    std::cout << "Received signal [" << signum << "]\n";
    RECEIVED_INTERRUPT = true;

    char signal[3];
    unsigned long n = snprintf(signal, sizeof(signal), "%d", signum);
    if (n > 0 && n < sizeof(signal)) {
        // If the signal can't be delivered then the pipe is full, so
        // just drop it.
        ssize_t ret = write(PIPE_FDS[1], signal, n + 1);
        (void)ret;
    }
}

int setupSignals()
{
    struct sigaction sa{};
    sa.sa_handler = handleSignal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_SIGINFO;

    if (sigaction(SIGINT, &sa, nullptr) == -1) {
        std::cout << "Unable to register signal handler for SIGINT\n";
        return RC_SIGNAL_HANDLER_ERROR;
    }
    if (sigaction(SIGTERM, &sa, nullptr) == -1) {
        std::cout << "Unable to register signal handler for SIGTERM\n";
        return RC_SIGNAL_HANDLER_ERROR;
    }
    if (sigaction(SIGHUP, &sa, nullptr) == -1) {
        std::cout << "Unable to register signal handler for SIGHUP\n";
        return RC_SIGNAL_HANDLER_ERROR;
    }
    return RC_OK;
}

int main(int argc, char *argv[])
{
    try {
        auto args = std::span(argv, argc);

        const std::string additionalHelpText(
            "rumbad is a daemon for the recc compiler wrapper. It listens "
            "to compilation data protobuf messages sent to a UDP port, and "
            "handles writing them to some persistent storage in batches.\n"
            "Environment Variables:\n"
            "   RUMBAD_RECV_BUFFER_SIZE: Used to set the value of SO_RCVBUF. "
            "Will be ignored if set to a non-integer value.\n"
            "   RUMBAD_PUBLISH_INTERVAL: Number of seconds to wait between "
            "writing batches of messages to the persistent storage.\n"
            "   RUMBAD_USE_STDOUT: If set to non-null, then rumbad will "
            "send records to stdout\n");

        int port = 0;
        std::string stats_filepath = "";
        ArgumentSpec spec[] = {
            {"log-file-path", "File path to store the log.",
             TypeInfo(DataType::COMMANDLINE_DT_STRING),
             ArgumentSpec::O_OPTIONAL, ArgumentSpec::C_WITH_ARG},
            {"", "port", TypeInfo(&port), ArgumentSpec::O_REQUIRED},
            {"", "stats_filepath", TypeInfo(&stats_filepath),
             ArgumentSpec::O_OPTIONAL}};
        CommandLine commandLine(spec, additionalHelpText);
        const bool success = commandLine.parse(argc, argv);

        if (!success) {
            commandLine.usage();
            return 1;
        }

        if (commandLine.exists("help") || commandLine.exists("version")) {
            return 0;
        }

        if (port > std::numeric_limits<uint16_t>::max() ||
            port < std::numeric_limits<uint16_t>::min()) {
            BUILDBOX_LOG_ERROR(
                "Port number out of range("
                << port << "), use a positive integer lower than "
                << std::numeric_limits<uint16_t>::max() << " and greater than "
                << std::numeric_limits<uint16_t>::min());
            commandLine.usage();
            return RC_PORT_ERROR;
        }

        if (pipe(PIPE_FDS) < 0) {
            std::cout << "Error creating pipe for signal handling, errno="
                      << errno << '\n';
            return RC_PIPE_ERROR;
        }

        int rcode = setupSignals();
        if (rcode != 0) {
            return rcode;
        }

        setupLogger(args[0]);

        BUILDBOX_LOG_INFO("Starting up rumba daemon on port [" << port << "]");

        int socketFd = socket(AF_INET, SOCK_DGRAM, 0);
        if (socketFd < 0) {
            BUILDBOX_LOG_ERROR("Error creating socket: rcode="
                               << socketFd << ", errno=" << errno);
            return RC_SOCKET_CREATION_ERROR;
        }

        // Set the socket to be non-blocking, since we're using select to read
        // from it
        int socketFlags = fcntl(socketFd, F_GETFL);
        if (socketFlags < 0) {
            BUILDBOX_LOG_ERROR("Error getting socket flags: rcode="
                               << socketFlags << ", errno=" << errno);
            return RC_SOCKET_FLAGS_ERROR;
        }

        rcode = fcntl(socketFd, F_SETFL, socketFlags | O_NONBLOCK);
        if (rcode < 0) {
            BUILDBOX_LOG_ERROR("Error setting socket flags: rcode="
                               << rcode << ", errno=" << errno);
            return RC_SOCKET_FLAGS_ERROR;
        }

        char *configuredBufSize = getenv("RUMBAD_RECV_BUFFER_SIZE");
        if (configuredBufSize != NULL) {
            try {
                int bufferSize = std::stoi(std::string(configuredBufSize));
                rcode = setsockopt(socketFd, SOL_SOCKET, SO_RCVBUF,
                                   &bufferSize, sizeof(int));
                if (rcode < 0) {
                    BUILDBOX_LOG_ERROR(
                        "Error setting socket buffer size: rcode="
                        << rcode << ", errno=" << errno);
                    return RC_SOCKET_BUFFER_SIZE_ERROR;
                }
            }
            catch (std::invalid_argument const &ex) {
                BUILDBOX_LOG_ERROR("Invalid socket buffer size ["
                                   << configuredBufSize << "], " << ex.what());
                return RC_SOCKET_BUFFER_SIZE_ERROR;
            }
            catch (std::out_of_range const &ex) {
                BUILDBOX_LOG_ERROR("Specified socket buffer size ["
                                   << configuredBufSize << "] out of range");
                return RC_SOCKET_BUFFER_SIZE_ERROR;
            }
        }

        struct sockaddr_in socketAddr{};
        memset(&socketAddr, 0, sizeof(socketAddr));
        socketAddr.sin_family = AF_INET;
        socketAddr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        socketAddr.sin_port = htons(port);

        // NOLINTBEGIN (cppcoreguidelines-pro-type-reinterpret-cast)
        rcode = bind(socketFd,
                     reinterpret_cast<const struct sockaddr *>(&socketAddr),
                     sizeof(socketAddr));
        // NOLINTEND (cppcoreguidelines-pro-type-reinterpret-cast)

        if (rcode < 0) {
            BUILDBOX_LOG_ERROR("Error binding socket: rcode="
                               << rcode << ", errno=" << errno);
            return RC_SOCKET_BIND_ERROR;
        }

        char *interval_str = getenv("RUMBAD_PUBLISH_INTERVAL");
        std::chrono::seconds publish_interval = DEFAULT_PUBLISH_INTERVAL;
        if (interval_str != NULL) {
            try {
                int interval_sec = std::stoi(std::string(interval_str));
                publish_interval = std::chrono::seconds(interval_sec);
            }
            catch (std::invalid_argument const &ex) {
                BUILDBOX_LOG_WARNING("Invalid RUMBAD_PUBLISH_INTERVAL, using "
                                     << DEFAULT_PUBLISH_INTERVAL.count()
                                     << " instead.");
            }
            catch (std::out_of_range const &ex) {
                BUILDBOX_LOG_WARNING(
                    "RUMBAD_PUBLISH_INTERVAL value out of range. Using "
                    << DEFAULT_PUBLISH_INTERVAL.count()
                    << " instead. Value must be between 0 and " << INT_MAX
                    << ".");
            }
        }

        // Configure a list of writers we want to use.
        std::vector<std::unique_ptr<Writer>> writersToUse;

        // 1. StdoutWriter - for sure
        writersToUse.emplace_back(std::make_unique<StdoutWriter>());
        BUILDBOX_LOG_INFO("Writing received messages to stdout.");

        // 2. FileWriter - if the logfile path is provided
        if (commandLine.exists("log-file-path")) {
            std::string logFilepath = commandLine.getString("log-file-path");
            writersToUse.emplace_back(
                std::make_unique<FileWriter>(logFilepath));
            BUILDBOX_LOG_INFO("Writing received messages to logfile "
                              << logFilepath);
        }

        // Configure publisher
        char *buffer = new char[MSG_SIZE];
        std::unique_ptr<Writer> writer =
            std::make_unique<CompositeWriter>(std::move(writersToUse));
        BufferedPublisher publisher(BATCH_SIZE, publish_interval,
                                    writer.get());

        // Configure writing stats to file if needed
        StatsManager *statsManager = &publisher.statsManager;
        if (stats_filepath.size() > 0) {
            BUILDBOX_LOG_INFO(
                "Found stats_filepath option. Stats will save to: " +
                stats_filepath);
            statsManager->setStatsFilepath(stats_filepath);
        }

        // Main listening loop
        while (!RECEIVED_INTERRUPT) {
            receiveData(socketFd, &publisher, buffer);
        }
        delete[] buffer;

        close(socketFd);
        return RC_OK;
    }
    catch (const std::exception &exception) {
        // This catch should not be relied on and is here as a backstop.
        BUILDBOX_LOG_ERROR("Caught exception in main " << exception.what());
        return 1;
    }
}
