#include "core/parallel_for.h"

#include <atomic>
#include <condition_variable>
#include <exception>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

namespace fn {

struct ThreadPool::Impl {
    std::vector<std::thread> workers;
    std::queue<std::function<void()>> jobs;
    std::mutex m;
    std::condition_variable cv;
    bool stop = false;

    int n_threads = 0;

    // Estado por parallel_for:
    std::atomic<int> active_jobs{0};
    std::mutex done_m;
    std::condition_variable done_cv;
    std::exception_ptr first_exc;
    std::mutex exc_m;

    void worker_loop() {
        for (;;) {
            std::function<void()> job;
            {
                std::unique_lock<std::mutex> lk(m);
                cv.wait(lk, [this] { return stop || !jobs.empty(); });
                if (stop && jobs.empty()) return;
                job = std::move(jobs.front());
                jobs.pop();
            }
            try {
                job();
            } catch (...) {
                std::lock_guard<std::mutex> lk(exc_m);
                if (!first_exc) first_exc = std::current_exception();
            }
            if (active_jobs.fetch_sub(1) == 1) {
                std::lock_guard<std::mutex> lk(done_m);
                done_cv.notify_all();
            }
        }
    }
};

ThreadPool::ThreadPool(int n_threads) : impl_(new Impl) {
    if (n_threads <= 0) {
        n_threads = static_cast<int>(std::thread::hardware_concurrency());
        if (n_threads <= 0) n_threads = 1;
    }
    impl_->n_threads = n_threads;
    impl_->workers.reserve(static_cast<std::size_t>(n_threads));
    for (int i = 0; i < n_threads; ++i) {
        impl_->workers.emplace_back([this] { impl_->worker_loop(); });
    }
}

ThreadPool::~ThreadPool() {
    if (!impl_) return;
    {
        std::lock_guard<std::mutex> lk(impl_->m);
        impl_->stop = true;
    }
    impl_->cv.notify_all();
    for (auto& t : impl_->workers) {
        if (t.joinable()) t.join();
    }
    delete impl_;
    impl_ = nullptr;
}

int ThreadPool::n_threads() const { return impl_->n_threads; }

void ThreadPool::parallel_for_chunks(std::size_t begin, std::size_t end,
                                     const std::function<void(int, std::size_t, std::size_t)>& fn) {
    if (end <= begin) return;
    int N = impl_->n_threads;
    std::size_t total = end - begin;
    std::size_t chunk = (total + static_cast<std::size_t>(N) - 1) /
                        static_cast<std::size_t>(N);
    int actual_chunks = 0;
    for (int t = 0; t < N; ++t) {
        std::size_t lo = begin + chunk * static_cast<std::size_t>(t);
        if (lo >= end) break;
        std::size_t hi = lo + chunk;
        if (hi > end) hi = end;
        ++actual_chunks;
    }
    if (actual_chunks == 0) return;

    impl_->first_exc = nullptr;
    impl_->active_jobs.store(actual_chunks);

    {
        std::lock_guard<std::mutex> lk(impl_->m);
        for (int t = 0; t < N; ++t) {
            std::size_t lo = begin + chunk * static_cast<std::size_t>(t);
            if (lo >= end) break;
            std::size_t hi = lo + chunk;
            if (hi > end) hi = end;
            int tid = t;
            impl_->jobs.emplace([&fn, tid, lo, hi] { fn(tid, lo, hi); });
        }
    }
    impl_->cv.notify_all();

    {
        std::unique_lock<std::mutex> lk(impl_->done_m);
        impl_->done_cv.wait(lk, [this] {
            return impl_->active_jobs.load() == 0;
        });
    }

    if (impl_->first_exc) {
        std::exception_ptr e = impl_->first_exc;
        impl_->first_exc = nullptr;
        std::rethrow_exception(e);
    }
}

void ThreadPool::parallel_for(std::size_t begin, std::size_t end,
                              const std::function<void(std::size_t)>& fn) {
    parallel_for_chunks(begin, end,
        [&fn](int /*tid*/, std::size_t lo, std::size_t hi) {
            for (std::size_t i = lo; i < hi; ++i) fn(i);
        });
}

} // namespace fn
