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

#include "pythonlanguageclient.h"

#include "pipsupport.h"
#include "pythonbuildconfiguration.h"
#include "pysideuicextracompiler.h"
#include "pythonconstants.h"
#include "pythonproject.h"
#include "pythonsettings.h"
#include "pythontr.h"
#include "pythonutils.h"

#include <coreplugin/editormanager/editormanager.h>
#include <coreplugin/icore.h>
#include <coreplugin/progressmanager/progressmanager.h>

#include <languageclient/languageclientinterface.h>
#include <languageclient/languageclientmanager.h>
#include <languageserverprotocol/textsynchronization.h>
#include <languageserverprotocol/workspace.h>

#include <projectexplorer/buildconfiguration.h>
#include <projectexplorer/buildsteplist.h>
#include <projectexplorer/extracompiler.h>
#include <projectexplorer/projectmanager.h>
#include <projectexplorer/target.h>

#include <QtTaskTree/QMappedTaskTreeRunner>

#include <texteditor/textdocument.h>
#include <texteditor/texteditor.h>

#include <utils/async.h>
#include <utils/infobar.h>
#include <utils/qtcprocess.h>

#include <QJsonDocument>

using namespace LanguageClient;
using namespace LanguageServerProtocol;
using namespace ProjectExplorer;
using namespace QtTaskTree;
using namespace TextEditor;
using namespace Utils;

namespace Python::Internal {

static constexpr char installPylsInfoBarId[] = "Python::InstallPyls";
static constexpr char updatePylsInfoBarId[] = "Python::updatePyls";
static constexpr char alwaysUpdateKey[] = "Python/AlwaysUpdatePyls";

class PythonLanguageServerState
{
public:
    enum {
        NotInstallable,
        Installable,
        Updatable,
        Installed
    } state;
    FilePath pylsModulePath;
};

static QHash<FilePath, PyLSClient*> &pythonClients()
{
    static QHash<FilePath, PyLSClient*> clients;
    return clients;
}

static FilePath pyLspPath(const FilePath &python)
{
    const QString version = pythonVersion(python);
    if (python.isLocal())
        return Core::ICore::userResourcePath() / "pylsp" / version;
    if (const Result<FilePath> tmpDir = python.tmpDir())
        return *tmpDir / "qc-pylsp" / version;
    return {};
}

static PythonLanguageServerState checkPythonLanguageServer(const FilePath &python, bool pipUsable)
{
    static QHash<FilePath, PythonLanguageServerState> m_stateCache;
    using namespace LanguageClient;
    using namespace std::chrono;

    if (auto it = m_stateCache.find(python); it != m_stateCache.end())
        return it.value();

    const FilePath lspPath = pyLspPath(python);
    if (!lspPath.isEmpty()) {
        const FilePath pylsp = (lspPath / "bin" / "pylsp").withExecutableSuffix();
        if (pylsp.exists()) {
            if (pipUsable) {
                Process pythonProcess;
                Environment env = pylsp.deviceEnvironment();
                env.appendOrSet("PYTHONPATH", lspPath.toUserOutput());
                pythonProcess.setEnvironment(env);
                pythonProcess.setCommand(
                    {python, {"-m", "pip", "list", "--outdated", "--format=json"}});
                pythonProcess.runBlocking(20s);
                QString output = pythonProcess.allOutput();

                // Only the first line contains the json data. Following lines might contain warnings.
                if (int index = output.indexOf('\n'); index >= 0)
                    output.truncate(index);

                const QJsonDocument doc = QJsonDocument::fromJson(output.toUtf8());
                for (const QJsonValue &value : doc.array()) {
                    if (value.toObject().value("name") == "python-lsp-server") {
                        Process pylsProcess;
                        Environment env = pylsp.deviceEnvironment();
                        env.appendOrSet("PYTHONPATH", lspPath.toUserOutput());
                        pylsProcess.setEnvironment(env);
                        pylsProcess.setCommand({pylsp, {"--version"}});
                        pylsProcess.runBlocking(20s);
                        output = pylsProcess.allOutput();
                        if (!output.contains(value.toObject().value("latest_version").toString()))
                            return {PythonLanguageServerState::Updatable, lspPath};
                        break;
                    }
                }
            }
            return m_stateCache.insert(python, {PythonLanguageServerState::Installed, lspPath})
                .value();
        }

        if (pipUsable)
            return {PythonLanguageServerState::Installable, lspPath};
    }
    return m_stateCache.insert(python, {PythonLanguageServerState::NotInstallable, FilePath()})
        .value();
}


class PyLSInterface : public StdIOClientInterface
{
public:
    PyLSInterface()
        : m_extraPythonPath("QtCreator-pyls-XXXXXX")
    { }

