#include "KMContainer.h"
#include "KMConfigData.h"
#include "KMCrun.h"

#include <fmt/format.h>
#include <unistd.h>
#include <memory>
#include <filesystem>
#include <unordered_set>
#include <uuid/uuid.h>

#include "common/KMUtil.h"
#include "common/KMLogger.h"

extern char **environ;

using namespace std;

class KMContainer::Private
{
public:
    KMConfigDbata cfg;
    std::string id;
    std::string baseDir;
};

KMContainer::KMContainer()
    : d_ptr(make_unique<Private>())
{
    // 设置模板参数
    initTemplate();
    // 设置基本环境变量
    initBaseEnv();
    // 设置process默认参数
    initProcessDefault();
}

KMContainer::~KMContainer()
{
}

/**
 * @brief 初始化模板
 */
void KMContainer::initTemplate()
{
    d_ptr->cfg.ociVersion = "1.0.1";
    d_ptr->cfg.hostname = "kaiming";
    d_ptr->cfg.domainname = "kaiming";
    d_ptr->cfg.root = Root{ "rootfs", false };
    d_ptr->cfg.linuxConfig = Linux{};
    d_ptr->cfg.linuxConfig->uidMappings = std::vector<IDMap>{ { getuid(), getuid(), 1 } };
    d_ptr->cfg.linuxConfig->gidMappings = std::vector<IDMap>{ { getgid(), getgid(), 1 } };
    d_ptr->cfg.linuxConfig->namespaces = { { "pid" }, { "mount" }, { "uts" }, { "user" } };
}

/**
 * @brief 初始化环境变量
 */
void KMContainer::initBaseEnv()
{
    d_ptr->cfg.process = Process{};
    d_ptr->cfg.process->env = std::vector<std::string>{};
    const std::vector<std::string> envDefault = { "DISPLAY",
                                                  "LANG",
                                                  "LANGUAGE",
                                                  "SHELL",
                                                  "TEMP",
                                                  "TEMPDIR",
                                                  "PWD",
                                                  "EMAIL_ADDRESS",
                                                  "HOME",
                                                  "HOSTNAME",
                                                  "LOGNAME",
                                                  "REAL_NAME",
                                                  "TERM",
                                                  "USER",
                                                  "USERNAME",
                                                  "XAUTHORITY",
                                                  "XDG_SESSION_DESKTOP",
                                                  "D_DISABLE_RT_SCREEN_SCALE",
                                                  "XMODIFIERS",
                                                  "DESKTOP_SESSION",
                                                  "DEEPIN_WINE_SCALE",
                                                  "XDG_CURRENT_DESKTOP",
                                                  "XDG_DATA_HOME",
                                                  "XIM",
                                                  "XDG_SESSION_TYPE",
                                                  "XDG_RUNTIME_DIR",
                                                  "CLUTTER_IM_MODULE",
                                                  "QT4_IM_MODULE",
                                                  "GTK_IM_MODULE",
                                                  "auto_proxy",
                                                  "http_proxy",
                                                  "https_proxy",
                                                  "ftp_proxy",
                                                  "SOCKS_SERVER",
                                                  "no_proxy",
                                                  "QT_IM_MODULE",
                                                  "QT_QPA_PLATFORM",
                                                  "LINGLONG_ROOT",
                                                  "WAYLAND_DISPLAY",
                                                  "QT_QPA_PLATFORM",
                                                  "QT_WAYLAND_SHELL_INTEGRATION",
                                                  "GDMSESSION",
                                                  "QT_WAYLAND_FORCE_DPI",
                                                  "GIO_LAUNCHED_DESKTOP_FILE",
                                                  "GNOME_DESKTOP_SESSION_ID",
                                                  "PYTHONPATH",
                                                  "PERLLIB",
                                                  "PERL5LIB",
                                                  "XCURSOR_PATH",
                                                  "GST_PLUGIN_PATH_1_0",
                                                  "GST_REGISTRY",
                                                  "GST_REGISTRY_1_0",
                                                  "GST_PLUGIN_PATH",
                                                  "GST_PLUGIN_SYSTEM_PATH",
                                                  "GST_PLUGIN_SCANNER",
                                                  "GST_PLUGIN_SCANNER_1_0",
                                                  "GST_PLUGIN_SYSTEM_PATH_1_0",
                                                  "GST_PRESET_PATH",
                                                  "GST_PTP_HELPER",
                                                  "GST_PTP_HELPER_1_0",
                                                  "GST_INSTALL_PLUGINS_HELPER",
                                                  "KRB5CCNAME",
                                                  "XKB_CONFIG_ROOT",
                                                  "GIO_EXTRA_MODULES",
                                                  "GDK_BACKEND" };

    for (const auto &filter : envDefault)
    {
        for (int i = 0; environ[i] != nullptr; ++i)
        {
            if (environ[i][filter.length() + 1] == '\0')
                continue;

            // 解决环境变量长短名相同时问题
            string key = filter + string("=");
            if (std::strncmp(environ[i], key.c_str(), key.length()) == 0)
            {
                d_ptr->cfg.process->env.emplace_back(environ[i]);
                break;
            }
        }
    }
}

