/*
 * Copyright (c) KylinSoft  Co., Ltd. 2024. All rights reserved.
 *
 * kaiming is licensed under the GPL v2.0+.
 * 
 * See the LICENSE file for more details.
 */

#include "KMOKPackager.h"

#include <vector>
#include <optional>
#include <uuid/uuid.h>

#include "KMConfig.h"
#include "KMException.h"
#include "KMLogger.h"
#include "KMUtil.h"
#include "KMStringUtils.h"
#include "KMProcessUtils.h"
#include "KMFileUtils.h"
#include "KMDirMd5sum.h"

class KMOKPackager::Private
{
public:
    std::string m_buildDir;
    std::string m_cacheDir;
    bool m_isExportAll;
    std::string m_okPackage;
    std::string m_module;
    std::vector<std::string> m_modulePaths;

    std::string m_icon;
    std::string m_bundleFilesystem{ BUNDLE_FILESYSTE_EROFS };
    std::string m_compression;
    std::vector<Layer> m_layers;

    KMOABMetadata m_metadata;
};

KMOKPackager::KMOKPackager(const std::string &buildDir, bool exportAll)
    : d(std::make_shared<Private>())
{
    d->m_buildDir = buildDir;
    d->m_cacheDir = KMStringUtils::buildFilename(buildDir, ".kaiming-builder/package");
    d->m_isExportAll = exportAll;
}

KMOKPackager::~KMOKPackager()
{}

void KMOKPackager::setIcon(const std::string &icon)
{
    d->m_icon = icon;
}

void KMOKPackager::setBundleFilesystem(const std::string &filesystem)
{
    d->m_bundleFilesystem = filesystem;
}

void KMOKPackager::setCompression(const std::string &compression)
{
    d->m_compression = compression;
}

void KMOKPackager::setModulePaths(const std::string &moduleName, const std::string& paths)
{
    if (paths != "*" || paths != ".*")
    {
        std::vector<std::string> splits = KMStringUtils::splitString(paths,"\n");

        if (!KMStringUtils::contains(splits, "*") && !KMStringUtils::contains(splits, ".*"))
        {
            d->m_module = moduleName;
            d->m_modulePaths = KMStringUtils::splitString(paths,"\n");
        }
    }
}

void KMOKPackager::appendLayer(Layer layer)
{
    d->m_layers.push_back(layer);
}

void KMOKPackager::pack(const std::string &okFilename)
{
    KMTrace("KMOKPackager::pack invoke begin");

    std::string okEngine = KMStringUtils::buildFilename(KAIMING_TOOLS_BIN_PATH, "ok-engine");
    if (!fs::exists(okEngine))
    {
        throw KMException(okEngine + " is missing.");
    }

    if (!KMFileUtils::isEmptyDirectory(d->m_cacheDir))
    {
        // KMFileUtils::rmRfAll(d->m_cacheDir);
        fs::remove_all(d->m_cacheDir);
        KMInfo("clear cache : " + d->m_cacheDir);
    }
    KMInfo("mkpath : " + d->m_cacheDir);
    KMFileUtils::mkpath(d->m_cacheDir);

    d->m_okPackage = KMStringUtils::buildFilename(d->m_cacheDir, okFilename);
    KMDebug("delete old " + d->m_okPackage);
    std::error_code ec;
    if (fs::exists(d->m_okPackage) && !fs::remove(d->m_okPackage, ec))
    {
        throw KMException(ec.message());
    }

    KMDebug("copy file " + okEngine + " to " + d->m_okPackage);
    if (!fs::copy_file(okEngine, d->m_okPackage, ec))
    {
        throw KMException(ec.message());
    }

    if (!d->m_icon.empty())
    {
        packIcon();
    }

    packBundle();

    packMetadata();

    // 从缓存目录移动到目标目录
    std::string okPackage = KMStringUtils::buildFilename(d->m_buildDir, okFilename);
    if (fs::exists(okPackage) && !fs::remove(okPackage, ec))
    {
        throw KMException(ec.message());
    }

    std::string result;
    if (!KMProcessUtils::spawn("mv", { d->m_okPackage, okPackage }, result))
    {
        throw KMException(result);
    }

    fs::permissions(okPackage, fs::perms::all);

    KMFileUtils::rmRfAll(d->m_cacheDir);

    KMInfo("Packaging completed. Generate file: " + okPackage);

    KMTrace("KMOKPackager::pack invoke end");
}

