/*
 * Copyright 2023 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "appmanagercgroupd.h"
#include <QDebug>
#include <QDBusMetaType>
#include <QThread>
#include <QCoreApplication>
#include <QDir>
#include <libcgroup.h>
#include <unistd.h>
#include "common.h"
#include "confmanager.h"

#define CPU_MAX "cpu.max"
#define CGROUP_FREEZE "cgroup.freeze"

AppManagerCgroupd::AppManagerCgroupd(QObject *parent)
    : QObject{parent}
    , m_cgroupInited(false)
{
    qDBusRegisterMetaType<QMap<QString,QVariant>>();
}

QList<int> AppManagerCgroupd::pids(const QString &cgroupName)
{
    int size = 0;
    int ret = 0;
    const char *controller = "pids";
    int retNumber = 0;
    QList<int> result;

    pid_t *pids = nullptr;
    QByteArray ba = cgroupName.toLocal8Bit();
    retNumber = cgroup_get_procs(ba.data(), (char *)controller, &pids, &size);
    if (retNumber != 0) {
        qWarning() << "cgroup_get_procs error, " << cgroup_strerror(cgroup_get_last_errno());
        return QList<int>();
    }
    
    if (size == 0 || !pids) {
        return QList<int>();
    }

    for (int i=0; i<size; ++i) {
        result.push_back(pids[i]);
    }

    delete[] pids;

    return result;
}

void AppManagerCgroupd::thawProcess(struct cgroup_file_info *info, char *root)
{
	if (info->type == CGROUP_FILE_TYPE_DIR) {
		printf("path %s, parent %s, relative %s, full %s\n",
			info->path, info->parent,
			info->full_path + strlen(root) - 1,
			info->full_path);
        QString cgroupName = QString(info->full_path + strlen(root) - 1).prepend("kylin-app-manager");
        if (cgroupName.contains("instance_")) {
            SetProcessCGroupResourceLimit(cgroupName, CGROUP_FREEZE, 0);
        }
	}
}

void AppManagerCgroupd::thawAllProcesses(bool walk)
{
    if (!walk) {
        for (auto &cgroupName : qAsConst(m_frozenCGroups)) {
            SetProcessCGroupResourceLimit(cgroupName, CGROUP_FREEZE, 0);
        }
        return;
    }

	struct cgroup_file_info info;
	char root[FILENAME_MAX];
	char *controller;
	void *handle;
	int lvl, i;
	int ret;
	ret = cgroup_init();
	if (ret != 0) {
		fprintf(stderr, "Init failed\n");
		exit(EXIT_FAILURE);
	}

	ret = cgroup_walk_tree_begin("cgroup", "kylin-app-manager", 0, &handle, &info, &lvl);

	if (ret != 0) {
		fprintf(stderr, "Walk failed\n");
		exit(EXIT_FAILURE);
	}

	strcpy(root, info.full_path);
	printf("Begin pre-order walk\n");
	printf("root is %s\n", root);
	thawProcess(&info, root);
	while ((ret = cgroup_walk_tree_next(0, &handle, &info, lvl)) !=
			ECGEOF) {
		thawProcess(&info, root);
	}
	cgroup_walk_tree_end(&handle);
}

QMap<QString, QVariant> AppManagerCgroupd::InitCGroup(const QString &rootPath)
{
    return {
        { kResult, initCgroup(rootPath) },
        { kErrMsg, m_cgroupInitErrMsg }
    };
}

QMap<QString, QVariant> AppManagerCgroupd::CreateProcessCGroup(const QString &appId, int pid)
{
    auto result = initResult();

    if (!m_cgroupInited) {
        result[kErrMsg] = m_cgroupInitErrMsg;
        return result;
    }

    if (appId.isEmpty() || pid <=0) {
        result[kErrMsg] = "The parameter is invalid";
        return result;
    }

    cgroup *appIdCgroup = createCgroup(m_rootCgroup + "/" + appId);
    if (!appIdCgroup) {
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        return result;
    }
    cgroup_free(&appIdCgroup);

    QString retCgroupName = m_rootCgroup + "/" + appId + "/instance_" + QString::number(pid);
    cgroup *appInstCgroup = createCgroup(retCgroupName, kControllerNames);
    if (!appInstCgroup) {
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        return result;
    }

    int ret = cgroup_attach_task_pid(appInstCgroup, pid);
    if (ret != 0) {
        result[kErrMsg] = cgroup_strerror(ret);
        cgroup_free(&appInstCgroup);
        return result;
    }

    cgroup_free(&appInstCgroup);
    result[kResult] = retCgroupName;

    m_createdCGroupNames.push_back(retCgroupName);
    return result;
}

QMap<QString, QVariant> AppManagerCgroupd::CreateProcessCGroup(const QString &appId, const QList<int> &pids)
{
    auto result = initResult();

    if (!m_cgroupInited) {
        result[kErrMsg] = m_cgroupInitErrMsg;
        return result;
    }

    if (appId.isEmpty() || pids.isEmpty()) {
        result[kErrMsg] = "The parameter is invalid";
        return result;
    }

    cgroup *appIdCgroup = createCgroup(m_rootCgroup + "/" + appId);
    if (!appIdCgroup) {
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        return result;
    }
    cgroup_free(&appIdCgroup);

    QString retCgroupName = m_rootCgroup + "/" + appId + "/instance_" + QString::number(pids.first());
    cgroup *appInstCgroup = createCgroup(retCgroupName, kControllerNames);
    if (!appInstCgroup) {
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        return result;
    }

    for (auto const &pid : pids) {
        int ret = cgroup_attach_task_pid(appInstCgroup, pid);
        if (ret != 0) {
            result[kErrMsg] = cgroup_strerror(ret);
            cgroup_free(&appInstCgroup);
            return result;
        }        
    }

    cgroup_free(&appInstCgroup);
    result[kResult] = retCgroupName;

    m_createdCGroupNames.push_back(retCgroupName);
    return result;
}

QMap<QString, QVariant> AppManagerCgroupd::DeleteProcessCGroup(const QString &cgroupName)
{
    auto result = initResult();
    int retNumber = 0;
    cgroup *cgroup = cgroup_new_cgroup(cgroupName.toLatin1().data());
    if (!cgroup) {
        result[kResult] = false;
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        qWarning() << "DeleteProcessCGroup error " << result[kErrMsg];
        return result;
    }

    retNumber = cgroup_get_cgroup(cgroup);
    if (retNumber != 0) {
        result[kResult] = false;
        result[kErrMsg] = cgroup_strerror(retNumber);
        qWarning() << "DeleteProcessCGroup cgroup_get_cgroup error " << result[kErrMsg];
        return result;
    }

    int deleteCounter = 0;
    do {
        retNumber = cgroup_delete_cgroup_ext(cgroup, CGFLAG_DELETE_RECURSIVE);
        qApp->processEvents();
        usleep(10000);
        ++ deleteCounter;
    } while (retNumber == ECGOTHER && deleteCounter <= 1000);

    if (retNumber != 0) {
        result[kResult] = false;
        result[kErrMsg] = cgroup_strerror(retNumber);
        qWarning() << "DeleteProcessCGroup cgroup_delete_cgroup error " 
                   << result[kErrMsg] << retNumber;
        return result;
    }
    result[kResult] = true;

    m_createdCGroupNames.removeOne(cgroupName);

    return result;
}

QMap<QString, QVariant> AppManagerCgroupd::SetProcessCGroupResourceLimit(const QString &cgroupName, 
                                                                         const QString &attrName, 
                                                                         int value)
{
    auto result = initResult();
    int retNumber = 0;
    
    auto &config = common::Singleton<ConfManager>::GetInstance();
    if (attrName == CGROUP_FREEZE && 
        !config.features().contains(FEAT_FROZEN) &&
        value == 1) {
        result[kErrMsg] = "The 'frozen' feature is not enabled.";
        result[kResult] = false;
        return result;
    }

    if (!m_cgroupInited) {
        result[kErrMsg] = m_cgroupInitErrMsg;
        result[kResult] = false;
        return result;
    }

    if (cgroupName.isEmpty() || attrName.isEmpty() || attrName.split(".").size() < 2) {
        result[kErrMsg] = "The parameter is invalid";
        result[kResult] = false;
        return result;
    }

    cgroup *cgroup = cgroup_new_cgroup(cgroupName.toLocal8Bit().data());
    if (!cgroup) {
        qWarning() << cgroup_get_last_errno();
    }

    QString resourceName = attrName.split(".").first();
    QByteArray ba = resourceName.toLocal8Bit();
    cgroup_controller *controller = cgroup_add_controller(cgroup, ba.data());
    if (!controller) {
        result[kErrMsg] = cgroup_strerror(cgroup_get_last_errno());
        result[kResult] = false;
        cgroup_free(&cgroup);
        qWarning() << "cgroup_get_controller failed, "
                   << result[kErrMsg]
                   << attrName.toLocal8Bit().data()
                   << ba.data()
                   << value;
        return result;
    }

    if (attrName == CPU_MAX) {
        QString strVale = value == 100000 ? 
                          QString("max 100000"):
                          QString("%1 100000").arg(value);
        retNumber = cgroup_set_value_string(controller, attrName.toLocal8Bit().data(), strVale.toLocal8Bit().data());
    } else {
        retNumber = cgroup_set_value_int64(controller, attrName.toLocal8Bit().data(), value);
    }
    
    if (retNumber != 0) {
        result[kErrMsg] = cgroup_strerror(retNumber);
        result[kResult] = false;
        cgroup_free_controllers(cgroup);
        cgroup_free(&cgroup);
        qWarning() << "cgroup_add_value failed, "
                   << cgroup_strerror(retNumber)
                   << attrName.toLocal8Bit().data()
                   << value;
        return result;
    }

    retNumber = cgroup_modify_cgroup(cgroup);
    if (retNumber != 0) {
        result[kErrMsg] = cgroup_strerror(retNumber);
        result[kResult] = false;
        cgroup_free_controllers(cgroup);
        cgroup_free(&cgroup);
        qWarning() << "cgroup_modify_cgroup failed, "
                   << cgroup_strerror(retNumber)
                   << attrName.toLocal8Bit().data()
                   << value;
        return result;
    }

    if (attrName == CGROUP_FREEZE) {
        if (value == 1) {
            m_frozenCGroups.push_back(cgroupName);
        } else {
            m_frozenCGroups.removeOne(cgroupName);
        }
    }
    
    cgroup_free_controllers(cgroup);
    cgroup_free(&cgroup);
    result[kResult] = true; 
    return result;
}

QMap<QString, QVariant> AppManagerCgroupd::CGroupNameWithPid(int pid)
{   
    int size = 0;
    int ret = 0;
    const char *controller = "pids";
    auto result = initResult();
    int retNumber = 0;
    for (auto const &cgroupName : qAsConst(m_createdCGroupNames)) {
        pid_t *pids = nullptr;
        QByteArray ba = cgroupName.toLocal8Bit();
        retNumber = cgroup_get_procs(ba.data(), (char *)controller, &pids, &size);
        if (retNumber != 0) {
            qWarning() << "cgroup_get_procs error, " << cgroup_strerror(cgroup_get_last_errno());
            continue;
        }
        
        if (size == 0 || !pids) {
            continue;
        }

        for (int i=0; i<size; ++i) {
            if (pid == pids[i]) {
                result[kResult] = cgroupName;
                delete[] pids;
                return result;
            }
        }

        delete[] pids;
    }
    return result;
}

QMap<QString, QVariant> AppManagerCgroupd::PidsOfCGroup(const QString &cgroupName)
{
    pid_t *pids = nullptr;
    const char *controller = "pids";
    int size = 0;
    auto result = initResult();
    QList<int> retPids;
    QByteArray ba = cgroupName.toLocal8Bit();
    int retNumber = cgroup_get_procs(ba.data(), (char *)controller, &pids, &size);
    if (retNumber != 0) {
        result[kErrMsg] = cgroup_strerror(retNumber);
        qWarning() << "cgroup_get_procs error, " << result[kErrMsg];
        return result;
    }
    
    if (size == 0 || !pids) {
        return result;
    }

    for (int i=0; i<size; ++i) {
        retPids.push_back(pids[i]);
    }

    result[kResult] = QVariant::fromValue<QList<int>>(retPids);
    delete[] pids;
    return result;
}

bool AppManagerCgroupd::initCgroup(const QString &rootPath)
{
    if (m_cgroupInited && m_rootCgroup == rootPath) {
        return true;
    }

    auto initFailed = [this] {
        int errorNo = cgroup_get_last_errno();
        m_cgroupInitErrMsg = cgroup_strerror(errorNo);
        m_cgroupInited = false;  
        qWarning() << "Cgroup init failed, " << m_cgroupInitErrMsg;     
    };

    if (cgroup_init() != 0) {
        initFailed();
        return false;
    }

    cgroup *rootCgroup = createCgroup(rootPath, kControllerNames);
    if (!rootCgroup) {
        initFailed();
        return false;
    }

    m_cgroupInited = true;
    m_rootCgroup = rootPath;
    qDebug() << "init cgroup success. path is " << rootPath;

    return true;
}

cgroup *AppManagerCgroupd::createCgroup(const QString &cgroupName, const QStringList &controllers)
{
    cgroup *cgroup = cgroup_new_cgroup(cgroupName.toLocal8Bit().data());
    if (!cgroup) {
        return nullptr;
    }

    if (!controllers.isEmpty()) {
        for (auto const &controlName : qAsConst(controllers)) {
            cgroup_controller *controller = cgroup_add_controller(cgroup, controlName.toLocal8Bit().data());
            if (!controller) {
                qWarning() << "cgroup_add_controller failed " << cgroup_strerror(cgroup_get_last_errno());
                cgroup_free(&cgroup);
                return nullptr;
            }
        }
    }

    if (cgroup_create_cgroup(cgroup, 0) != 0) { 
        cgroup_free_controllers(cgroup);
        cgroup_free(&cgroup);
        return nullptr;
    }

    return cgroup;
}

QMap<QString, QVariant> AppManagerCgroupd::initResult()
{
    QMap<QString, QVariant> result;
    result.insert(kResult, "");
    result.insert(kResult, "");
    return result;
}
