// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include <jxl/jxl_threads_export.h>
#include <jxl/memory_manager.h>
#include <jxl/parallel_runner.h>
#include <jxl/resizable_parallel_runner.h>

#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <mutex>
#include <thread>
#include <vector>

namespace jpegxl {
namespace {

// A thread pool that allows changing the number of threads it runs. It also
// runs tasks on the calling thread, which can work better on schedulers for
// heterogeneous architectures.
struct ResizeableParallelRunner {
  void SetNumThreads(size_t num) {
    if (num > 0) {
      num -= 1;
    }
    {
      std::unique_lock<std::mutex> l(state_mutex_);
      num_desired_workers_ = num;
      workers_can_proceed_.notify_all();
    }
    if (workers_.size() < num) {
      for (size_t i = workers_.size(); i < num; i++) {
        workers_.emplace_back([this, i]() { WorkerBody(i); });
      }
    }
    if (workers_.size() > num) {
      for (size_t i = num; i < workers_.size(); i++) {
        workers_[i].join();
      }
      workers_.resize(num);
    }
  }

  ~ResizeableParallelRunner() { SetNumThreads(0); }

  JxlParallelRetCode Run(void* jxl_opaque, JxlParallelRunInit init,
                         JxlParallelRunFunction func, uint32_t start,
                         uint32_t end) {
    if (start + 1 == end) {
      JxlParallelRetCode ret = init(jxl_opaque, 1);
      if (ret != 0) return ret;

      func(jxl_opaque, start, 0);
      return ret;
    }

    size_t num_workers = std::min<size_t>(workers_.size() + 1, end - start);
    JxlParallelRetCode ret = init(jxl_opaque, num_workers);
    if (ret != 0) {
      return ret;
    }

    {
      std::unique_lock<std::mutex> l(state_mutex_);
      // Avoid waking up more workers than needed.
      max_running_workers_ = end - start - 1;
      next_task_ = start;
      end_task_ = end;
      func_ = func;
      jxl_opaque_ = jxl_opaque;
      work_available_ = true;
      num_running_workers_++;
      workers_can_proceed_.notify_all();
    }

    DequeueTasks(0);

    while (true) {
      std::unique_lock<std::mutex> l(state_mutex_);
      if (num_running_workers_ == 0) break;
      work_done_.wait(l);
    }

    return ret;
  }

 private:
  void WorkerBody(size_t worker_id) {
    while (true) {
      {
        std::unique_lock<std::mutex> l(state_mutex_);
        // Worker pool was reduced, resize down.
        if (worker_id >= num_desired_workers_) {
          return;
        }
        // Nothing to do this time.
        if (!work_available_ || worker_id >= max_running_workers_) {
          workers_can_proceed_.wait(l);
          continue;
        }
        num_running_workers_++;
      }
      DequeueTasks(worker_id + 1);
    }
  }

  void DequeueTasks(size_t thread_id) {
    while (true) {
      uint32_t task = next_task_++;
      if (task >= end_task_) {
        std::unique_lock<std::mutex> l(state_mutex_);
        num_running_workers_--;
        work_available_ = false;
        if (num_running_workers_ == 0) {
          work_done_.notify_all();
        }
        break;
      }
      func_(jxl_opaque_, task, thread_id);
    }
  }

  // Checks when the worker has something to do, which can be one of:
  // - quitting (when worker_id >= num_desired_workers_)
  // - having work available for them (work_available_ is true and worker_id >=
  // max_running_workers_)
  std::condition_variable workers_can_proceed_;

  // Workers are done, and the main thread can proceed (num_running_workers_ ==
  // 0)
  std::condition_variable work_done_;

  std::vector<std::thread> workers_;

  // Protects all the remaining variables, except for func_, jxl_opaque_ and
  // end_task_ (for which only the write by the main thread is protected, and
  // subsequent uses by workers happen-after it) and next_task_ (which is
  // atomic).
  std::mutex state_mutex_;

  // Range of tasks still need to be done.
  std::atomic<uint32_t> next_task_;
  uint32_t end_task_;

  // Function to run and its argument.
  JxlParallelRunFunction func_;
  void* jxl_opaque_;  // not owned

  // Variables that control the workers:
  // - work_available_ is set to true after a call to Run() and to false at the
  // end of it.
  // - num_desired_workers_ represents the number of workers that should be
  // present.
  // - max_running_workers_ represents the number of workers that should be
  // executing tasks.
  // - num_running_workers_ represents the number of workers that are executing
  // tasks.
  size_t num_desired_workers_ = 0;
  size_t max_running_workers_ = 0;
  size_t num_running_workers_ = 0;
  bool work_available_ = false;
};
}  // namespace
}  // namespace jpegxl

extern "C" {
JXL_THREADS_EXPORT JxlParallelRetCode JxlResizableParallelRunner(
    void* runner_opaque, void* jpegxl_opaque, JxlParallelRunInit init,
    JxlParallelRunFunction func, uint32_t start_range, uint32_t end_range) {
  return static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque)
      ->Run(jpegxl_opaque, init, func, start_range, end_range);
}

JXL_THREADS_EXPORT void* JxlResizableParallelRunnerCreate(
    const JxlMemoryManager* memory_manager) {
  return new jpegxl::ResizeableParallelRunner();
}

JXL_THREADS_EXPORT void JxlResizableParallelRunnerSetThreads(
    void* runner_opaque, size_t num_threads) {
  static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque)
      ->SetNumThreads(num_threads);
}

JXL_THREADS_EXPORT void JxlResizableParallelRunnerDestroy(void* runner_opaque) {
  delete static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque);
}

JXL_THREADS_EXPORT uint32_t
JxlResizableParallelRunnerSuggestThreads(uint64_t xsize, uint64_t ysize) {
  // ~one thread per group.
  return std::min<uint64_t>(std::thread::hardware_concurrency(),
                            xsize * ysize / (256 * 256));
}
}