void KMOKPackager::addNewSection(const std::string &sectionName, const std::string &dataFile)
{
    std::string result;
    std::vector<std::string> args;
    args.push_back("--add-section");
    args.push_back(sectionName + "=" + dataFile);
    args.push_back(d->m_okPackage);
    args.push_back(d->m_okPackage);

    if (!KMProcessUtils::spawn("objcopy", args, result))
    {
        throw KMException(result);
    }
    else if (!result.empty())
    {
        KMDebug(result);
    }
}

void KMOKPackager::packIcon()
{
    KMTrace("KMOKPackager::packIcon invoke begin");

    KMInfo("Packaging icom ...");

    std::string result;
    std::vector<std::string> args;
    args.push_back("q");
    std::string iconAchieve = KMStringUtils::buildFilename(d->m_cacheDir, "icon.a");
    args.push_back(iconAchieve);
    args.push_back(d->m_icon);

    if (!KMProcessUtils::spawn("ar", args, result))
    {
        throw KMException(result);
    }
    else if (!result.empty())
    {
        KMDebug(result);
    }

    std::string iconSection(KMOABMETADATA_PREFIX ".icon");
    addNewSection(iconSection, iconAchieve);

    d->m_metadata.sections.icon = iconSection;

    KMTrace("KMOKPackager::packIcon invoke end");
}

void KMOKPackager::packBundle()
{
    KMTrace("KMOKPackager::packBundle invoke begin");

    KMInfo("Packaging bundle ...");

    std::string bundleDir = KMStringUtils::buildFilename(d->m_cacheDir, "bundle");
    std::error_code ec;
    if (!KMFileUtils::mkpath(bundleDir, ec))
    {
        throw KMException(ec.message());
    }
    prepareBundle(bundleDir);

    std::string bundleFile;
    if (d->m_bundleFilesystem == BUNDLE_FILESYSTE_SQUASHFS)
    {
        bundleFile = KMStringUtils::buildFilename(d->m_cacheDir, "bundle.squashfs");
        generateSquashfsBundle(bundleFile, bundleDir);
    }
    else //if (d->m_bundleFilesystem == BUNDLE_FILESYSTE_EROFS)
    {
        bundleFile = KMStringUtils::buildFilename(d->m_cacheDir, "bundle.ef");
        generateErofsBundle(bundleFile, bundleDir);
    }

    // calculate digest
    std::string digest = KMFileUtils::getFileSha256sum(bundleFile);
    if (digest.empty())
    {
        throw KMException("Failed to sha256sum " + bundleFile);
    }
    d->m_metadata.digest = digest;

    std::string bundleSection(KMOABMETADATA_PREFIX ".bundle");
    addNewSection(bundleSection, bundleFile);

    d->m_metadata.sections.bundle = bundleSection;

    KMTrace("KMOKPackager::packBundle invoke end");
}

static void filterSubPathsForSplitPackage(const std::string& rootPath, const std::vector<std::string>& paths, const std::string& outputFile, const std::vector<std::string>& appendPaths) 
{
    std::ofstream outFile(outputFile, std::ios::trunc);

    if (!outFile.is_open()) 
    {
        throw KMException("Failed to open file: " + outputFile);
    }

    for (const auto& path : paths) 
    {
        fs::path targetPath = KMStringUtils::buildFilename(rootPath, "files", path);

        if (fs::exists(targetPath)) 
        {
            std::string relativePath = fs::relative(targetPath, rootPath).string();
            outFile << relativePath << "\n";
        } 
        else 
        {
            std::regex regexPattern(path);

            for (const auto& entry : fs::recursive_directory_iterator(rootPath + "/files")) 
            {
                if (std::regex_match(entry.path().string(), regexPattern)) 
                {
                    std::string relativePath = fs::relative(entry.path(), rootPath).string();
                    outFile << relativePath << "\n";
                }
            }
        }
    }

    for (auto& append : appendPaths)
    {
        outFile << append << "\n";
    }

    outFile.close();
}