    TemporaryDirectory m_extraPythonPath;
protected:
    void startImpl() override
    {
        const FilePath python = m_cmd.executable();
        Environment env = python.deviceEnvironment();
        const FilePath lspPath = pyLspPath(python);
        if (!lspPath.isEmpty() && lspPath.exists() && QTC_GUARD(lspPath.isSameDevice(python))) {
            env.appendOrSet("PYTHONPATH", lspPath.path());
        }
        if (python.isLocal()) {
            // todo check where to put this tempdir in remote setups
            env.appendOrSet("PYTHONPATH", m_extraPythonPath.path().toUrlishString());
        }
        if (env.hasChanges())
            setEnvironment(env);
        StdIOClientInterface::startImpl();
    }
};

PyLSClient *clientForPython(const FilePath &python)
{
    if (auto client = pythonClients()[python])
        return client;
    auto interface = new PyLSInterface;
    interface->setCommandLine({python, {"-m", "pylsp"}});
    auto client = new PyLSClient(interface);
    client->setName(Tr::tr("Python Language Server (%1)").arg(python.toUserOutput()));
    client->updateConfiguration();
    LanguageFilter filter;
    filter.mimeTypes = QStringList() << Constants::C_PY_MIMETYPE << Constants::C_PY3_MIMETYPE;
    client->setSupportedLanguage(filter);
    client->start();
    pythonClients()[python] = client;
    for (Project *project : ProjectManager::projects()) {
        if (auto *bc = qobject_cast<PythonBuildConfiguration *>(project->activeBuildConfiguration())) {
            if (bc->python() == python)
                client->updateExtraCompilers(project);
        }
    }
    return client;
}

PyLSClient::PyLSClient(PyLSInterface *interface)
    : Client(interface)
    , m_extraCompilerOutputDir(interface->m_extraPythonPath.path())
{
    connect(this, &Client::initialized, this, &PyLSClient::updateConfiguration);
    connect(PythonSettings::instance(), &PythonSettings::pylsConfigurationChanged,
            this, &PyLSClient::updateConfiguration);
    connect(PythonSettings::instance(), &PythonSettings::pylsEnabledChanged,
            this, [this](const bool enabled){
                if (!enabled)
                    LanguageClientManager::shutdownClient(this);
            });

}

PyLSClient::~PyLSClient()
{
    pythonClients().remove(pythonClients().key(this));
    for (const QPointer<PySideUicExtraCompiler> &compiler : std::as_const(m_extraCompilers))
        deleteExtraCompiler(compiler);
}

void PyLSClient::updateConfiguration()
{
    const auto doc = QJsonDocument::fromJson(PythonSettings::pylsConfiguration().toUtf8());
    if (doc.isArray())
        Client::updateConfiguration(doc.array());
    else if (doc.isObject())
        Client::updateConfiguration(doc.object());
}

void PyLSClient::openDocument(TextDocument *document)
{
    using namespace LanguageServerProtocol;
    if (reachable()) {
        const FilePath documentPath = document->filePath();
        if (isSupportedDocument(document) && !pythonProjectForFile(documentPath)) {
            const FilePath workspacePath = documentPath.parentDir();
            if (!m_extraWorkspaceDirs.contains(workspacePath)) {
                WorkspaceFoldersChangeEvent event;
                event.setAdded({WorkSpaceFolder(hostPathToServerUri(workspacePath),
                                                workspacePath.fileName())});
                DidChangeWorkspaceFoldersParams params;
                params.setEvent(event);
                DidChangeWorkspaceFoldersNotification change(params);
                sendMessage(change);
                m_extraWorkspaceDirs.append(workspacePath);
            }
        }
    }
    Client::openDocument(document);
}

void PyLSClient::buildConfigurationClosed(BuildConfiguration *bc)
{
    if (auto pythonBuildConfig = qobject_cast<PythonBuildConfiguration *>(bc)) {
        const FilePath python = pythonBuildConfig->python();
        if (!python.isEmpty() && clientForPython(python) == this) {
            const QList<QPointer<PySideUicExtraCompiler>> compilers = Utils::takeAll(
                m_extraCompilers,
                [project = bc->project()](const QPointer<PySideUicExtraCompiler> &compiler) {
                    return compiler && compiler->project() == project;
                });
            for (PySideUicExtraCompiler *compiler : compilers)
                deleteExtraCompiler(compiler);
        }
    }
    Client::buildConfigurationClosed(bc);
}

void PyLSClient::updateExtraCompilers(const Project *project)
{
    QList<QPointer<PySideUicExtraCompiler>> oldCompilers = Utils::takeAll(
        m_extraCompilers, [project](const QPointer<PySideUicExtraCompiler> &compiler) {
            return compiler && compiler->project() == project;
        });
    if (auto pythonBuildConfig = qobject_cast<PythonBuildConfiguration *>(
            project->activeBuildConfiguration())) {
        if (auto step = pythonBuildConfig->buildSteps()->firstOfType<PySideBuildStep>()) {
            const FilePath pySideUic = step->pySideUicPath();
            const FilePaths uiFiles = step->uiFiles();

            for (const FilePath &uiFile : uiFiles) {
                const FilePath generated = uiFile.parentDir().pathAppended(
                    "/ui_" + uiFile.baseName() + ".py");

                const int index = Utils::indexOf(oldCompilers, [&](PySideUicExtraCompiler *c) {
                    return c->pySideUicPath() == pySideUic && c->source() == uiFile;
                });

                PySideUicExtraCompiler *compiler;
                if (index < 0) {
                    compiler = new PySideUicExtraCompiler(
                        pySideUic, step->project(), uiFile, {generated}, this);
                    connect(
                        compiler,
                        &ExtraCompiler::contentsChanged,
                        this,
                        [this, compiler](const FilePath &file) {
                            updateExtraCompilerContents(compiler, file);
                        });
                    if (compiler->isDirty())
                        compiler->compileFile();
                } else {
                    compiler = oldCompilers.takeAt(index);
                }
                m_extraCompilers.append(compiler);
            }
        }
    }

    for (PySideUicExtraCompiler *compiler : std::as_const(oldCompilers))
        deleteExtraCompiler(compiler);
}

void PyLSClient::updateExtraCompilerContents(ExtraCompiler *compiler, const FilePath &file)
{
    const FilePath target = m_extraCompilerOutputDir.pathAppended(file.fileName());

    target.writeFileContents(compiler->content(file));
}

void PyLSClient::deleteExtraCompiler(ExtraCompiler *compiler)
{
    if (!compiler)
        return;
    if (const FilePaths targets = compiler->targets(); !targets.isEmpty())
        m_extraCompilerOutputDir.pathAppended(targets.first().fileName()).removeFile();
    compiler->disconnect(this);
    delete compiler;
}

PyLSClient *PyLSClient::clientForPython(const FilePath &python)
{
    return pythonClients()[python];
}

class PyLSConfigureAssistant : public QObject
{
public:
    PyLSConfigureAssistant();

