int sched_setaffinity(pid_t pid, size_t cpusetsize, const cpu_set_t *mask);

#include <fstream>
#include <string>
#include <map>
#include <vector>
#include <dirent.h>
#include <regex>

#include "CPUUtils.h"
#include "../log/Log.h"

#define TAG "CPUUtils"

const std::regex CPUUtils::cpuNamePattern("^cpu[0-9]+$");

/**
  * 获取CPU频率分组map。
  *
  * @return 返回非null的CPU频率分组map，key为频率，value为同频CPU列表。例如：{"888":["cpu0","cpu1"], "666":["cpu2","cpu3"]}。
  */
std::map<int, std::vector<std::string>, std::greater<>> CPUUtils::findCPUFreqMap() {
    const std::string cpuDirPath = "/sys/devices/system/cpu";
    std::map<int, std::vector<std::string>, std::greater<>> freqCPUMap;

    const char *dirPath = cpuDirPath.c_str();
    DIR *dir = opendir(cpuDirPath.c_str());
    if (!dir) {
        LOGE(TAG, "findCpuFreqMap, the system/cpu [%s] can be not found.", dirPath);
        return freqCPUMap;
    }

    // 遍历CPU目录下的CPU列表
    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        const std::string cpuFileName = entry->d_name;

        // 只看 CPU+数字的目录，这才是CPU信息目录
        if (!std::regex_match(cpuFileName, CPUUtils::cpuNamePattern)) {
            continue;
        }

        // 获取CPU最大频率文件
        std::string maxFreqFilePath = cpuDirPath;
        maxFreqFilePath.append("/");
        maxFreqFilePath.append(cpuFileName);
        maxFreqFilePath.append("/cpufreq/cpuinfo_max_freq");

        std::ifstream maxFreqFile(maxFreqFilePath);
        if (!maxFreqFile) {
            LOGE(TAG, "findCpuFreqMap, the cpuinfo_max_freq [%s] can be not found.", maxFreqFilePath.c_str());
            continue;
        }

        // 读取CPU最大频率，并按照频率分组记录
        std::string line;
        if (std::getline(maxFreqFile, line)) {
            try {
                int frequency = std::stoi(line);
                freqCPUMap[frequency].push_back(cpuFileName);
            } catch (const std::exception &e) {
                LOGE(TAG, "findCpuFreqMap, error occurred on parse frequency: %s.", e.what());
            }
        }
    }

    closedir(dir);
    return freqCPUMap;
}

/**
  * 获取大中小核CPU列表。
  *
  * @param level CPU等级，0为大核，1为中核，2为小核。
  * @param upward CPU序号是否向上兼容，例如如果2找不到则向上找1或者0。
  *
  * @return 返回非null的CPU列表，例如：["cpu0", "cpu1"]。
  */
std::vector<std::string> CPUUtils::findCPUListByLevel(const int level, const bool upward) {
    // 获取CPU频率分组
    const auto findCPUFreqMap = CPUUtils::findCPUFreqMap();
    if (findCPUFreqMap.empty()) {
        return {};
    }

    // 把std::map的key（频率）按照升序存到std::vector
    std::vector<int> freqList;
    freqList.reserve(findCPUFreqMap.size());
    for (const auto &pair: findCPUFreqMap) {
        freqList.push_back(pair.first);
    }

    // 如果指定级别不存在，考虑是否向上兼容
    const int size = static_cast<int>(freqList.size());
    const int freqIndex = size > level ? level : (upward ? (size - 1) : -1);

    // 若指定级别存在，则返回频率对应的CPU列表
    if (freqIndex > -1) {
        const int freq = freqList[freqIndex];
        auto it = findCPUFreqMap.find(freq);

        if (it != findCPUFreqMap.end()) {
            return it->second;
        }
    }
    return {};
}

#include <unistd.h>

int ThreadUtils::bindThreadToCPU(const int cpuNumber) {
    pid_t tid = gettid();

    ...;
}

void* thread_function(void* arg) {
    ...;
    return nullptr;
}

void startWork() {
    pthread_t thread;
    int result = pthread_create(&thread, nullptr, thread_function, nullptr);
    if (result != 0) {
        LOGI(TAG, "pthread_create failed: %d", result);
        return;
    }

    pid_t tid = pthread_gettid_np(thread);
    LOGI(TAG, "Created thread tid: %d", tid);

    pthread_join(thread, nullptr);
}

/**
  * 根据线程名称获取线程id。
  *
  * @param threadName 目标线程名称。
  *
  * @return 没找到返回-1，找到返回线程id。
  */
int ThreadUtils::findThreadId(const std::string &threadName) {
    const int pid = getpid();
    const std::string taskDirPath = "/proc/" + std::to_string(pid) + "/task";
    LOGD(TAG, "findThreadId on path: %s", taskDirPath.c_str());

    DIR *dir = opendir(taskDirPath.c_str());
    if (dir == nullptr) {
        LOGE(TAG, "findThreadId, dir [%s] can be not found.", taskDirPath.c_str());
        return -1;
    }

    // 遍历线程目录列表
    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string threadId = entry->d_name;
        if (threadId == "." || threadId == "..") {
            continue;
        }
        std::string statusFilePath = taskDirPath;
        statusFilePath.append("/");
        statusFilePath.append(threadId);
        statusFilePath.append("/status");
        std::ifstream statusFile(statusFilePath);
        if (statusFile) {
            std::string line;
            // 遍历status文件中的每一行，找到Name行，并对比名称是否和指定相同
            while (std::getline(statusFile, line)) {
                if (line.find("Name") != std::string::npos) {
                    size_t found = line.find(':');
                    if (found != std::string::npos) {
                        std::string name = Utils::trim(line.substr(found + 1));
                        if (threadName == name) {
                            closedir(dir);
                            LOGI(TAG,
                                 "findThreadId, id of [%s] has been found: %s.",
                                 threadName.c_str(),
                                 threadId.c_str()
                            );
                            // 找到目标线程，返回目标线程tid
                            return std::stoi(threadId);
                        }
                    }
                    break;
                }
            }
        }
    }
    closedir(dir);
    LOGE(TAG, "findThreadId, id of [%s] can be not found.", threadName.c_str());
    return -1;
}