void KMOKPackager::prepareBundle(const std::string &bundleDir)
{
    KMTrace("KMOKPackager::prepareBundle invoke begin");

    std::error_code ec;
    if (d->m_isExportAll)
    {
        // copy extra data
        std::string extraDir = KMStringUtils::buildFilename(bundleDir, "extra");
        if (!KMFileUtils::mkpath(extraDir, ec))
        {
            throw KMException(ec.message());
        }

        // copy crun
        const auto copyOptions = fs::copy_options::overwrite_existing;
        std::string crunFile = KMStringUtils::buildFilename(extraDir, "crun");
        fs::copy_file("/usr/bin/crun", crunFile, copyOptions, ec);
        if (0 == ec.value())
        {
            fs::permissions(crunFile, fs::perms::all);
        }
        else
        {
            KMWarn("Failed to copy /usr/bin/crun, error: " + ec.message());
        }
    }

    // export layers，m_isExportAll=true时，layers中包含base、runtime
    std::string layersDir = KMStringUtils::buildFilename(bundleDir, "layers");
    if (!KMFileUtils::mkpath(layersDir, ec))
    {
        throw KMException(ec.message());
    }
    for (Layer layer : d->m_layers)
    {
        KMInfo("Copying " + layer.m_info->id + " ...");

        if (layer.m_info->kind == BASE_TYPE_APP && !d->m_module.empty())
        {
            std::string appDir = KMStringUtils::buildFilename(layersDir, layer.m_info->id, layer.m_info->module);
            if (!KMFileUtils::mkpath(appDir, ec))
            {
                throw KMException(ec.message());
            }

            std::vector<std::string> appendPaths;
            std::vector<std::string> args;
            args.push_back("-rlpt");
            args.push_back("--info=progress2");
            args.push_back("--ignore-errors");
            if (layer.m_info->module == "binary")
            {
                appendPaths.push_back("entries");
            }
            else
            {
                layer.m_info->commands.clear();
            }
            std::string modulePathsFile = KMStringUtils::buildFilename(d->m_cacheDir, "paths.txt");
            filterSubPathsForSplitPackage(layer.m_layerDir, d->m_modulePaths, modulePathsFile, appendPaths);
            args.push_back("--files-from=" + modulePathsFile);
            args.push_back(layer.m_layerDir + "/");
            args.push_back(appDir + "/");

            if (!KMProcessUtils::spawn("rsync", args, true))
            {
                throw KMException("Failed to copy " + layer.m_layerDir + " to " + appDir);
            }
            KMInfo("\nSyncing, please wait a moment ...\n");
            KMProcessUtils::spawn("sync", {}, true);

            std::string metadataFile = KMStringUtils::buildFilename(appDir, "info.json");
            layer.m_info->size = KMFileUtils::getDirectorySize(fs::path(appDir));
            layer.m_info->saveFile(metadataFile);

            // 更新更精确
            layer.m_info->loadFile(metadataFile);
            layer.m_info->size = KMFileUtils::getDirectorySize(fs::path(appDir));
            layer.m_info->saveFile(metadataFile);

            // 求md5校验和
            KMDirMd5sum md5sum(appDir);
            md5sum.md5sum("files");
        }
        else
        {
            std::string appDir = KMStringUtils::buildFilename(layersDir, layer.m_info->id, layer.m_info->module);
            if (!KMFileUtils::mkpath(appDir, ec))
            {
                throw KMException(ec.message());
            }

            std::vector<std::string> args;
            args.push_back("-rlpt");
            args.push_back("--info=progress2");
            args.push_back("--ignore-errors");
            args.push_back("--exclude=deploy");
            args.push_back("--exclude=modules.json");
            args.push_back(layer.m_layerDir + "/");
            args.push_back(appDir + "/");
            if (!KMProcessUtils::spawn("rsync", args, true))
            {
                throw KMException("Failed to copy " + layer.m_layerDir + " to " + appDir);
            }

            KMInfo("\nSyncing, please wait a moment ...\n");
            KMProcessUtils::spawn("sync", {}, true);

            // 求md5校验和
            KMDirMd5sum md5sum(appDir);
            md5sum.md5sum("files");

            std::string metadataFile = KMStringUtils::buildFilename(appDir, "info.json");
            layer.m_info->size = KMFileUtils::getDirectorySize(fs::path(appDir));
            layer.m_info->saveFile(metadataFile);
        }

        d->m_metadata.layers.push_back(*(layer.m_info));
    }

    KMProcessUtils::spawn("sync", {}, true);

    KMTrace("KMOKPackager::prepareBundle invoke end");
}

