// Copyright (C) 2022 The Qt Company Ltd.
// SPDX-License-Identifier: LicenseRef-Qt-Commercial OR GPL-3.0-only WITH Qt-GPL-exception-1.0

#include "deviceshell.h"

#include "qtcprocess.h"
#include "processinterface.h"
#include "qtcassert.h"
#include "utilstr.h"

#include <QLoggingCategory>

Q_LOGGING_CATEGORY(deviceShellLog, "qtc.utils.deviceshell", QtWarningMsg)

using namespace std::chrono_literals;

namespace Utils {

/*
 * The multiplex script waits for input via stdin.
 *
 * To start a command, a message is sent with
 * the format "<cmd-id> "<base64-encoded-stdin-data>" <commandline>\n"
 * To stop the script, simply send "exit\n" via stdin
 *
 * Once a message is received, two new streams are created that the new process redirects
 * its output to ( $stdoutraw and $stderrraw ).
 *
 * These streams are piped through base64 into the two streams stdoutenc and stderrenc.
 *
 * Two subshells read from these base64 encoded streams, and prepend the command-id,
 * as well as either "O:" or "E:" depending on whether its the stdout or stderr stream.
 *
 * Once the process exits its exit code is send to stdout with the command-id and the type "R".
 *
 */

DeviceShell::DeviceShell(bool forceFailScriptInstallation)
: m_forceFailScriptInstallation(forceFailScriptInstallation)
{
    m_thread.setObjectName("DeviceShell");
    m_thread.start();
}

DeviceShell::~DeviceShell()
{
    if (m_thread.isRunning()) {
        m_thread.quit();
        m_thread.wait();
    }

    QTC_CHECK(!m_shellProcess);
}

/*!
 * \brief DeviceShell::runInShell
 * \param cmd The command to run
 * \param stdInData Data to send to the stdin of the command
 * \return true if the command finished with EXIT_SUCCESS(0)
 *
 * Runs the cmd inside the internal shell process and return stdout, stderr and exit code
 *
 * Will automatically defer to the internal thread
 */
RunResult DeviceShell::runInShell(const CommandLine &cmd, const QByteArray &stdInData)
{
    Q_ASSERT(QThread::currentThread() != &m_thread);

    return run(cmd, stdInData);
}

DeviceShell::State DeviceShell::state() const { return m_shellScriptState; }

QStringList DeviceShell::missingFeatures() const { return m_missingFeatures; }

RunResult DeviceShell::run(const CommandLine &cmd, const QByteArray &stdInData)
{
    // If the script failed to install, use Process directly instead.
    bool useProcess = m_shellScriptState == State::Failed;

    // Transferring large amounts of stdInData is slow via the shell script.
    // Use Process directly if the size exceeds 100kb.
    useProcess |= stdInData.size() > (1024 * 100);

    if (useProcess) {
        Process proc;
        const CommandLine fallbackCmd = createFallbackCommand(cmd);
        qCDebug(deviceShellLog) << "Running fallback:" << fallbackCmd;
        proc.setCommand(fallbackCmd);
        proc.setWriteData(stdInData);
        // Practically unlimited timeout since we are most probably running dd for copying
        // data to a device here (=deployment), so a) we have no idea how long that might take
        // for large files, and b) the user can cancel this manually.
        proc.runBlocking(/*timout=*/1h);
        // TODO This misses interesting data like proc.errorMessage() and/or proc.exitMessage()
        return RunResult{proc.exitCode(), proc.rawStdOut(), proc.rawStdErr()};
    }

    const RunResult errorResult{-1, {}, {}};
    QTC_ASSERT(m_shellProcess, return errorResult);
    QTC_ASSERT(m_shellProcess->isRunning(), return errorResult);
    QTC_ASSERT(m_shellScriptState == State::Succeeded, return errorResult);

    QMutexLocker lk(&m_commandMutex);

    QWaitCondition waiter;
    const int id = ++m_currentId;
    m_commandOutput.insert(id, CommandRun{{-1, {}, {}}, &waiter});

    QMetaObject::invokeMethod(m_shellProcess.get(), [this, id, cmd, stdInData] {
        const QString command = QString("%1 \"%2\" %3\n").arg(id)
                                .arg(QString::fromLatin1(stdInData.toBase64()), cmd.toUserOutput());
        qCDebug(deviceShellLog) << "Running via shell:" << command;
        m_shellProcess->writeRaw(command.toUtf8());
    });

    waiter.wait(&m_commandMutex);

    const auto it = m_commandOutput.constFind(id);
    const RunResult result = *it;
    m_commandOutput.erase(it);

    return result;
}

void DeviceShell::close()
{
    QTC_ASSERT(QThread::currentThread() == thread(), return );
    QTC_ASSERT(m_thread.isRunning(), return );

    m_thread.quit();
    m_thread.wait();
}

/*!
 * \brief DeviceShell::setupShellProcess
 *
 * Override this function to setup the shell process.
 * The default implementation just sets the command line to "bash"
 */
void DeviceShell::setupShellProcess(Process *shellProcess)
{
    shellProcess->setCommand(CommandLine{"bash"});
}

/*!
* \brief DeviceShell::createFallbackCommand
* \param cmd The command to run
* \return The command to run in case the shell script is not available
*
* Creates a command to run in case the shell script is not available
*/
CommandLine DeviceShell::createFallbackCommand(const CommandLine &cmd)
{
    return cmd;
}

/*!
 * \brief DeviceShell::start
 * \return Returns true if starting the Shell process succeeded
 *
 * \note You have to call this function when deriving from DeviceShell. Current implementations call the function from their constructor.
 */
Result<> DeviceShell::start()
{
    m_shellProcess = std::make_unique<Process>();
    connect(m_shellProcess.get(), &Process::done, m_shellProcess.get(),
            [this] { emit done(m_shellProcess->resultData()); });
    connect(&m_thread, &QThread::finished, m_shellProcess.get(), [this] { closeShellProcess(); }, Qt::DirectConnection);

    setupShellProcess(m_shellProcess.get());

    m_shellProcess->setProcessMode(ProcessMode::Writer);

    // Moving the process into its own thread ...
    m_shellProcess->moveToThread(&m_thread);

    Result<> result;
    QMetaObject::invokeMethod(
        m_shellProcess.get(),
        [this]() -> Result<> {
            qCDebug(deviceShellLog)
                << "Starting shell process:" << m_shellProcess->commandLine().toUserOutput();
            m_shellProcess->start();

            if (!m_shellProcess->waitForStarted()) {
                closeShellProcess();
                return ResultError(msgProcessFailedToStart());
            }

            Result<> installResult = installShellScript();
            if (installResult) {
                connect(m_shellProcess.get(),
                        &Process::readyReadStandardOutput,
                        m_shellProcess.get(),
                        [this] { onReadyRead(); });
                connect(m_shellProcess.get(),
                        &Process::readyReadStandardError,
                        m_shellProcess.get(),
                        [this] {
                            const QByteArray stdErr = m_shellProcess->readAllRawStandardError();
                            qCWarning(deviceShellLog)
                                << "Received unexpected output on stderr:" << stdErr;
                        });

                connect(m_shellProcess.get(), &Process::done, m_shellProcess.get(), [this] {
                    if (m_shellProcess->resultData().m_exitCode != EXIT_SUCCESS
                        || m_shellProcess->resultData().m_exitStatus != QProcess::NormalExit) {
                        qCWarning(deviceShellLog) << "Shell exited with error code:"
                                                  << m_shellProcess->resultData().m_exitCode << "("
                                                  << m_shellProcess->exitMessage() << ")";
                        emit exitedIrregularly();
                    }
                });

                return ResultOk;
            }

            const CommandLine savedShellCmd = m_shellProcess->commandLine();

            if (m_shellProcess->isRunning())
                m_shellProcess->kill();
            m_shellProcess.reset();

            qCDebug(deviceShellLog()) << "Failed to install shell script: " << installResult.error();

            // Shell script setup might not have worked, e.g. due to a missing 'base64'
            // on the device, but the connection could in principle still be usable due
            // to the fallback to "single shot" mode.
            // Try to execute a simple command to check:

            QStringList args = savedShellCmd.splitArguments();
            QTC_ASSERT(!args.empty(), return ResultError(ResultAssert));
            args.removeLast(); // drop the traling /bin/bash, so we don't depend on bash either.
            args.append("echo");
            args.append("TEST_OK");

            const CommandLine fallbackTestCommand = {savedShellCmd.executable(), args};
            qCDebug(deviceShellLog()) << "Checking plain access as fallback"
                                      << fallbackTestCommand.toUserOutput();
            Process process;
            process.setProcessChannelMode(QProcess::MergedChannels);
            process.setCommand(fallbackTestCommand);
            process.runBlocking(5s);

            const QByteArray out = process.readAllRawStandardOutput().trimmed();

            if (process.result() != ProcessResult::FinishedWithSuccess || out != "TEST_OK") {
                qCDebug(deviceShellLog()) << "Fallback failed:" << process.error();
                return ResultError(installResult.error() + process.errorString() + out);
            }

            return ResultOk;
        },
        Qt::BlockingQueuedConnection,
        &result);

    return result;
}

static std::optional<QByteArray> waitAndReadLine(Process *process, QProcess::ProcessChannel channel)
{
    QByteArray result;
    while (process->waitForReadyRead(5s)) {
        if (channel == QProcess::StandardOutput)
            result += process->readAllRawStandardOutput();
        if (channel == QProcess::StandardError)
            result += process->readAllRawStandardError();
        if (result.contains('\n')) {
            QTC_CHECK(result.endsWith('\n'));
            return result;
        }
    }
    return {};
}

Result<QByteArray> DeviceShell::checkCommand(const QByteArray &command)
{
    const QByteArray checkCmd = "(type " + command + " || echo '<missing>')\n";

    m_shellProcess->writeRaw(checkCmd);
    const std::optional<QByteArray> out
        = waitAndReadLine(m_shellProcess.get(), QProcess::StandardOutput);
    if (!out) {
        if (!m_shellProcess->isRunning()) {
            return ResultError(
                m_shellProcess->exitMessage(Process::FailureMessageFormat::WithStdErr));
        }
        return ResultError(
            Tr::tr("Timeout while trying to check for %1.").arg(QString::fromUtf8(command)));
    }
    if (out->contains("<missing>")) {
        m_shellScriptState = State::Failed;
        m_missingFeatures.append(QString::fromUtf8(command));
        return ResultError(
            Tr::tr("Command \"%1\" was not found.").arg(QString::fromUtf8(command)));
    }

    return *out;
}

Result<> DeviceShell::installShellScript()
{
    if (m_forceFailScriptInstallation) {
        m_shellScriptState = State::Failed;
        return ResultError(Tr::tr("Script installation was forced to fail."));
    }

    static const QList<QByteArray> requiredCommands
        = {"base64", "cat", "echo", "kill", "mkfifo", "mktemp", "rm"};

    for (const QByteArray &command : requiredCommands) {
        auto checkResult = checkCommand(command);
        if (!checkResult)
            return ResultError(checkResult.error());
    }

    const static QByteArray shellScriptBase64 = FilePath(":/utils/scripts/deviceshell.sh")
                                                    .fileContents()
                                                    .value()
                                                    .replace("\r\n", "\n")
                                                    .toBase64();
    const QByteArray scriptCmd = "(scriptData=$(echo " + shellScriptBase64
                                 + " | base64 -d 2>/dev/null ) && /bin/sh -c \"$scriptData\") || "
                                   "echo ERROR_INSTALL_SCRIPT >&2\n";

    qCDebug(deviceShellLog) << "Installing shell script:" << scriptCmd;
    m_shellProcess->writeRaw(scriptCmd);

    while (m_shellScriptState == State::Unknown) {
        const std::optional<QByteArray> out
            = waitAndReadLine(m_shellProcess.get(), QProcess::StandardError);
        if (!out)
            return ResultError(Tr::tr("Timeout while waiting for shell script installation."));

        if (out->contains("SCRIPT_INSTALLED") && !out->contains("ERROR_INSTALL_SCRIPT")) {
            m_shellScriptState = State::Succeeded;
            return ResultOk;
        }
        if (out->contains("ERROR_INSTALL_SCRIPT")) {
            m_shellScriptState = State::Failed;
            return ResultError(
                Tr::tr("Failed to install shell script: %1").arg(QString::fromUtf8(*out)));
        }
        if (!out->isEmpty()) {
            qCDebug(deviceShellLog)
                << "Unexpected output while installing device shell script:" << *out;
        }
    }

    return ResultOk;
}

void DeviceShell::closeShellProcess()
{
    if (m_shellProcess) {
        if (m_shellProcess->isRunning()) {
            m_shellProcess->write("exit\nexit\n");
            if (!m_shellProcess->waitForFinished(2s))
                m_shellProcess->terminate();
        }
        m_shellProcess.reset();
    }
}

QByteArray::const_iterator next(const QByteArray::const_iterator &bufferEnd,
                                const QByteArray::const_iterator &itCurrent)
{
    for (QByteArray::const_iterator it = itCurrent; it != bufferEnd; ++it) {
        if (*it == '\n')
            return it;
    }
    return bufferEnd;
}

QByteArray byteArrayFromRange(QByteArray::const_iterator itStart, QByteArray::const_iterator itEnd)
{
    return QByteArray(itStart, std::distance(itStart, itEnd));
}

QList<std::tuple<int, DeviceShell::ParseType, QByteArray>> parseShellOutput(const QByteArray &data)
{
    auto itStart = data.cbegin();
    const auto itEnd = data.cend();

    QList<std::tuple<int, DeviceShell::ParseType, QByteArray>> result;

    for (auto it = next(itEnd, itStart); it != itEnd; ++it, itStart = it, it = next(itEnd, it)) {
        const QByteArray lineView = byteArrayFromRange(itStart, it);
        QTC_ASSERT(lineView.size() > 0, continue);

        const auto pidEnd = lineView.indexOf(':');
        const auto typeEnd = lineView.indexOf(':', pidEnd + 1);

        QTC_ASSERT(pidEnd != -1 && typeEnd != -1, continue);

        bool ok = false;
        const QLatin1String sId(lineView.begin(), pidEnd);
        const int id = QString(sId).toInt(&ok);
        QTC_ASSERT(ok, continue);

        const QByteArray data = byteArrayFromRange(lineView.begin() + typeEnd + 1, lineView.end());
        const QByteArray decoded = QByteArray::fromBase64(data);

        DeviceShell::ParseType t;
        char type = lineView.at(typeEnd - 1);
        switch (type) {
        case 'O':
            t = DeviceShell::ParseType::StdOut;
            break;
        case 'E':
            t = DeviceShell::ParseType::StdErr;
            break;
        case 'R':
            t = DeviceShell::ParseType::ExitCode;
            break;
        default:
            QTC_CHECK(false);
            continue;
        }

        result.append(std::make_tuple(id, t, decoded));
    }

    return result;
}

/*!
 * \brief DeviceShell::onReadyRead
 *
 * Reads lines coming from the multiplex script.
 *
 * The format is: "<command-id>:<type>:base64-encoded-text-or-returnvalue"
 * The possible <type>'s are:
 * O for stdout
 * E for stderr
 * R for exit code
 *
 * Multiple O/E messages may be received for a process. Once
 * a single "R" is received, the exit code is reported back
 * and no further messages from that process are expected.
 */
void DeviceShell::onReadyRead()
{
    m_commandBuffer += m_shellProcess->readAllRawStandardOutput();
    const qsizetype lastLineEndIndex = m_commandBuffer.lastIndexOf('\n') + 1;

    if (lastLineEndIndex == 0)
        return;

    const QByteArray input(m_commandBuffer.cbegin(), lastLineEndIndex);

    const auto result = parseShellOutput(input);

    QMutexLocker lk(&m_commandMutex);
    for (const auto &line : result) {
        const auto &[cmdId, type, data] = line;

        const auto itCmd = m_commandOutput.find(cmdId);
        QTC_ASSERT(itCmd != m_commandOutput.end(), continue);

        switch (type) {
        case ParseType::StdOut:
            itCmd->stdOut.append(data);
            break;
        case ParseType::StdErr:
            itCmd->stdErr.append(data);
            break;
        case ParseType::ExitCode: {
            bool ok = false;
            int exitCode;
            exitCode = data.toInt(&ok);
            QTC_ASSERT(ok, exitCode = -1);
            itCmd->exitCode = exitCode;
            itCmd->waiter->wakeOne();
            break;
        }
        }
    };

    if (lastLineEndIndex == m_commandBuffer.size())
        m_commandBuffer.clear();
    else
        m_commandBuffer = m_commandBuffer.mid(lastLineEndIndex);
}

QString msgProcessFailedToStart()
{
    return Tr::tr("The process failed to start.");
}

} // namespace Utils
