#pragma once

#include <cstdint>
#include <string>
#include <vector>

#include "utils/simple_json.h"

#if defined(__linux__)
#include <pthread.h>
#include <sched.h>
#elif defined(_WIN32)
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#endif

namespace rk3588 {

inline std::vector<int> ParseCpuAffinity(const SimpleJson& node_cfg) {
    std::vector<int> out;
    const SimpleJson* a = node_cfg.Find("cpu_affinity");
    if (!a || !a->IsArray()) return out;
    for (const auto& v : a->AsArray()) {
        if (!v.IsNumber()) continue;
        const int cpu = v.AsInt(-1);
        if (cpu >= 0) out.push_back(cpu);
    }
    return out;
}

inline bool SetCurrentThreadAffinity(const std::vector<int>& cpus, std::string& err) {
    err.clear();
    if (cpus.empty()) return true;

#if defined(__linux__)
    cpu_set_t set;
    CPU_ZERO(&set);
    for (int cpu : cpus) {
        if (cpu < 0 || cpu >= CPU_SETSIZE) continue;
        CPU_SET(cpu, &set);
    }
    const int rc = pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
    if (rc != 0) {
        err = "pthread_setaffinity_np failed";
        return false;
    }
    return true;
#elif defined(_WIN32)
    // Windows uses a bitmask; best-effort for up to 64 CPUs in the current group.
    DWORD_PTR mask = 0;
    for (int cpu : cpus) {
        if (cpu < 0 || cpu >= static_cast<int>(8 * sizeof(DWORD_PTR))) continue;
        mask |= (static_cast<DWORD_PTR>(1) << cpu);
    }
    if (mask == 0) return true;
    const DWORD_PTR prev = SetThreadAffinityMask(GetCurrentThread(), mask);
    if (prev == 0) {
        err = "SetThreadAffinityMask failed";
        return false;
    }
    return true;
#else
    (void)cpus;
    err = "thread affinity not supported on this platform";
    return false;
#endif
}

}  // namespace rk3588