void KMOKPackager::generateErofsBundle(const std::string &bundleFile, const std::string &bundleDir)
{
    KMTrace("KMOKPackager::generateErofsBundle invoke begin");

    KMInfo("Doing mkfs.erofs ...");

    std::vector<std::string> args;
    if (d->m_compression.empty())
    {
        args.push_back("-zlz4hc");
    }
    else
    {
        args.push_back("-z" + d->m_compression);
    }
    // args.push_back("-b4096");
    args.push_back("-d3");
    args.push_back(bundleFile);
    args.push_back(bundleDir);

    if (!KMProcessUtils::spawn("mkfs.erofs", args))
    {
        throw KMException("Failed to generate bundle file.");
    }

    uuid_t uuidtime;
    // 生成基于时间的UUID
    uuid_generate_time(uuidtime);
    // 将UUID转换为字符串格式
    std::array<char, 64> uuidarray;
    uuid_unparse(uuidtime, uuidarray.data());
    std::string uuid = std::string(uuidarray.data());

    d->m_metadata.uuid = uuid;

    KMTrace("KMOKPackager::generateErofsBundle invoke end");
}

void KMOKPackager::generateSquashfsBundle(const std::string &bundleFile, const std::string &bundleDir)
{
    KMTrace("KMOKPackager::generateSquashfsBundle invoke begin");

    KMInfo("Doing mksquashfs ...");

    std::vector<std::string> args;
    args.push_back(bundleDir);
    args.push_back(bundleFile);
    // args.push_back("-info");

    if (!d->m_compression.empty())
    {
        std::vector<std::string> splits = KMStringUtils::splitString(d->m_compression, " ");
        std::vector<std::string> comps{ "gzip", "lzo", "lz4", "xz", "zstd", "lzma"};
        if (splits.size() > 0 && KMStringUtils::contains(comps, splits.at(0)))
        {
            args.push_back("-comp");
            for (const std::string& item : splits)
            {
                args.push_back(item);
            }
        }
    }

    if (!KMProcessUtils::spawn("mksquashfs", args))
    {
        throw KMException("Failed to generate bundle file.");
    }

    uuid_t uuidtime;
    // 生成基于时间的UUID
    uuid_generate_time(uuidtime);
    // 将UUID转换为字符串格式
    std::array<char, 64> uuidarray;
    uuid_unparse(uuidtime, uuidarray.data());
    std::string uuid = std::string(uuidarray.data());

    d->m_metadata.uuid = uuid;

    KMTrace("KMOKPackager::generateSquashfsBundle invoke end");
}

void KMOKPackager::packMetadata()
{
    KMTrace("KMOKPackager::packMetadata invoke begin");

    KMInfo("Packaging meta ...");

    std::string metadataFile = KMStringUtils::buildFilename(d->m_cacheDir, "metadata.json");
    d->m_metadata.bundleFilesystem = d->m_bundleFilesystem;
    d->m_metadata.saveFile(metadataFile);

    std::string metaSection(KMOABMETADATA_PREFIX ".meta");
    addNewSection(metaSection, metadataFile);

    KMTrace("KMOKPackager::packMetadata invoke end");
}