// Utils.cpp
std::string Utils::trim(const std::string &str) {
    auto start = str.find_first_not_of(" \t");
    auto end = str.find_last_not_of(" \t");

    if (start == std::string::npos) {
        return "";
    }
    return str.substr(start, end - start + 1);
}

int sched_setaffinity(pid_t pid, size_t cpusetsize, const cpu_set_t *mask);

/**
  * 绑定当前线程到指定级别的CPU上。
  *
  * @param cpuNumber CPU序号。
  *
  * @return 绑定成功返回0，否则返回-1。
  */
int ThreadUtils::bindThreadToCPU(const int cpuNumber) {
    pid_t tid = gettid();

    int bind = ThreadUtils::bindThreadToCPU(tid, cpuNumber);
    return bind;
}

/**
  * 绑定指定线程到指定级别的CPU上。
  *
  * @param threadId 线程id。
  * @param cpuNumber CPU序号。
  *
  * @return 绑定成功返回0，否则返回-1。
  */
int ThreadUtils::bindThreadToCPU(const int threadId, const int cpuNumber) {
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(cpuNumber, &mask);
    long max_cpus = sysconf(_SC_NPROCESSORS_CONF);
    if (cpuNumber < 0 || cpuNumber >= max_cpus) {
        LOGE(TAG, "bindThreadToCPU, cpu %d is invalid.", cpuNumber);
        return false;
    }

    int bind = sched_setaffinity(threadId, sizeof(mask), &mask);
    if (bind == 0) {
        LOGI(TAG, "bindThreadToCPU, %d, success.", threadId);
    } else {
        LOGE(TAG, "bindThreadToCPU, %d, error: %d, %s", bind, errno, strerror(errno));
    }

    return bind;
}

/**
* 绑定当前线程到指定级别的CPU上。
*
* @param env
* @param clazz
* @param cpuLevel CPU级别，0为大核，1为中核，2为小核。
* @param upward CPU序号是否向上兼容，例如如果2找不到则向上找1或者0。
*
* @return 绑定成功返回0，否则返回-1。
*/
JNIEXPORT jint JNICALL Java_com_yanzhenjie_android_Harry_bindCurrentThreadToCPU(
    JNIEnv *env, jclass clazz, jint cpuLevel, jboolean upward) {

    // 获取指定级别的CPU列表
    const auto cpuList = CPUUtils::findCPUListByLevel(cpuLevel, upward);

    if (!cpuList.empty()) {
        const int cpuSize = static_cast<int>(cpuList.size());
        const int cpuIndex = Utils::randomInt(0, cpuSize - 1);
        const auto &cpu = cpuList[cpuIndex];
        // 去掉前面CPU，只获取num
        const int cpuNum = std::stoi(cpu.substr(3));

        // 设定CPU亲和性
        int bound = ThreadUtils::bindThreadToCPU(cpuNum);
        Log::i(env, TAG,
               "bindCurrentThreadToCPU, level-%s cpu %s: %s",
               {std::to_string(cpuLevel), std::to_string(cpuNum), std::to_string(bound)}
        );
        return bound;
    }

    return -1;
}

/**
 * 绑定指定线程到指定级别的CPU上。
 *
 * @param env
 * @param clazz
 * @param threadName 线程名。
 * @param cpuLevel CPU级别，0为大核，1为中核，2为小核。
 * @param upward CPU序号是否向上兼容，例如如果2找不到则向上找1或者0。
 *
 * @return 绑定成功返回0，否则返回-1。
 */
JNIEXPORT jint JNICALL Java_com_yanzhenjie_android_Harry_bindThreadToCPU(
    JNIEnv *env, jclass clazz, jstring threadName, jint cpuLevel, jboolean upward) {

    const char *chars = env->GetStringUTFChars(threadName, nullptr);
    std::string name(chars);
    env->ReleaseStringUTFChars(threadName, chars);

    // 根据传入线程名称获取线程tid
    const int threadId = ThreadUtils::findThreadId(name);
    // 根据传入CPU级别获取CPU列表
    const auto cpuList = CPUUtils::findCPUListByLevel(cpuLevel, upward);

    if (threadId > -1 && !cpuList.empty()) {
        const int cpuSize = static_cast<int>(cpuList.size());
        // 从对应级别的CPU列表随机选一个CPU
        const int cpuIndex = Utils::randomInt(0, cpuSize - 1);
        const auto &cpu = cpuList[cpuIndex];
        // 去掉前面CPU，只获取num
        const int cpuNum = std::stoi(cpu.substr(3));

        // 设定CPU亲和性
        int bound = ThreadUtils::bindThreadToCPU(threadId, cpuNum);
        Log::i(env, TAG,
               "bindThreadToCPU, bind [%s] to level-[%s] cpu %s : %s",
               {name, std::to_string(cpuLevel), std::to_string(cpuNum), std::to_string(bound)}
        );
        return bound;
    }

    return -1;
}

// Utils.cpp
int Utils::randomInt(const int min, const int max) {
    std::random_device rd;
    std::mt19937 gen(rd());

    std::uniform_int_distribution<> distr(min, max);

    int random_num = distr(gen);
    return random_num;
}