/*
 * ekp_thread_pool.c - Work-stealing thread pool
 *
 * Copyright (C) 2024-2026 Kinney Zhang
 * SPDX-License-Identifier: GPL-3.0-or-later
 *
 * This file is part of emacs-kp, which is free software: you can
 * redistribute it and/or modify it under the terms of the GNU General
 * Public License as published by the Free Software Foundation, either
 * version 3 of the License, or (at your option) any later version.
 * It is distributed WITHOUT ANY WARRANTY; see the GNU General Public
 * License (COPYING) for details.
 *
 * Simple but effective: fixed thread count, lock-free queue would be
 * overkill for our batch sizes. Keep it simple, stupid.
 */

#include "ekp_module.h"
#include <stdlib.h>
#include <string.h>
#if !defined(_WIN32)
#include <unistd.h>
#endif

#define QUEUE_CAPACITY 1024

static void *worker_thread(void *arg)
{
    ekp_thread_pool_t *pool = (ekp_thread_pool_t *)arg;

    while (1) {
        pthread_mutex_lock(&pool->queue_lock);

        /* Wait for work */
        while (pool->queue_head == pool->queue_tail && !pool->shutdown) {
            pthread_cond_wait(&pool->queue_cond, &pool->queue_lock);
        }

        if (pool->shutdown && pool->queue_head == pool->queue_tail) {
            pthread_mutex_unlock(&pool->queue_lock);
            break;
        }

        /* Dequeue work */
        void (*func)(void *) = pool->queue[pool->queue_head].func;
        void *work_arg = pool->queue[pool->queue_head].arg;
        bool was_full =
            ((pool->queue_tail + 1) % pool->queue_size) == pool->queue_head;
        pool->queue_head = (pool->queue_head + 1) % pool->queue_size;
        pool->active_count++;
        if (was_full)
            pthread_cond_broadcast(&pool->done_cond);

        pthread_mutex_unlock(&pool->queue_lock);

        /* Execute */
        if (func)
            func(work_arg);

        /* Mark done */
        pthread_mutex_lock(&pool->queue_lock);
        pool->active_count--;
        if (pool->active_count == 0 && pool->queue_head == pool->queue_tail) {
            pthread_cond_signal(&pool->done_cond);
        }
        pthread_mutex_unlock(&pool->queue_lock);
    }

    return NULL;
}

size_t ekp_pool_default_threads(void)
{
    long n = 0;
#if defined(_SC_NPROCESSORS_ONLN)
    n = sysconf(_SC_NPROCESSORS_ONLN);
#endif
    if (n <= 0)
        n = 4;
    if (n > EKP_THREAD_POOL_MAX)
        n = EKP_THREAD_POOL_MAX;
    return (size_t)n;
}

ekp_thread_pool_t *ekp_pool_create(size_t num_threads)
{
    if (num_threads == 0)
        num_threads = ekp_pool_default_threads();
    if (num_threads > EKP_THREAD_POOL_MAX)
        num_threads = EKP_THREAD_POOL_MAX;

    ekp_thread_pool_t *pool = calloc(1, sizeof(*pool));
    if (!pool)
        return NULL;

    pool->queue_size = QUEUE_CAPACITY;
    pool->queue = calloc(pool->queue_size, sizeof(pool->queue[0]));
    if (!pool->queue) {
        free(pool);
        return NULL;
    }

    pthread_mutex_init(&pool->queue_lock, NULL);
    pthread_cond_init(&pool->queue_cond, NULL);
    pthread_cond_init(&pool->done_cond, NULL);

    /* Start worker threads */
    for (size_t i = 0; i < num_threads; i++) {
        if (pthread_create(&pool->threads[i], NULL, worker_thread, pool) != 0) {
            /* Cleanup on failure */
            pool->shutdown = true;
            pthread_cond_broadcast(&pool->queue_cond);
            for (size_t j = 0; j < i; j++) {
                pthread_join(pool->threads[j], NULL);
            }
            pthread_mutex_destroy(&pool->queue_lock);
            pthread_cond_destroy(&pool->queue_cond);
            pthread_cond_destroy(&pool->done_cond);
            free(pool->queue);
            free(pool);
            return NULL;
        }
    }

    pool->thread_count = num_threads;
    return pool;
}

void ekp_pool_destroy(ekp_thread_pool_t *pool)
{
    if (!pool)
        return;

    pthread_mutex_lock(&pool->queue_lock);
    pool->shutdown = true;
    pthread_cond_broadcast(&pool->queue_cond);
    pthread_mutex_unlock(&pool->queue_lock);

    for (size_t i = 0; i < pool->thread_count; i++) {
        pthread_join(pool->threads[i], NULL);
    }

    pthread_mutex_destroy(&pool->queue_lock);
    pthread_cond_destroy(&pool->queue_cond);
    pthread_cond_destroy(&pool->done_cond);
    free(pool->queue);
    free(pool);
}

void ekp_pool_submit(ekp_thread_pool_t *pool, void (*func)(void *), void *arg)
{
    if (!pool || !func)
        return;

    pthread_mutex_lock(&pool->queue_lock);

    /* Queue full: wait for a worker to make room.  Dropping the task
     * here used to silently degrade the batch to the Elisp fallback
     * exactly when parallelism mattered most. */
    while (((pool->queue_tail + 1) % pool->queue_size) == pool->queue_head
           && !pool->shutdown) {
        pthread_cond_wait(&pool->done_cond, &pool->queue_lock);
    }
    if (pool->shutdown) {
        pthread_mutex_unlock(&pool->queue_lock);
        return;
    }

    pool->queue[pool->queue_tail].func = func;
    pool->queue[pool->queue_tail].arg = arg;
    pool->queue_tail = (pool->queue_tail + 1) % pool->queue_size;

    pthread_cond_signal(&pool->queue_cond);
    pthread_mutex_unlock(&pool->queue_lock);
}

void ekp_pool_wait(ekp_thread_pool_t *pool)
{
    if (!pool)
        return;

    pthread_mutex_lock(&pool->queue_lock);
    while (pool->active_count > 0 || pool->queue_head != pool->queue_tail) {
        pthread_cond_wait(&pool->done_cond, &pool->queue_lock);
    }
    pthread_mutex_unlock(&pool->queue_lock);
}