/**
 * @brief 设置appID
 * @param id appID
 */
void KMContainer::setAppID(const std::string &id)
{
    d_ptr->id = id;
}

/**
 * @brief 设置base目录,必须设置用于修订挂载点信息
 * @param baseDir 目录
 */
void KMContainer::setBaseDir(const std::string &baseDir)
{
    d_ptr->baseDir = baseDir;
}

/**
 * @brief 获取appID
 * @return appID
 */
std::string KMContainer::getAppID()
{
    return d_ptr->id;
}

/**
 * @brief 设置oci版本
 * @param version 版本号
 */
void KMContainer::setOciVersion(const std::string &version)
{
    d_ptr->cfg.ociVersion = version;
}

/**
 * @brief 获取版本
 * @return string 版本号
 */
std::string KMContainer::getOciVersion()
{
    if (d_ptr->cfg.ociVersion)
    {
        return d_ptr->cfg.ociVersion.value();
    }

    return {};
}

/**
 * @brief 设置主机名
 * @param hostName 主机名
 */
void KMContainer::setHostName(const std::string &hostName)
{
    d_ptr->cfg.hostname = hostName;
}

/**
 * @brief 获取容器主机名
 * @return string 主机名
 */
std::string KMContainer::setHostName()
{
    if (d_ptr->cfg.hostname)
    {
        return d_ptr->cfg.hostname.value();
    }

    return {};
}

/**
 * @brief 设置域名
 * @param domainName 域名
 */
void KMContainer::setDomainName(const std::string &domainName)
{
    d_ptr->cfg.domainname = domainName;
}

/**
 * @brief 获取域名
 * @return string 域名
 */
std::string KMContainer::getDomainName()
{
    if (d_ptr->cfg.domainname)
    {
        return d_ptr->cfg.domainname.value();
    }

    return {};
}

/**
 * @brief 设置根目录
 * @param root 根目录
 * @param readonly 是否只读系统
 */
void KMContainer::setRoot(const std::string &root, bool readonly)
{
    if (!d_ptr->cfg.root)
    {
        d_ptr->cfg.root = Root{};
    }
    d_ptr->cfg.root->path = root;
    d_ptr->cfg.root->readonly = readonly;
}

/**
 * @brief 获取根结构数据
 * @param Root 数据
 */
Root KMContainer::getRoot()
{
    if (d_ptr->cfg.root)
    {
        return d_ptr->cfg.root.value();
    }

    return Root{};
}

/**
 * @brief 添加注释；包含任意传递给容器的元数据
 * @param key 键
 * @param value 值
 */
void KMContainer::addAnnotation(const std::string &key, const std::string &value)
{
    if (d_ptr->cfg.annotations)
    {
        d_ptr->cfg.annotations.value()[key] = value;
    }
    else
    {
        d_ptr->cfg.annotations = std::map<std::string, std::string>{ { key, value } };
    }
}

/**
 * @brief 获取注释；包含任意传递给容器的元数据
 * @return 注释
 */
std::map<std::string, std::string> KMContainer::getAnnotations()
{
    if (d_ptr->cfg.annotations)
    {
        return d_ptr->cfg.annotations.value();
    }

    return std::map<std::string, std::string>{};
}

/**
 * @brief 删除注释
 * @param key 键
 */
