#include "pthrpoll.h"
#include "common.h"

// 初始化任务队列
static int task_queue_init(TaskQueue *queue) {
    queue->head = NULL;
    queue->tail = NULL;
    queue->size = 0;
    if (pthread_mutex_init(&queue->mutex, NULL) != 0) {
        return -1;
    }
    if (pthread_cond_init(&queue->cond, NULL) != 0) {
        pthread_mutex_destroy(&queue->mutex);
        return -1;
    }
    return 0;
}

// 向任务队列中添加任务
// 注意：现在不负责复制参数，调用者需要保证 argument 的生命周期
static int task_queue_add(TaskQueue *queue, void (*function)(void *), void *argument) {
    Task *task = (Task *)malloc(sizeof(Task));
    if (task == NULL) {
        return -1;
    }
    
    task->function = function;
    task->argument = argument;
    task->next = NULL;

    pthread_mutex_lock(&queue->mutex);
    
    if (queue->tail == NULL) {
        queue->head = queue->tail = task;
    } else {
        queue->tail->next = task;
        queue->tail = task;
    }
    queue->size++;
    
    // 唤醒一个等待的线程
    pthread_cond_signal(&queue->cond);
    pthread_mutex_unlock(&queue->mutex);
    return 0;
}

// 从任务队列中取出任务
static Task *task_queue_remove(TaskQueue *queue, ThreadPool_t *pool) {
    pthread_mutex_lock(&queue->mutex);
    
    // 如果队列为空且线程池未关闭，则等待
    while (queue->size == 0 && !pool->shutdown) {
        pthread_cond_wait(&queue->cond, &queue->mutex);
    }

    // 如果线程池已关闭且队列为空，则线程可以退出了
    if (pool->shutdown && queue->size == 0) {
        pthread_mutex_unlock(&queue->mutex);
        return NULL;
    }

    Task *task = queue->head;
    queue->head = queue->head->next;
    
    if (queue->head == NULL) {
        queue->tail = NULL;
    }
    queue->size--;
    
    pthread_mutex_unlock(&queue->mutex);
    return task;
}

// 工作线程函数
static void *worker_thread(void *arg) {
    ThreadPool_t *pool = (ThreadPool_t *)arg;
    
    while (1) {
        // 从队列获取任务，如果返回 NULL，说明需要退出
        Task *task = task_queue_remove(&pool->task_queue, pool);
        if (task == NULL) {
            break; 
        }
        
        // 执行任务
        task->function(task->argument);
        // 释放任务结构体本身
        free(task);
        // 注意：task->argument 的释放由任务函数内部或调用者负责
    }
    
    // 线程退出前，减少存活线程计数
    pthread_mutex_lock(&pool->mutex);
    pool->live_threads--;
    pthread_mutex_unlock(&pool->mutex);
    
    return NULL;
}

// 初始化线程池 (修正版，非单例)
ThreadPool_t *thread_pool_init(int min_threads, int max_threads) {
    if (min_threads <= 0 || max_threads <= 0 || min_threads > max_threads) {
        errno = EINVAL;
        return NULL;
    }

    ThreadPool_t *pool = (ThreadPool_t *)malloc(sizeof(ThreadPool_t));
    if (pool == NULL) {
        return NULL;
    }

    if (task_queue_init(&pool->task_queue) != 0) {
        free(pool);
        return NULL;
    }

    pool->threads = (pthread_t *)malloc(max_threads * sizeof(pthread_t));
    if (pool->threads == NULL) {
        pthread_mutex_destroy(&pool->task_queue.mutex);
        pthread_cond_destroy(&pool->task_queue.cond);
        free(pool);
        return NULL;
    }

    pool->min_threads = min_threads;
    pool->max_threads = max_threads;
    pool->live_threads = 0;   // 初始为0，创建一个增加一个
    pool->busy_threads = 0;   // 可以增加这个来跟踪繁忙线程（本次简化实现中未使用）
    pool->shutdown = 0;
    
    if (pthread_mutex_init(&pool->mutex, NULL) != 0) {
        pthread_mutex_destroy(&pool->task_queue.mutex);
        pthread_cond_destroy(&pool->task_queue.cond);
        free(pool->threads);
        free(pool);
        return NULL;
    }

    // 创建核心（最小）数量的工作线程
    for (int i = 0; i < min_threads; i++) {
        if (pthread_create(&pool->threads[i], NULL, worker_thread, pool) != 0) {
            // 如果创建失败，销毁已创建的资源并返回
            thread_pool_destroy(pool);
            return NULL;
        }
        pool->live_threads++;
    }
    
    return pool;
}

// 向线程池添加任务
int thread_pool_add_task(ThreadPool_t *pool, void (*function)(void *), void *argument) {
    if (pool == NULL || function == NULL) {
        return -1;
    }

    pthread_mutex_lock(&pool->mutex);
    if (pool->shutdown) {
        pthread_mutex_unlock(&pool->mutex);
        return -1;
    }
    
    // 动态扩容：如果当前任务数大于存活线程数，且尚未达到最大线程数
    // 这是一个简单的扩容策略，可以根据需要调整
    if (pool->task_queue.size > 0 && pool->live_threads < pool->max_threads) {
        // 尝试创建一个新线程
        if (pthread_create(&pool->threads[pool->live_threads], NULL, worker_thread, pool) == 0) {
            pool->live_threads++;
        }
        // 如果创建失败，也没关系，现有线程会处理任务
    }
    pthread_mutex_unlock(&pool->mutex);

    // 直接添加任务，不再复制参数
    return task_queue_add(&pool->task_queue, function, argument);
}

// 销毁线程池
void thread_pool_destroy(ThreadPool_t *pool) {
    if (pool == NULL) {
        return;
    }

    pthread_mutex_lock(&pool->mutex);
    // 如果已经调用过销毁，直接返回
    if (pool->shutdown) {
        pthread_mutex_unlock(&pool->mutex);
        return;
    }
    pool->shutdown = 1;
    pthread_mutex_unlock(&pool->mutex);

    // 唤醒所有可能在等待任务的线程，让他们检查 shutdown 标志并退出
    pthread_cond_broadcast(&pool->task_queue.cond);

    // 等待所有存活的线程退出
    // live_threads 在这里是创建过的线程数，因为我们去掉了回收功能
    for (int i = 0; i < pool->live_threads; i++) {
        pthread_join(pool->threads[i], NULL);
    }

    // 清理任务队列中未执行的任务
    Task *task = pool->task_queue.head;
    while (task != NULL) {
        Task *next = task->next;
        // 注意：这里我们只释放任务结构体
        // 参数 argument 的内存应该由调用者在确认线程池销毁后自己管理
        // 或者设计一个统一的清理回调
        free(task);
        task = next;
    }

    // 释放所有资源
    free(pool->threads);
    pthread_mutex_destroy(&pool->task_queue.mutex);
    pthread_cond_destroy(&pool->task_queue.cond);
    pthread_mutex_destroy(&pool->mutex);
    free(pool);
}