#include "core_workers.h"

#include <algorithm>
#include <stdexcept>

#if defined(__linux__)
#include <pthread.h>
#include <sched.h>
#endif

namespace ling3 {

CoreWorkers & CoreWorkers::Instance() {
    static CoreWorkers workers;
    return workers;
}

CoreWorkers::CoreWorkers() {
    for (int core = 0; core < 4; ++core) {
        workers_[core] = std::thread([this, core]() { Worker(core); });
    }
}

CoreWorkers::~CoreWorkers() {
    {
        std::lock_guard lock(state_mutex_);
        stop_ = true;
    }
    start_.notify_all();
    for (auto & worker : workers_) if (worker.joinable()) worker.join();
}

void CoreWorkers::Worker(int core) {
#if defined(__linux__) && defined(__aarch64__)
    cpu_set_t affinity;
    CPU_ZERO(&affinity);
    CPU_SET(4 + core, &affinity);
    pthread_setaffinity_np(pthread_self(), sizeof(affinity), &affinity);
    const char * names[] = {"l3-npu0", "l3-npu1", "l3-npu2", "l3-cpu3"};
    pthread_setname_np(pthread_self(), names[core]);
#endif
    std::size_t observed = 0;
    while (true) {
        std::function<void(int)> task;
        {
            std::unique_lock lock(state_mutex_);
            start_.wait(lock, [this, observed]() {
                return stop_ || generation_ != observed;
            });
            if (stop_) return;
            observed = generation_;
            if (!active_[core]) continue;
            task = task_;
        }
        try {
            task(core);
        } catch (...) {
            errors_[core] = std::current_exception();
        }
        {
            std::lock_guard lock(state_mutex_);
            ++completed_;
        }
        done_.notify_one();
    }
}

void CoreWorkers::Run(
    std::span<const int> cores,
    const std::function<void(int)> & task) {
    if (cores.empty() || cores.size() > 4 || !task) {
        throw std::invalid_argument("CoreWorkers requires one to four active cores and a task");
    }
    std::array<bool, 4> selected {};
    for (int core : cores) {
        if (core < 0 || core > 3 || selected[core]) {
            throw std::invalid_argument("CoreWorkers cores must be unique values in [0, 3]");
        }
        selected[core] = true;
    }

    std::lock_guard run_lock(run_mutex_);
    {
        std::lock_guard state_lock(state_mutex_);
        active_ = selected;
        errors_.fill(nullptr);
        task_ = task;
        active_count_ = cores.size();
        completed_ = 0;
        ++generation_;
    }
    start_.notify_all();
    {
        std::unique_lock lock(state_mutex_);
        done_.wait(lock, [this]() { return completed_ == active_count_; });
        task_ = {};
    }
    for (int core : cores) {
        if (errors_[core] != nullptr) std::rethrow_exception(errors_[core]);
    }
}

} // namespace ling3