void KMContainer::delAnnotation(const std::string &key)
{
    if (d_ptr->cfg.annotations)
        d_ptr->cfg.annotations->erase(key);
}

/**
 * @brief 设置挂载点
 * @param mounts 挂载点
 */
void KMContainer::setMounts(const std::vector<Mount> &mounts)
{
    d_ptr->cfg.mounts = mounts;
}

/**
 * @brief 添加挂载点
 * @param mount 挂载点
 */
void KMContainer::addMount(const Mount &mount)
{
    if (d_ptr->cfg.mounts)
    {
        d_ptr->cfg.mounts->emplace_back(mount);
    }
    else
    {
        d_ptr->cfg.mounts = std::vector<Mount>{ mount };
    }
}

/**
 * @brief 获取Mounts
 * @return mounts 挂载数据
 */
std::vector<Mount> KMContainer::getMounts()
{
    if (d_ptr->cfg.mounts)
    {
        return d_ptr->cfg.mounts.value();
    }
    else
    {
        return std::vector<Mount>{};
    }
}

/**
 * @brief 设置进程启动参数
 * @param args 启动参数
 */
void KMContainer::setProcessArgs(const std::vector<std::string> &args)
{
    if (!d_ptr->cfg.process)
    {
        d_ptr->cfg.process = Process{};
    }

    d_ptr->cfg.process->args = args;
}

/**
 * @brief 设置进程默认参数
 */
void KMContainer::initProcessDefault()
{
    if (!d_ptr->cfg.process)
    {
        d_ptr->cfg.process = Process{};
    }

    if (isatty(fileno(stdin)) != 0)
    {
        d_ptr->cfg.process->terminal = true;
    }
    else
    {
        d_ptr->cfg.process->terminal = false;
    }

    d_ptr->cfg.process->user = User{ .gid = getgid(), .uid = getuid() };

    // 工作目录默认为根目录
    d_ptr->cfg.process->cwd = string("/");
}

/**
 * @brief 获取进程启动参数
 * @return args 启动参数
 */
std::vector<std::string> KMContainer::getProcessArgs()
{
    if (d_ptr->cfg.process)
    {
        return d_ptr->cfg.process->args;
    }

    return std::vector<std::string>{};
}

/**
 * @brief 设置进程环境变量
 * @param env 环境变量
 */
void KMContainer::setProcessEnv(const std::vector<std::string> &env)
{
    if (!d_ptr->cfg.process)
        d_ptr->cfg.process = Process{};

    d_ptr->cfg.process->env = env;
}

/**
 * @brief 添加进程环境变量
 * @param env 环境变量
 */
void KMContainer::addProcessEnv(const std::vector<std::string> &env)
{
    if (!d_ptr->cfg.process)
        d_ptr->cfg.process = Process{};

    d_ptr->cfg.process->env.insert(d_ptr->cfg.process->env.end(), std::make_move_iterator(env.begin()), std::make_move_iterator(env.end()));
}

/**
 * @brief 获取进程启动环境变量
 * @return env 环境变量
 */
std::vector<std::string> KMContainer::getProcessEnv()
{
    if (d_ptr->cfg.process)
    {
        return d_ptr->cfg.process->env;
    }

    return std::vector<std::string>{};
}

/**
 * @brief 设置进程配置
 * @param process 进程配置
 */
void KMContainer::setProcess(const Process &process)
{
    d_ptr->cfg.process = process;
}

/**
 * @brief 获取进程配置
 * @return Process 进程配置
 */
Process KMContainer::getProcess()
{
    if (d_ptr->cfg.process)
    {
        return d_ptr->cfg.process.value();
    }

    return Process{};
}

/**
 * @brief 设置工作目录
 * @param cwd 工作目录
 */
void KMContainer::setCwd(const std::string &cwd)
{
    if (!d_ptr->cfg.process)
    {
        d_ptr->cfg.process = Process{};
    }

    d_ptr->cfg.process->cwd = cwd;
}

/**
 * @brief 设置hooks挂载
 * @param hooks hooks挂载
 */
void KMContainer::setHooks(const Hooks &hooks)
{
    d_ptr->cfg.hooks = hooks;
}

/**
 * @brief 获取Hooks挂载数据
 * @return hooks 挂载数据
 */
