/*****************************************************************************
* smsub Server Manager Subprocess
* Copyright (C) 2020 Syping
*
* Redistribution and use in source and binary forms, with or without modification,
* are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* This software is provided as-is, no warranties are given to you, we are not
* responsible for anything with use of the software, you are self responsible.
*****************************************************************************/

#include <QCommandLineParser>
#include <QCommandLineOption>
#include <QCoreApplication>
#include <QJsonDocument>
#include <QJsonObject>
#include <QTextStream>
#include <QJsonValue>
#include <QJsonArray>
#include <QFileInfo>
#include <QFile>
#include "SMSubProcess.h"
#include "SMSubServer.h"
#include "smsub.h"

#ifdef Q_OS_UNIX
#include <initializer_list>
#include "signal.h"
#include "unistd.h"
#endif

#ifdef Q_OS_UNIX
void catchUnixSignals(std::initializer_list<int> quitSignals) {
    auto handler = [](int sig) -> void {
        QString unixSignal;
        switch (sig) {
        case SIGINT:
            unixSignal = QLatin1String("SIGINT");
            break;
        case SIGHUP:
            unixSignal = QLatin1String("SIGHUP");
            break;
        case SIGQUIT:
            unixSignal = QLatin1String("SIGQUIT");
            break;
        case SIGTERM:
            unixSignal = QLatin1String("SIGTERM");
            break;
        default:
            unixSignal = QString::number(sig);
        }
        QTextStream(stderr) << "Received Unix signal: " << unixSignal << endl;
        QCoreApplication::quit();
    };

    sigset_t blocking_mask;
    sigemptyset(&blocking_mask);
    for (int sig : quitSignals)
        sigaddset(&blocking_mask, sig);

    struct sigaction sa;
    sa.sa_handler = handler;
    sa.sa_mask = blocking_mask;
    sa.sa_flags = 0;

    for (int sig : quitSignals)
        sigaction(sig, &sa, nullptr);
}
#endif

