// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (c) 2025 NVIDIA Corporation.
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <unistd.h>
#include <signal.h>
#include <time.h>
#include <bpf/bpf.h>
#include <scx/common.h>
#include <sys/wait.h>
#include <sched.h>
#include <pthread.h>
#include "scx_test.h"
#include "dequeue.bpf.skel.h"

#define NUM_WORKERS 8
#define AFFINITY_HAMMER_MS 500

/*
 * Worker function that creates enqueue/dequeue events via CPU work and
 * sleep.
 */
static void worker_fn(int id)
{
	int i;
	volatile int sum = 0;

	for (i = 0; i < 1000; i++) {
		volatile int j;

		/* Do some work to trigger scheduling events */
		for (j = 0; j < 10000; j++)
			sum += j;

		/* Sleep to trigger dequeue */
		usleep(1000 + (id * 100));
	}

	exit(0);
}

/*
 * This thread changes workers' affinity from outside so that some changes
 * hit tasks while they are still in the scheduler's queue and trigger
 * property-change dequeues.
 */
static void *affinity_hammer_fn(void *arg)
{
	pid_t *pids = arg;
	cpu_set_t cpuset;
	int i = 0, n = NUM_WORKERS;
	struct timespec start, now;

	clock_gettime(CLOCK_MONOTONIC, &start);
	while (1) {
		int w = i % n;
		int cpu = (i / n) % 4;

		CPU_ZERO(&cpuset);
		CPU_SET(cpu, &cpuset);
		sched_setaffinity(pids[w], sizeof(cpuset), &cpuset);
		i++;

		/* Check elapsed time every 256 iterations to limit gettime cost */
		if ((i & 255) == 0) {
			long long elapsed_ms;

			clock_gettime(CLOCK_MONOTONIC, &now);
			elapsed_ms = (now.tv_sec - start.tv_sec) * 1000LL +
				     (now.tv_nsec - start.tv_nsec) / 1000000;
			if (elapsed_ms >= AFFINITY_HAMMER_MS)
				break;
		}
	}
	return NULL;
}

static enum scx_test_status run_scenario(struct dequeue *skel, u32 scenario,
					 const char *scenario_name)
{
	struct bpf_link *link;
	pid_t pids[NUM_WORKERS];
	pthread_t hammer;

	int i, status;
	u64 enq_start, deq_start,
	    dispatch_deq_start, change_deq_start, bpf_queue_full_start;
	u64 enq_delta, deq_delta,
	    dispatch_deq_delta, change_deq_delta, bpf_queue_full_delta;

	/* Set the test scenario */
	skel->bss->test_scenario = scenario;

	/* Record starting counts */
	enq_start = skel->bss->enqueue_cnt;
	deq_start = skel->bss->dequeue_cnt;
	dispatch_deq_start = skel->bss->dispatch_dequeue_cnt;
	change_deq_start = skel->bss->change_dequeue_cnt;
	bpf_queue_full_start = skel->bss->bpf_queue_full;

	link = bpf_map__attach_struct_ops(skel->maps.dequeue_ops);
	SCX_FAIL_IF(!link, "Failed to attach struct_ops for scenario %s", scenario_name);

	/* Fork worker processes to generate enqueue/dequeue events */
	for (i = 0; i < NUM_WORKERS; i++) {
		pids[i] = fork();
		SCX_FAIL_IF(pids[i] < 0, "Failed to fork worker %d", i);

		if (pids[i] == 0) {
			worker_fn(i);
			/* Should not reach here */
			exit(1);
		}
	}

	/*
	 * Run an "affinity hammer" so that some property changes hit tasks
	 * while they are still in BPF custody (e.g., in user DSQ or BPF
	 * queue), triggering SCX_DEQ_SCHED_CHANGE dequeues.
	 */
	SCX_FAIL_IF(pthread_create(&hammer, NULL, affinity_hammer_fn, pids) != 0,
		    "Failed to create affinity hammer thread");
	pthread_join(hammer, NULL);

	/* Wait for all workers to complete */
	for (i = 0; i < NUM_WORKERS; i++) {
		SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i],
			    "Failed to wait for worker %d", i);
		SCX_FAIL_IF(status != 0, "Worker %d exited with status %d", i, status);
	}

	bpf_link__destroy(link);

	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_UNREG));

	/* Calculate deltas */
	enq_delta = skel->bss->enqueue_cnt - enq_start;
	deq_delta = skel->bss->dequeue_cnt - deq_start;
	dispatch_deq_delta = skel->bss->dispatch_dequeue_cnt - dispatch_deq_start;
	change_deq_delta = skel->bss->change_dequeue_cnt - change_deq_start;
	bpf_queue_full_delta = skel->bss->bpf_queue_full - bpf_queue_full_start;

	printf("%s:\n", scenario_name);
	printf("  enqueues: %lu\n", (unsigned long)enq_delta);
	printf("  dequeues: %lu (dispatch: %lu, property_change: %lu)\n",
	       (unsigned long)deq_delta,
	       (unsigned long)dispatch_deq_delta,
	       (unsigned long)change_deq_delta);
	printf("  BPF queue full: %lu\n", (unsigned long)bpf_queue_full_delta);

	/*
	 * Validate enqueue/dequeue lifecycle tracking.
	 *
	 * For scenarios 0, 1, 3, 4 (local and global DSQs from
	 * ops.select_cpu() and ops.enqueue()), both enqueues and dequeues
	 * should be 0 because tasks bypass the BPF scheduler entirely:
	 * tasks never enter BPF scheduler's custody.
	 *
	 * For scenarios 2, 5, 6 (user DSQ or BPF internal queue) we expect
	 * both enqueues and dequeues.
	 *
	 * The BPF code does strict state machine validation with
	 * scx_bpf_error() to ensure the workflow semantics are correct.
	 *
	 * If we reach this point without errors, the semantics are
	 * validated correctly.
	 */
	if (scenario == 0 || scenario == 1 ||
	    scenario == 3 || scenario == 4) {
		/* Tasks bypass BPF scheduler completely */
		SCX_EQ(enq_delta, 0);
		SCX_EQ(deq_delta, 0);
		SCX_EQ(dispatch_deq_delta, 0);
		SCX_EQ(change_deq_delta, 0);
	} else {
		/*
		 * User DSQ from ops.enqueue() or ops.select_cpu(): tasks
		 * enter BPF scheduler's custody.
		 *
		 * Also validate 1:1 enqueue/dequeue pairing.
		 */
		SCX_GT(enq_delta, 0);
		SCX_GT(deq_delta, 0);
		SCX_EQ(enq_delta, deq_delta);
	}

	return SCX_TEST_PASS;
}

