#include <socks/cpu.h>
#include <socks/percpu.h>
#include <socks/bitmap.h>

DECLARE_BITMAP(cpu_available, CPU_MAX);
DECLARE_BITMAP(cpu_online, CPU_MAX);

DEFINE_PERCPU_VAR(struct cpu_data, cpu_data);

struct cpu_data *get_this_cpu(void)
{
	return percpu_get(&cpu_data);
}

struct cpu_data *get_cpu(unsigned int id)
{
	return percpu_get_from(id, &cpu_data);
}

void put_cpu(struct cpu_data *cpu)
{
	percpu_put(cpu);
}

bool cpu_is_available(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return false;
	}

	return bitmap_check(cpu_available, cpu_id);
}

bool cpu_is_online(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return false;
	}

	return bitmap_check(cpu_online, cpu_id);
}

void cpu_set_available(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return;
	}

	bitmap_set(cpu_available, cpu_id);
}

void cpu_set_online(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return;
	}

	bitmap_set(cpu_online, cpu_id);
}

unsigned int cpu_nr_available(void)
{
	return bitmap_count_set(cpu_available, CPU_MAX);
}

unsigned int cpu_nr_online(void)
{
	return bitmap_count_set(cpu_online, CPU_MAX);
}

void preempt_disable(void)
{
	ml_cpu_block *ml_cpu = ml_this_cpu();
	struct cpu_data *cpu_data = ml_cpu_block_get_data(ml_cpu);
	if (!cpu_data) {
		return;
	}

	int preempt = READ_ONCE(cpu_data->c_preempt_count);
	preempt++;
	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
}

void preempt_enable(void)
{
	ml_cpu_block *ml_cpu = ml_this_cpu();
	struct cpu_data *cpu_data = ml_cpu_block_get_data(ml_cpu);
	if (!cpu_data) {
		return;
	}

	int preempt = READ_ONCE(cpu_data->c_preempt_count);

	if (preempt > 0) {
		preempt--;
	}

	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
}

int preempt_count(void)
{
	ml_cpu_block *ml_cpu = ml_this_cpu();
	struct cpu_data *cpu_data = ml_cpu_block_get_data(ml_cpu);
	if (!cpu_data) {
		return 0;
	}

	return READ_ONCE(cpu_data->c_preempt_count);
}

unsigned int cpu_get_highest_available(void)
{
	return bitmap_highest_set(cpu_available, CPU_MAX);
}

cycles_t __weak get_cycles(void)
{
	return 0;
}