Hooks KMContainer::getHooks()
{
    if (d_ptr->cfg.hooks)
    {
        return d_ptr->cfg.hooks.value();
    }

    return Hooks{};
}

/**
 * @brief 设置prestart钩子
 * @param perstart prestart结构
 */
void KMContainer::setPrestart(const std::vector<Hook> &prestart)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }

    d_ptr->cfg.hooks->prestart = prestart;
}

/**
 * @brief 添加createRuntime钩子
 * @param createRuntime createRuntime结构
 */
void KMContainer::setCreateRuntime(const std::vector<Hook> &createRuntime)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }
    d_ptr->cfg.hooks->createRuntime = createRuntime;
}

/**
 * @brief 设置createContainer钩子
 * @param createContainer createContainer结构
 */
void KMContainer::setCreateContainer(const std::vector<Hook> &createContainer)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }
    d_ptr->cfg.hooks->createContainer = createContainer;
}

/**
 * @brief 添加startContainer钩子
 * @param startContainer startContainer结构
 */
void KMContainer::setStartContainer(const std::vector<Hook> &startContainer)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }
    d_ptr->cfg.hooks->startContainer = startContainer;
}

/**
 * @brief 设置poststart钩子
 * @param postStart postStart结构
 */
void KMContainer::setPoststart(const std::vector<Hook> &postStart)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }
    d_ptr->cfg.hooks->postStart = postStart;
}

/**
 * @brief 添加poststop钩子
 * @param postStop postStop结构
 */
void KMContainer::setPostStop(const std::vector<Hook> &postStop)
{
    if (!d_ptr->cfg.hooks)
    {
        d_ptr->cfg.hooks = Hooks{};
    }
    d_ptr->cfg.hooks->postStop = postStop;
}

/**
 * @brief 设置linux平台配置
 * @param linuxConfig linux配置
 */
void KMContainer::setLinux(const Linux &linuxConfig)
{
    d_ptr->cfg.linuxConfig = linuxConfig;
}

/**
 * @brief 获取linux平台配置
 * @return Linux 平台信息
 */
Linux KMContainer::getLinux()
{
    if (d_ptr->cfg.linuxConfig)
    {
        return d_ptr->cfg.linuxConfig.value();
    }
    return Linux{};
}

/**
 * @brief 运行容器应用
 * @param dir 工作目录
 * @param containerID 容器ID
 * @return 运行成功返回true，否则返回false
 */
bool KMContainer::run()
{
    string progress = KMUtil::getKMBox();
    if (!fs::exists(progress))
    {
        KMError("crun do not exist, please install crun first");
        return false;
    }

    string runDir{}, containerID{};
    if (!genRunDir(runDir, containerID))
    {
        KMError("generate run directory failed");
        return false;
    }

    if (!genConfig(runDir))
    {
        KMError("generate config.json failed");
        return false;
    }

    vector<string> args = { "--cgroup-manager=disabled", "run", "-b", runDir, containerID };
    if (!KMCrun::exec(args))
    {
        clearRunDir(runDir);
        return false;
    }

    if (!clearRunDir(runDir))
    {
        KMInfo("clear run directory failed");
    }

    return true;
}

/**
 * @brief 修正挂载信息
 * @return true 成功 false 失败
 */
