// SPDX-License-Identifier: GPL-2.0
/*
 * Test for DSQ operations including create, destroy, and peek operations.
 *
 * Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2025 Ryan Newton <ryan.newton@alum.mit.edu>
 */
#include <bpf/bpf.h>
#include <scx/common.h>
#include <sys/wait.h>
#include <unistd.h>
#include <pthread.h>
#include <string.h>
#include <sched.h>
#include "peek_dsq.bpf.skel.h"
#include "scx_test.h"

#define NUM_WORKERS 4

static bool workload_running = true;
static pthread_t workload_threads[NUM_WORKERS];

/**
 * Background workload thread that sleeps and wakes rapidly to exercise
 * the scheduler's enqueue operations and ensure DSQ operations get tested.
 */
static void *workload_thread_fn(void *arg)
{
        while (workload_running) {
                /* Sleep for a very short time to trigger scheduler activity */
                usleep(1000); /* 1ms sleep */
                /* Yield to ensure we go through the scheduler */
                sched_yield();
        }
        return NULL;
}

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

        skel = peek_dsq__open();
        SCX_FAIL_IF(!skel, "Failed to open");
        SCX_ENUM_INIT(skel);
        SCX_FAIL_IF(peek_dsq__load(skel), "Failed to load skel");

        *ctx = skel;

        return SCX_TEST_PASS;
}

static int print_observed_pids(struct bpf_map *map, int max_samples, const char *dsq_name)
{
        long count = 0;

        printf("Observed %s DSQ peek pids:\n", dsq_name);
        for (int i = 0; i < max_samples; i++) {
                long pid;
                int err;

                err = bpf_map_lookup_elem(bpf_map__fd(map), &i, &pid);
                if (err == 0) {
                        if (pid == 0) {
                                printf("  Sample %d: NULL peek\n", i);
                        } else if (pid > 0) {
                                printf("  Sample %d: pid %ld\n", i, pid);
                                count++;
                        }
                } else {
                        printf("  Sample %d: error reading pid (err=%d)\n", i, err);
                }
        }
        printf("Observed ~%ld pids in the %s DSQ(s)\n", count, dsq_name);
        return count;
}

static enum scx_test_status run(void *ctx)
{
        struct peek_dsq *skel = ctx;
        bool failed = false;
        int seconds = 3;
        int err;

        /* Enable the scheduler to test DSQ operations */
        printf("Enabling scheduler to test DSQ insert operations...\n");

        struct bpf_link *link =
                bpf_map__attach_struct_ops(skel->maps.peek_dsq_ops);

        if (!link) {
                SCX_ERR("Failed to attach struct_ops");
                return SCX_TEST_FAIL;
        }

        printf("Starting %d background workload threads...\n", NUM_WORKERS);
        workload_running = true;
        for (int i = 0; i < NUM_WORKERS; i++) {
                err = pthread_create(&workload_threads[i], NULL, workload_thread_fn, NULL);
                if (err) {
                        SCX_ERR("Failed to create workload thread %d: %s", i, strerror(err));
                        /* Stop already created threads */
                        workload_running = false;
                        for (int j = 0; j < i; j++)
                                pthread_join(workload_threads[j], NULL);
                        bpf_link__destroy(link);
                        return SCX_TEST_FAIL;
                }
        }

        printf("Waiting for enqueue events.\n");
        sleep(seconds);
        while (skel->data->enqueue_count <= 0) {
                printf(".");
                fflush(stdout);
                sleep(1);
                seconds++;
                if (seconds >= 30) {
                        printf("\n\u2717 Timeout waiting for enqueue events\n");
                        /* Stop workload threads and cleanup */
                        workload_running = false;
                        for (int i = 0; i < NUM_WORKERS; i++)
                                pthread_join(workload_threads[i], NULL);
                        bpf_link__destroy(link);
                        return SCX_TEST_FAIL;
                }
        }

        workload_running = false;
        for (int i = 0; i < NUM_WORKERS; i++) {
                err = pthread_join(workload_threads[i], NULL);
                if (err) {
                        SCX_ERR("Failed to join workload thread %d: %s", i, strerror(err));
                        bpf_link__destroy(link);
                        return SCX_TEST_FAIL;
                }
        }
        printf("Background workload threads stopped.\n");

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

        /* Detach the scheduler */
        bpf_link__destroy(link);

        printf("Enqueue/dispatch count over %d seconds: %d / %d\n", seconds,
                skel->data->enqueue_count, skel->data->dispatch_count);
        printf("Debug: ksym_exists=%d\n",
               skel->bss->debug_ksym_exists);

        /* Check DSQ insert result */
        printf("DSQ insert test done on cpu: %d\n", skel->data->insert_test_cpu);
        if (skel->data->insert_test_cpu != -1)
                printf("\u2713 DSQ insert succeeded !\n");
        else {
                printf("\u2717 DSQ insert failed or not attempted\n");
                failed = true;
        }

        /* Check DSQ peek results */
        printf("  DSQ peek result 1 (before insert): %d\n",
               skel->data->dsq_peek_result1);
        if (skel->data->dsq_peek_result1 == 0)
                printf("\u2713 DSQ peek verification success: peek returned NULL!\n");
        else {
                printf("\u2717 DSQ peek verification failed\n");
                failed = true;
        }

        printf("  DSQ peek result 2 (after insert): %ld\n",
               skel->data->dsq_peek_result2);
        printf("  DSQ peek result 2, expected: %ld\n",
               skel->data->dsq_peek_result2_expected);
        if (skel->data->dsq_peek_result2 ==
            skel->data->dsq_peek_result2_expected)
                printf("\u2713 DSQ peek verification success: peek returned the inserted task!\n");
        else {
                printf("\u2717 DSQ peek verification failed\n");
                failed = true;
        }

        printf("  Inserted test task -> pid: %ld\n", skel->data->dsq_inserted_pid);
        printf("  DSQ peek result 2 -> pid: %ld\n", skel->data->dsq_peek_result2_pid);

        int pid_count;

        pid_count = print_observed_pids(skel->maps.peek_results,
                                        skel->data->max_samples, "DSQ pool");
        printf("Total non-null peek observations: %ld out of %ld\n",
               skel->data->successful_peeks, skel->data->total_peek_attempts);

        if (skel->bss->debug_ksym_exists && pid_count == 0) {
                printf("\u2717 DSQ pool test failed: no successful peeks in native mode\n");
                failed = true;
        }
        if (skel->bss->debug_ksym_exists && pid_count > 0)
                printf("\u2713 DSQ pool test success: observed successful peeks in native mode\n");

        if (failed)
                return SCX_TEST_FAIL;
        else
                return SCX_TEST_PASS;
}

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

        if (workload_running) {
                workload_running = false;
                for (int i = 0; i < NUM_WORKERS; i++)
                        pthread_join(workload_threads[i], NULL);
        }

        peek_dsq__destroy(skel);
}

struct scx_test peek_dsq = {
        .name = "peek_dsq",
        .description =
                "Test DSQ create/destroy operations and future peek functionality",
        .setup = setup,
        .run = run,
        .cleanup = cleanup,
};
REGISTER_SCX_TEST(&peek_dsq)