/*****************************************************************************
* smsub Server Manager Subprocess
* Copyright (C) 2020-2024 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"
#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 << smsub_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

QStringList parseStringArguments(const QString &string)
{
    QString argument;
    bool slashMode = false;
    bool dQuoteMode = false;
    bool sQuoteMode = false;
    QStringList argumentList;
    for (const QChar &strChar : string) {
        if (!slashMode && !dQuoteMode && !sQuoteMode) {
            if (strChar == ' ') {
                if (!argument.isEmpty()) {
                    argumentList << argument;
                    argument.clear();
                }
            }
            else if (strChar == '\"') {
                dQuoteMode = true;
            }
            else if (strChar == '\'') {
                sQuoteMode = true;
            }
            else if (strChar == '\\') {
                slashMode = true;
            }
            else {
                argument += strChar;
            }
        }
        else if (slashMode) {
            argument += strChar;
            slashMode = false;
        }
        else if (dQuoteMode) {
            if (strChar == '\"') {
                dQuoteMode = false;
            }
            else {
                argument += strChar;
            }
        }
        else if (sQuoteMode) {
            if (strChar == '\'') {
                sQuoteMode = false;
            }
            else {
                argument += strChar;
            }
        }
    }
    if (slashMode || dQuoteMode || sQuoteMode)
        return QStringList();
    if (!argument.isEmpty())
        argumentList << argument;
    return argumentList;
}

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

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

    bool rportSet = false;
    bool autoStart = true;
    bool keepAlive = false;
    bool timeoutSet = false;
    int termTimeout = 60000;
    quint16 rport;
    QString socket;
    QString rsocket;
    QString executable;
    QString workingDirectory;
    QStringList argumentList;

    const QByteArray envEnvironmentMode = qgetenv("SMSUB_ENVIRONMENT_MODE");
    const QByteArray envAutoStart = qgetenv("SMSUB_AUTOSTART");
    const QByteArray envKeepAlive = qgetenv("SMSUB_KEEPALIVE");
    const QByteArray envManifest = qgetenv("SMSUB_JSON");
    const QByteArray envExecutable = qgetenv("SMSUB_EXEC");
    const QByteArray envArguments = qgetenv("SMSUB_ARGS");
    const QByteArray envSocket = qgetenv("SMSUB_SOCK");
    const QByteArray envRemotePort = qgetenv("SMSUB_RPORT");
    const QByteArray envRemoteSocket = qgetenv("SMSUB_RSOCK");
    const QByteArray envTimeout = qgetenv("SMSUB_TIMEOUT");
    const QByteArray envWorkDir = qgetenv("SMSUB_WORKDIR");

    if (envEnvironmentMode == "1" || envEnvironmentMode.toLower() == "true" || envEnvironmentMode.toLower() == "yes") {
        if (envExecutable.isEmpty() && envArguments.isEmpty()) {
            QStringList arguments = a.arguments();
            arguments.removeFirst();
            executable = arguments.takeFirst();
            argumentList = arguments;
        }
        else {
            if (!envExecutable.isEmpty()) {
                executable = QString::fromUtf8(envExecutable);
            }
            else {
                QTextStream(stderr) << "Executable is not defined in environment, aborting!" << smsub_endl;
                return 1;
            }

            if (!envArguments.isEmpty()) {
                argumentList = parseStringArguments(QString::fromUtf8(envArguments));
                if (argumentList.empty()) {
                    QTextStream(stderr) << "Arguments can't be parsed properly!" << smsub_endl;
                    return 1;
                }
            }
            else {
                QTextStream(stderr) << "Arguments are not defined in environment, aborting!" << smsub_endl;
                return 1;
            }
        }

        if (!envTimeout.isEmpty()) {
            bool ok;
            const int _termTimeout = envTimeout.toInt(&ok);
            if (ok) {
                termTimeout = _termTimeout;
                timeoutSet = true;
            }
            else {
                QTextStream(stderr) << "Termination timeout is not a number in environment, aborting!" << smsub_endl;
                return 1;
            }
        }

        if (!envAutoStart.isEmpty()) {
            if (envAutoStart == "0" || envAutoStart.toLower() == "false" || envAutoStart.toLower() == "no") {
                autoStart = false;
            }
        }

        if (!envKeepAlive.isEmpty()) {
            if (envKeepAlive == "1" || envKeepAlive.toLower() == "true" || envKeepAlive.toLower() == "yes") {
                keepAlive = true;
            }
        }

        if (!envWorkDir.isEmpty()) {
            workingDirectory = QString::fromUtf8(envWorkDir);
        }
        else {
            workingDirectory = QFileInfo(executable).absolutePath();
        }

        if (!envSocket.isEmpty()) {
            socket = QString::fromUtf8(envSocket);
        }
        else {
#ifdef Q_OS_WIN
            QTextStream(stderr) << "You must define at least a local IPC socket!" << smsub_endl;
#else
            QTextStream(stderr) << "You must define at least a local Unix socket!" << smsub_endl;
#endif
            return 1;
        }

        if (!envRemoteSocket.isEmpty()) {
            rsocket = QString::fromUtf8(envRemoteSocket);
        }

        if (!envRemotePort.isEmpty()) {
            bool ok;
            const quint16 _rport = envRemotePort.toUShort(&ok);
            if (ok) {
                rport = _rport;
                rportSet = true;
            }
            else {
                QTextStream(stderr) << "WebSockets port is not valid in environment!" << smsub_endl;
                return 1;
            }
        }
    }
    else {
        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 (commandLineParser.isSet(processManifest) && commandLineParser.isSet(processExecutable) ||
                     !envManifest.isEmpty() && !envExecutable.isEmpty() ||
                     commandLineParser.isSet(processManifest) && !envExecutable.isEmpty() ||
                     !envManifest.isEmpty() && commandLineParser.isSet(processExecutable)) {
            QTextStream(stderr) << "You can't define a Process executable and a JSON process manifest at the same time!" << smsub_endl;
            return 1;
        }

        if (commandLineParser.isSet(processManifest) && commandLineParser.isSet(processArguments) ||
                     !envManifest.isEmpty() && !envArguments.isEmpty() ||
                     commandLineParser.isSet(processManifest) && !envArguments.isEmpty() ||
                     !envManifest.isEmpty() && commandLineParser.isSet(processArguments)) {
            QTextStream(stderr) << "You can't define a Process arguments and a JSON process manifest at the same time!" << smsub_endl;
            return 1;
        }

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

        if (commandLineParser.isSet(processTimeout)) {
            bool ok;
            const int _termTimeout = commandLineParser.value(processTimeout).toInt(&ok);
            if (ok) {
                termTimeout = _termTimeout;
                timeoutSet = true;
            }
            else {
                QTextStream(stderr) << "Termination timeout is not a number in argument, aborting!" << smsub_endl;
                return 1;
            }
        }
        else if (!envTimeout.isEmpty()) {
            bool ok;
            const int _termTimeout = envTimeout.toInt(&ok);
            if (ok) {
                termTimeout = _termTimeout;
                timeoutSet = true;
            }
            else {
                QTextStream(stderr) << "Termination timeout is not a number in environment, aborting!" << smsub_endl;
                return 1;
            }
        }

        if (commandLineParser.isSet(processExecutable)) {
            executable = commandLineParser.value(processExecutable);
        }
        else if (!envExecutable.isEmpty()) {
            executable = QString::fromUtf8(envExecutable);
        }

        if (!envAutoStart.isEmpty()) {
            if (envAutoStart == "0" || envAutoStart.toLower() == "false" || envAutoStart.toLower() == "no") {
                autoStart = false;
            }
        }

        if (!envKeepAlive.isEmpty()) {
            if (envKeepAlive == "1" || envKeepAlive.toLower() == "true" || envKeepAlive.toLower() == "yes") {
                keepAlive = true;
            }
        }

        if (!envWorkDir.isEmpty()) {
            workingDirectory = QString::fromUtf8(envWorkDir);
        }
        else {
            workingDirectory = QFileInfo(executable).absolutePath();
        }

        QString manifestPath;
        if (commandLineParser.isSet(processManifest)) {
            manifestPath = commandLineParser.value(processManifest);
        }
        else if (!envManifest.isEmpty()) {
            manifestPath = QString::fromUtf8(envManifest);
        }

        if (!manifestPath.isEmpty()) {
            QFile manifestFile(manifestPath);
            if (manifestFile.open(QIODevice::ReadOnly)) {
                const QByteArray jsonData = manifestFile.readAll();
                QJsonDocument jsonDocument = QJsonDocument::fromJson(jsonData);
                QJsonObject jsonObject = jsonDocument.object();

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

                if (jsonObject.contains("AutoStart")) {
                    const QJsonValue jsonAutoStart = jsonObject.value("AutoStart");
                    if (!jsonAutoStart.isBool()) {
                        QTextStream(stderr) << "AutoStart is not a bool in manifest, aborting!" << smsub_endl;
                        manifestFile.close();
                        return 1;
                    }
                    autoStart = jsonAutoStart.toBool();
                }

                if (jsonObject.contains("KeepAlive")) {
                    const QJsonValue jsonKeepAlive = jsonObject.value("KeepAlive");
                    if (!jsonKeepAlive.isBool()) {
                        QTextStream(stderr) << "KeepAlive is not a bool in manifest, aborting!" << smsub_endl;
                        manifestFile.close();
                        return 1;
                    }
                    keepAlive = jsonKeepAlive.toBool();
                }

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

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

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

                manifestFile.close();
            }
        }
        else if (commandLineParser.isSet(processArguments)) {
            argumentList = parseStringArguments(commandLineParser.value(processArguments));
            if (argumentList.empty()) {
                QTextStream(stderr) << "Arguments can't be parsed properly!" << smsub_endl;
                return 1;
            }
        }
        else if (!envArguments.isEmpty()) {
            argumentList = parseStringArguments(QString::fromUtf8(envArguments));
            if (argumentList.empty()) {
                QTextStream(stderr) << "Arguments can't be parsed properly!" << smsub_endl;
                return 1;
            }
        }

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

        if (commandLineParser.isSet(subprocessRemoteSocket)) {
            rsocket = commandLineParser.value(subprocessRemoteSocket);
        }
        else if (!envRemoteSocket.isEmpty()) {
            rsocket = QString::fromUtf8(envRemoteSocket);
        }

        if (commandLineParser.isSet(subprocessRemotePort)) {
            bool ok;
            const quint16 _rport = commandLineParser.value(subprocessRemotePort).toUShort(&ok);
            if (ok) {
                rport = _rport;
                rportSet = true;
            }
            else {
                QTextStream(stderr) << "WebSockets port is not valid in arguments!" << smsub_endl;
                return 1;
            }
        }
        else if (!envRemotePort.isEmpty()) {
            bool ok;
            const quint16 _rport = envRemotePort.toUShort(&ok);
            if (ok) {
                rport = _rport;
                rportSet = true;
            }
            else {
                QTextStream(stderr) << "WebSockets port is not valid in arguments!" << smsub_endl;
                return 1;
            }
        }
    }

    SMSubServerSettings localSettings;
    localSettings.isLocal = true;

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

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

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

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

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

    if (autoStart)
        subProcess.startProcess();

    QTextStream(stderr) << QString("SMSub Version %1 initialized!").arg(QCoreApplication::applicationVersion()) << smsub_endl;

    return a.exec();
}