    void handlePyLSState(const FilePath &python,
                         const PythonLanguageServerState &state,
                         TextDocument *document);
    void resetEditorInfoBar(TextDocument *document);
    void installPythonLanguageServer(const FilePath &python,
                                     QPointer<TextDocument> document,
                                     const FilePath &pylsPath, bool silent, bool upgrade);
    void startCheckPythonLanguageServer(
        const FilePath &python, QPointer<TextDocument> document, bool pipUsable);
    void openDocument(const FilePath &python, TextDocument *document);

    QHash<FilePath, QList<TextDocument *>> m_infoBarEntries;
    QMappedTaskTreeRunner<TextDocument *> m_documentRunner;
    QSingleTaskTreeRunner m_pipInstallerRunner;
};

void PyLSConfigureAssistant::installPythonLanguageServer(const FilePath &python,
                                                         QPointer<TextDocument> document,
                                                         const FilePath &pylsPath,
                                                         bool silent,
                                                         bool upgrade)
{
    document->infoBar()->removeInfo(installPylsInfoBarId);

    // Hide all install info bar entries for this python, but keep them in the list
    // so the language server will be setup properly after the installation is done.
    for (TextDocument *additionalDocument : m_infoBarEntries[python])
        additionalDocument->infoBar()->removeInfo(installPylsInfoBarId);

    PipInstallerData data;
    data.python = python;
    data.targetPath = pylsPath;
    data.packages = {PipPackage{"python-lsp-server[all]", "Python Language Server"}};
    data.upgrade = upgrade;
    data.silent = silent;

    const auto onDone = [this, python, document](DoneWith result) {
        const QList<TextDocument *> additionalDocuments = m_infoBarEntries.take(python);
        if (result == DoneWith::Success) {
            if (PyLSClient *client = clientForPython(python)) {
                if (document)
                    LanguageClientManager::openDocumentWithClient(document, client);
                for (TextDocument *additionalDocument : additionalDocuments)
                    LanguageClientManager::openDocumentWithClient(additionalDocument, client);
            }
        }
    };

    m_pipInstallerRunner.start(pipInstallerTask(data), {}, onDone);
}

void PyLSConfigureAssistant::startCheckPythonLanguageServer(
    const FilePath &python, QPointer<TextDocument> document, bool pipUsable)
{
    if (!document)
        return;
    const auto onSetup = [python, pipUsable](Async<PythonLanguageServerState> &task) {
        task.setConcurrentCallData(&checkPythonLanguageServer, python, pipUsable);
    };
    const auto onDone =
        [this, python, document](const Async<PythonLanguageServerState> &task) {
            if (!document || !task.isResultAvailable())
                return;
            handlePyLSState(python, task.result(), document);
        };

    using namespace std::literals::chrono_literals;
    const Group recipe{AsyncTask<PythonLanguageServerState>(onSetup, onDone).withTimeout(10s)};
    m_documentRunner.start(document, recipe);
}

void PyLSConfigureAssistant::openDocument(const FilePath &python, TextDocument *document)
{
    resetEditorInfoBar(document);
    if (!PythonSettings::pylsEnabled() || !python.exists() || document->isTemporary())
        return;

    if (auto client = pythonClients().value(python)) {
        LanguageClientManager::openDocumentWithClient(document, client);
        return;
    }

    pipIsUsableAsync(
        python,
        [self = QPointer(this), python, document = QPointer<TextDocument>(document)](
            const bool pipUsable) {
            if (!self)
                return;
            self->startCheckPythonLanguageServer(python, document, pipUsable);
        });
}

void PyLSConfigureAssistant::handlePyLSState(const FilePath &python,
                                             const PythonLanguageServerState &state,
                                             TextDocument *document)
{
    if (state.state == PythonLanguageServerState::NotInstallable)
        return;

    InfoBar *infoBar = document->infoBar();
    if (state.state == PythonLanguageServerState::Installable
        && infoBar->canInfoBeAdded(installPylsInfoBarId)) {
        auto message = Tr::tr("Install Python language server (PyLS) for %1 (%2). "
                              "The language server provides Python specific completion and annotation.")
                           .arg(pythonVersion(python), python.toUserOutput());
        InfoBarEntry info(installPylsInfoBarId, message, InfoBarEntry::GlobalSuppression::Enabled);
        info.addCustomButton(Tr::tr("Install"), [this, python, document, state] {
            installPythonLanguageServer(python, document, state.pylsModulePath, false, false);
        });
        infoBar->addInfo(info);
        m_infoBarEntries[python] << document;
    } else if (state.state == PythonLanguageServerState::Updatable) {
        if (infoBar->canInfoBeAdded(updatePylsInfoBarId)) {
            auto message = Tr::tr("Update Python language server (PyLS) for %1 (%2).")
                               .arg(pythonVersion(python), python.toUserOutput());
            InfoBarEntry info(updatePylsInfoBarId, message);
            info.addCustomButton(
                Tr::tr("Always Update"),
                [this, python, document, state] {
                    Core::ICore::settings()->setValue(alwaysUpdateKey, true);
                    installPythonLanguageServer(python, document, state.pylsModulePath, false, true);
                },
                {},
                InfoBarEntry::ButtonAction::SuppressPersistently);
            info.addCustomButton(
                Tr::tr("Update"),
                [this, python, document, state] {
                    installPythonLanguageServer(python, document, state.pylsModulePath, false, true);
                },
                {},
                InfoBarEntry::ButtonAction::Hide);
            info.addCustomButton(
                Tr::tr("Never"),
                [document, python] {
                    if (auto client = clientForPython(python))
                        LanguageClientManager::openDocumentWithClient(document, client);
                },
                {},
                InfoBarEntry::ButtonAction::SuppressPersistently);
            info.setCancelButtonInfo([python, document]{
                if (auto client = clientForPython(python))
                    LanguageClientManager::openDocumentWithClient(document, client);
            });
            infoBar->addInfo(info);

            m_infoBarEntries[python] << document;
        } else if (Core::ICore::settings()->value(alwaysUpdateKey, false).toBool()) {
            installPythonLanguageServer(python, document, state.pylsModulePath, true, true);
        } else if (auto client = clientForPython(python)) {
            LanguageClientManager::openDocumentWithClient(document, client);
        }
    } else if (state.state == PythonLanguageServerState::Installed) {
        if (auto client = clientForPython(python))
            LanguageClientManager::openDocumentWithClient(document, client);
    }
}

void PyLSConfigureAssistant::resetEditorInfoBar(TextDocument *document)
{
    for (QList<TextDocument *> &documents : m_infoBarEntries)
        documents.removeAll(document);
    document->infoBar()->removeInfo(installPylsInfoBarId);
    m_documentRunner.resetKey(document);
}

PyLSConfigureAssistant::PyLSConfigureAssistant()
{
    Core::EditorManager::instance();

    connect(Core::EditorManager::instance(),
            &Core::EditorManager::documentClosed,
            this,
            [this](Core::IDocument *document) {
                if (auto textDocument = qobject_cast<TextDocument *>(document))
                    resetEditorInfoBar(textDocument);
            });
}

static PyLSConfigureAssistant &pyLSConfigureAssistant()
{
    static PyLSConfigureAssistant thePyLSConfigureAssistant;
    return thePyLSConfigureAssistant;
}

void openDocumentWithPython(const FilePath &python, TextDocument *document)
{
    pyLSConfigureAssistant().openDocument(python, document);
}

} // Python::Internal
