#include "Core/Threads.h"
#include <iostream>
#include <vector>
#include <cassert>

#ifdef __linux__
#include <pthread.h>
#include <sched.h>
#endif

using namespace uLib;

std::vector<int> GetAvailableCpus() {
    std::vector<int> available;
#ifdef __linux__
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    if (sched_getaffinity(0, sizeof(cpu_set_t), &cpuset) == 0) {
        for (int i = 0; i < CPU_SETSIZE; ++i) {
            if (CPU_ISSET(i, &cpuset)) {
                available.push_back(i);
            }
        }
    }
#endif
    return available;
}

class TestThread : public Thread {
public:
    void Run() override {
        Thread::Sleep(200);
    }
};

void TestThreadAffinity() {
    std::cout << "Testing Thread Affinity..." << std::endl;
#ifdef __linux__
    auto available = GetAvailableCpus();
    if (available.empty()) {
        std::cout << "  No CPUs available for affinity test, skipping." << std::endl;
        return;
    }
    int target_cpu = available[0];
    std::cout << "  Using CPU " << target_cpu << std::endl;

    TestThread t;
    t.Start();
    t.SetAffinity(target_cpu); 
    
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    int s = pthread_getaffinity_np(t.GetNativeHandle(), sizeof(cpu_set_t), &cpuset);
    if (s != 0) {
        std::cerr << "Error: pthread_getaffinity_np failed with code " << s << std::endl;
        assert(false);
    }
    assert(CPU_ISSET(target_cpu, &cpuset));
    
    t.Join();
    std::cout << "  Passed (Thread bound to CPU 0)." << std::endl;
#else
    std::cout << "  Affinity not supported on this OS, skipping." << std::endl;
#endif
}

void TestTeamAffinity() {
    std::cout << "Testing Team Affinity..." << std::endl;
#ifdef __linux__
    auto available = GetAvailableCpus();
    if (available.size() < 2) {
        std::cout << "  Not enough CPUs available for Team affinity test, skipping." << std::endl;
        return;
    }
    std::vector<int> cpus = {available[0], available[1]};
    std::cout << "  Using CPUs " << cpus[0] << ", " << cpus[1] << std::endl;

    Team team(2);
    team.SetAffinity(cpus);
    
    // We check affinity inside a parallel region
    #pragma omp parallel
    {
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
        int tid = Thread::GetThreadNum();
        int expected_cpu = cpus[tid % cpus.size()];
        assert(CPU_ISSET(expected_cpu, &cpuset));
    }
    std::cout << "  Passed (Team threads bound correctly)." << std::endl;
#else
    std::cout << "  Affinity not supported on this OS, skipping." << std::endl;
#endif
}

// Helper to get native handle if needed (oops, I forgot to add it to Thread class)
// I'll add GetNativeHandle() to Thread class in Threads.h

int main() {
    TestThreadAffinity();
    TestTeamAffinity();
    std::cout << "All Affinity tests finished!" << std::endl;
    return 0;
}