static enum scx_test_status setup(void **ctx)
{
	struct dequeue *skel;

	skel = dequeue__open();
	SCX_FAIL_IF(!skel, "Failed to open skel");
	SCX_ENUM_INIT(skel);
	SCX_FAIL_IF(dequeue__load(skel), "Failed to load skel");

	*ctx = skel;

	return SCX_TEST_PASS;
}

static enum scx_test_status run(void *ctx)
{
	struct dequeue *skel = ctx;
	enum scx_test_status status;

	status = run_scenario(skel, 0, "Scenario 0: Local DSQ from ops.select_cpu()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 1, "Scenario 1: Global DSQ from ops.select_cpu()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 2, "Scenario 2: User DSQ from ops.select_cpu()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 3, "Scenario 3: Local DSQ from ops.enqueue()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 4, "Scenario 4: Global DSQ from ops.enqueue()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 5, "Scenario 5: User DSQ from ops.enqueue()");
	if (status != SCX_TEST_PASS)
		return status;

	status = run_scenario(skel, 6, "Scenario 6: BPF queue from ops.enqueue()");
	if (status != SCX_TEST_PASS)
		return status;

	printf("\n=== Summary ===\n");
	printf("Total enqueues: %lu\n", (unsigned long)skel->bss->enqueue_cnt);
	printf("Total dequeues: %lu\n", (unsigned long)skel->bss->dequeue_cnt);
	printf("  Dispatch dequeues: %lu (no flag, normal workflow)\n",
	       (unsigned long)skel->bss->dispatch_dequeue_cnt);
	printf("  Property change dequeues: %lu (SCX_DEQ_SCHED_CHANGE flag)\n",
	       (unsigned long)skel->bss->change_dequeue_cnt);
	printf("  BPF queue full: %lu\n",
	       (unsigned long)skel->bss->bpf_queue_full);
	printf("\nAll scenarios passed - no state machine violations detected\n");
	printf("-> Validated: Local DSQ dispatch bypasses BPF scheduler\n");
	printf("-> Validated: Global DSQ dispatch bypasses BPF scheduler\n");
	printf("-> Validated: User DSQ dispatch triggers ops.dequeue() callbacks\n");
	printf("-> Validated: Dispatch dequeues have no flags (normal workflow)\n");
	printf("-> Validated: Property change dequeues have SCX_DEQ_SCHED_CHANGE flag\n");
	printf("-> Validated: No duplicate enqueues or invalid state transitions\n");

	return SCX_TEST_PASS;
}

static void cleanup(void *ctx)
{
	struct dequeue *skel = ctx;

	dequeue__destroy(skel);
}

struct scx_test dequeue_test = {
	.name = "dequeue",
	.description = "Verify ops.dequeue() semantics",
	.setup = setup,
	.run = run,
	.cleanup = cleanup,
};

REGISTER_SCX_TEST(&dequeue_test)