bool KMContainer::amendMountInfo()
{
    if (!d_ptr->cfg.mounts.has_value())
    {
        KMError("config mounts is empty");
        return false;
    }

    // 生成tmpfs目录,解决只读挂载时，文件写入失败的问题
    std::vector<Mount> &mounts = d_ptr->cfg.mounts.value();
    string oriRoot = d_ptr->baseDir + "/files";
    std::unordered_set<std::string> tmpfs;
    for (auto pos = mounts.begin(); pos != mounts.end(); ++pos)
    {
        std::string destination = pos->destination;
        if (destination.empty() || destination[0] != '/')
        {
            continue;
        }
        if (pos->type == "tmpfs")
        {
            continue;
        }

        string hostDestination = oriRoot + destination;
        std::string tmpfsDir;
        if (!getTmpfsDir(hostDestination, tmpfsDir))
        {
            continue;
        }

        if (tmpfs.find(tmpfsDir) != tmpfs.end())
        {
            continue;
        }
        else
        {
            tmpfs.insert(tmpfsDir);
        }

        string dst = tmpfsDir.substr(oriRoot.size());
        pos = mounts.insert(pos, Mount{ .source = "tmpfs", .destination = dst, .type = "tmpfs", .options = { { "nodev", "nosuid", "mode=755" } } });
        ++pos;

        if (!fs::exists(tmpfsDir))
            continue;

        for (auto &entry : fs::directory_iterator(tmpfsDir))
        {
            string dst = entry.path().string().substr(oriRoot.size());
            string src = entry.path().string();
            string type = "bind";
            std::vector<std::string> options = { "rbind", "ro" };
            if (entry.is_symlink())
            {
                options.emplace_back("copy-symlink");
            }

            pos = mounts.insert(pos, Mount{ .source = src, .destination = dst, .type = type, .options = options });
            ++pos;
        }
    }

    // 去除重复挂载
    std::unordered_set<std::string> unordset;
    for (auto it = mounts.crbegin(); it != mounts.crend(); ++it)
    {
        if (unordset.find(it->destination) != unordset.end())
        {
            mounts.erase(std::next(it).base());
            continue;
        }
        unordset.insert(it->destination);
    }

    setMounts(mounts);

    return true;
}

/**
 * @brief 获取需要生成的tmpfs目录
 * @param dir 目录
 * @return true 成功 false 失败
 */
bool KMContainer::getTmpfsDir(const std::string dir, std::string &tmpfsDir)
{
    if (fs::exists(dir))
    {
        return false;
    }

    fs::path dirpath(dir);
    std::string lastDir = dirpath.parent_path().string();
    while (!fs::exists(dirpath.parent_path()))
    {
        dirpath = dirpath.parent_path();
    }

    string existDir = dirpath.parent_path().string();
    string baseroot = d_ptr->baseDir + string("/files");
    if (existDir <= baseroot)
    {
        return false;
    }

    tmpfsDir = existDir;

    return true;
}

/**
 * @brief 生成运行参数
 * @return true 成功，false 失败
 */
bool KMContainer::genRunDir(std::string &dir, std::string &uuid)
{
    uuid_t uuidtime;
    // 生成基于时间的UUID
    uuid_generate_time(uuidtime);

    // 将UUID转换为字符串格式
    std::array<char, 64> uuidarray;
    uuid_unparse(uuidtime, uuidarray.data());
    uuid = string(uuidarray.data());

    if (uuid.empty())
    {
        KMError("generate uuid failed");
        return false;
    }

    // 创建运行目录
    dir = string("/run/user/") + to_string(getuid()) + string("/kaiming/") + d_ptr->id + string("/") + uuid;

    return true;
}

/**
 * @brief 生成配置文件
 * @param dir bundle目录
 * @return true 成功，false 失败
 */
bool KMContainer::genConfig(const std::string &dir)
{
    if (!fs::exists(dir))
    {
        std::error_code ec;
        fs::create_directories(dir, ec);
        if (ec)
        {
            KMError("create bundle directory failed: " + ec.message());
            return false;
        }
    }

    string rootPath = dir + string("/") + d_ptr->cfg.root->path;
    if (!fs::exists(rootPath))
    {
        std::error_code ec;
        fs::create_directories(rootPath, ec);
    }

    std::ofstream ofs(dir + "/config.json");
    if (!ofs.is_open())
    {
        KMError("create config.json in bundle directory failed");
        return false;
    }

    try
    {
        nlohmann::json json = d_ptr->cfg;
        ofs << std::setw(4) << json << std::endl;
        ofs.close();
    }
    catch (...)
    {
        ofs.close();
        KMError("generate config.json failed");
        return false;
    }

    return true;
}

/**
 * @brief 清空容器运行目录
 * @param dir 运行目录
 * @return true 成功 false 失败
 */
bool KMContainer::clearRunDir(const std::string &dir)
{
    if (dir.empty())
    {
        return true;
    }

    if (!fs::exists(dir))
    {
        return true;
    }

    std::error_code ec;
    if (!fs::remove_all(dir, ec))
    {
        KMError("clear run directory failed: " + ec.message());
        return false;
    }

    return true;
}