int main(int argc, char *argv[])
{
    QCoreApplication a(argc, argv);
    a.setApplicationName("Server Manager Subprocess");
    a.setApplicationVersion("0.3.1");

#ifdef Q_OS_UNIX
    catchUnixSignals({SIGINT, SIGHUP, SIGQUIT, SIGTERM});
#endif

    QCommandLineParser commandLineParser;
    commandLineParser.addHelpOption();
    commandLineParser.addVersionOption();

    QCommandLineOption processManifest("json", "JSON process manifest.", "json");
    commandLineParser.addOption(processManifest);

    QCommandLineOption processExecutable(QStringList() << "exec" << "executable", "Process executable to run.", "exec");
    commandLineParser.addOption(processExecutable);

    QCommandLineOption processArguments(QStringList() << "args" << "arguments", "Arguments given to process.", "args");
    commandLineParser.addOption(processArguments);

#ifdef Q_OS_WIN
    QCommandLineOption subprocessSocket(QStringList() << "sock" << "socket", "IPC socket used for local communication.", "sock");
#else
    QCommandLineOption subprocessSocket(QStringList() << "sock" << "socket", "Unix socket used for local communication.", "sock");
#endif
    commandLineParser.addOption(subprocessSocket);

    QCommandLineOption subprocessRemotePort("rport", "WebSockets port used for remote communication.", "rport");
    commandLineParser.addOption(subprocessRemotePort);

#ifdef Q_OS_WIN
    QCommandLineOption subprocessRemoteSocket(QStringList() << "rsock" << "rsocket", "IPC socket used for remote communication.", "rsock");
#else
    QCommandLineOption subprocessRemoteSocket(QStringList() << "rsock" << "rsocket", "Unix socket used for remote communication.", "rsock");
#endif
    commandLineParser.addOption(subprocessRemoteSocket);

    QCommandLineOption processTimeout("timeout", "SMSub termination timeout.", "timeout");
    commandLineParser.addOption(processTimeout);

    commandLineParser.process(a);

    if (unlikely(commandLineParser.isSet(processManifest) && commandLineParser.isSet(processExecutable))) {
        QTextStream(stderr) << "You can't define a Process executable and a JSON process manifest at the same time!" << endl;
        return 1;
    }

    if (unlikely(commandLineParser.isSet(subprocessRemotePort) && commandLineParser.isSet(subprocessRemoteSocket))) {
#ifdef Q_OS_WIN
        QTextStream(stderr) << "You can't define a WebSockets port and a IPC socket at same time!" << endl;
#else
        QTextStream(stderr) << "You can't define a WebSockets port and a Unix socket at same time!" << endl;
#endif
        return 1;
    }

    bool timeoutSet = false;
    int termTimeout = 60000;
    if (unlikely(commandLineParser.isSet(processTimeout))) {
        bool ok;
        const int _termTimeout = commandLineParser.value(processTimeout).toInt(&ok);
        if (ok) {
            termTimeout = _termTimeout;
            timeoutSet = true;
        }
    }

    QString executable;
    QString workingDirectory;
    QStringList argumentList;
    if (likely(commandLineParser.isSet(processManifest))) {
        QFile manifestFile(commandLineParser.value(processManifest));
        if (likely(manifestFile.open(QIODevice::ReadOnly))) {
            const QByteArray jsonData = manifestFile.readAll();
            QJsonDocument jsonDocument = QJsonDocument::fromJson(jsonData);
            QJsonObject jsonObject = jsonDocument.object();

            if (likely(jsonObject.contains("Executable"))) {
                const QJsonValue jsonExecutable = jsonObject.value("Executable");
                if (unlikely(!jsonExecutable.isString())) {
                    QTextStream(stderr) << "Executable is not a string in manifest, aborting!" << endl;
                    manifestFile.close();
                    return 1;
                }
                executable = jsonExecutable.toString();
            }
            else {
                QTextStream(stderr) << "Executable is not defined in manifest, aborting!" << endl;
                manifestFile.close();
                return 1;
            }

            if (likely(jsonObject.contains("WorkingDirectory"))) {
                const QJsonValue jsonWorkingDirectory = jsonObject.value("WorkingDirectory");
                if (unlikely(!jsonWorkingDirectory.isString())) {
                    QTextStream(stderr) << "Working Directory is not a string in manifest, aborting!" << endl;
                    manifestFile.close();
                    return 1;
                }
                workingDirectory = jsonWorkingDirectory.toString();
            }
            else {
                workingDirectory = QFileInfo(executable).absolutePath();
            }

            if (likely(jsonObject.contains("Arguments"))) {
                const QJsonValue jsonArguments = jsonObject.value("Arguments");
                if (likely(jsonArguments.isArray())) {
                    const QJsonArray jsonArray = jsonArguments.toArray();
                    QJsonArray::const_iterator it = jsonArray.constBegin();
                    QJsonArray::const_iterator end = jsonArray.constEnd();
                    while (it != end) {
                        argumentList << it->toString();
                        it++;
                    }
                }
                else {
                    QTextStream(stderr) << "Arguments is not a array in manifest, aborting!" << endl;
                    manifestFile.close();
                    return 1;
                }
            }

            if (unlikely(!timeoutSet && jsonObject.contains("TerminationTimeout"))) {
                const QJsonValue jsonTimeout = jsonObject.value("TerminationTimeout");
                if (unlikely(!jsonTimeout.isDouble())) {
                    termTimeout = qRound(jsonTimeout.toDouble());
                }
                else {
                    QTextStream(stderr) << "Termination Timeout is not a number in manifest, aborting!" << endl;
                    return 1;
                }
            }

            manifestFile.close();
        }
    }
    else if (unlikely(commandLineParser.isSet(processArguments))) {
        QTextStream(stderr) << "Arguments over command line are not supported yet!" << endl;
        return 1;
    }

    QString socket;
    if (likely(commandLineParser.isSet(subprocessSocket))) {
        socket = commandLineParser.value(subprocessSocket);
    }
    else {
#ifdef Q_OS_WIN
        QTextStream(stderr) << "You must define at least a local IPC socket!" << endl;
#else
        QTextStream(stderr) << "You must define at least a local Unix socket!" << endl;
#endif
        return 1;
    }

    SMSubServerSettings localSettings;
    localSettings.isLocal = true;

    SMSubServer subLocal(&localSettings, socket);
    if (unlikely(!subLocal.isListening())) {
#ifdef Q_OS_WIN
        QTextStream(stderr) << "Failed to start local IPC socket!" << endl;
#else
        QTextStream(stderr) << "Failed to start local Unix socket!" << endl;
#endif
        return 1;
    }

    QString rsocket;
    if (unlikely(commandLineParser.isSet(subprocessRemoteSocket))) {
        rsocket = commandLineParser.value(subprocessRemoteSocket);
    }

    bool rportSet = false;
    quint16 rport;
    if (unlikely(commandLineParser.isSet(subprocessRemotePort))) {
        bool ok;
        rport = commandLineParser.value(subprocessRemotePort).toUShort(&ok);
        if (!ok) {
            QTextStream(stderr) << "WebSockets port is not valid!" << endl;
            return 1;
        }
        else {
            rportSet = true;
        }
    }

    SMSubProcess subProcess(executable, argumentList, workingDirectory, termTimeout);

    SMSubServerSettings remoteSettings;
    remoteSettings.canRegister = false;
    remoteSettings.isLocal = false;

    if (unlikely(!rsocket.isEmpty())) {
        SMSubServer *subRemote = new SMSubServer(&remoteSettings, rsocket);
        if (unlikely(!subRemote->isListening())) {
#ifdef Q_OS_WIN
            QTextStream(stderr) << "Failed to start remote IPC socket!" << endl;
#else
            QTextStream(stderr) << "Failed to start remote Unix socket!" << endl;
#endif
            return 1;
        }
        localSettings.canRegister = true;
        QObject::connect(&subLocal, &SMSubServer::tokenRegistered, subRemote, &SMSubServer::registerToken);
        QObject::connect(&subProcess, &SMSubProcess::outputWritten, subRemote, &SMSubServer::writeOutput);
        QObject::connect(subRemote, &SMSubServer::inputWritten, &subProcess, &SMSubProcess::writeInput);
        QObject::connect(subRemote, &SMSubServer::killRequested, &subProcess, &SMSubProcess::killProcess);
        QObject::connect(subRemote, &SMSubServer::stopRequested, &subProcess, &SMSubProcess::stopProcess);
    }
    else if (unlikely(rportSet)) {
        SMSubServer *subRemote = new SMSubServer(&remoteSettings, QString(), rport);
        if (unlikely(!subRemote->isListening())) {
            QTextStream(stderr) << "Failed to start remote WebSockets server!" << endl;
            return 1;
        }
        localSettings.canRegister = true;
        QObject::connect(&subLocal, &SMSubServer::tokenRegistered, subRemote, &SMSubServer::registerToken);
        QObject::connect(&subProcess, &SMSubProcess::outputWritten, subRemote, &SMSubServer::writeOutput);
        QObject::connect(subRemote, &SMSubServer::inputWritten, &subProcess, &SMSubProcess::writeInput);
        QObject::connect(subRemote, &SMSubServer::killRequested, &subProcess, &SMSubProcess::killProcess);
        QObject::connect(subRemote, &SMSubServer::stopRequested, &subProcess, &SMSubProcess::stopProcess);
    }
    else {
        localSettings.canRegister = false;
    }

    QObject::connect(&subProcess, &SMSubProcess::outputWritten, &subLocal, &SMSubServer::writeOutput);
    QObject::connect(&subLocal, &SMSubServer::inputWritten, &subProcess, &SMSubProcess::writeInput);
    QObject::connect(&subLocal, &SMSubServer::killRequested, &subProcess, &SMSubProcess::killProcess);
    QObject::connect(&subLocal, &SMSubServer::stopRequested, &subProcess, &SMSubProcess::stopProcess);
    QObject::connect(&a, &QCoreApplication::aboutToQuit, &subProcess, &SMSubProcess::aboutToQuit);

    subProcess.start();

    return a.exec();
}