#include <frostbite2D/core/task_system.h>

#include <SDL2/SDL_log.h>
#include <algorithm>
#include <stdexcept>
#include <utility>

namespace frostbite2D {

TaskSystem& TaskSystem::get() {
  static TaskSystem instance;
  return instance;
}

bool TaskSystem::init(const TaskSystemConfig& config) {
  if (initialized_) {
    return true;
  }

  try {
    config_ = config;
    if (config_.workerCount == 0) {
      config_.workerCount = static_cast<uint32>(resolveWorkerCount(0));
    }

    {
      std::lock_guard<std::mutex> workerLock(workerMutex_);
      std::queue<Function<void()>> emptyWorkerTasks;
      std::swap(workerTasks_, emptyWorkerTasks);
    }

    {
      std::lock_guard<std::mutex> mainThreadLock(mainThreadMutex_);
      std::queue<Function<void()>> emptyMainThreadTasks;
      std::swap(mainThreadTasks_, emptyMainThreadTasks);
    }

    mainThreadId_ = std::this_thread::get_id();
    acceptingTasks_ = true;

    workers_.reserve(config_.workerCount);
    for (uint32 i = 0; i < config_.workerCount; ++i) {
      workers_.emplace_back([this]() {
        while (true) {
          Function<void()> task;

          {
            std::unique_lock<std::mutex> lock(workerMutex_);
            workerCondition_.wait(lock, [this]() {
              return !acceptingTasks_.load() || !workerTasks_.empty();
            });

            if (!acceptingTasks_.load() && workerTasks_.empty()) {
              return;
            }

            task = std::move(workerTasks_.front());
            workerTasks_.pop();
          }

          task();
        }
      });
    }

    initialized_ = true;
    SDL_Log("TaskSystem initialized with %u worker threads", config_.workerCount);
    return true;
  } catch (const std::exception& ex) {
    SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to initialize TaskSystem: %s", ex.what());
    shutdown();
    return false;
  }
}

void TaskSystem::shutdown() {
  if (!initialized_) {
    return;
  }

  acceptingTasks_ = false;
  workerCondition_.notify_all();

  for (std::thread& worker : workers_) {
    if (worker.joinable()) {
      worker.join();
    }
  }
  workers_.clear();

  {
    std::lock_guard<std::mutex> workerLock(workerMutex_);
    std::queue<Function<void()>> emptyWorkerTasks;
    std::swap(workerTasks_, emptyWorkerTasks);
  }

  {
    std::lock_guard<std::mutex> mainThreadLock(mainThreadMutex_);
    std::queue<Function<void()>> emptyMainThreadTasks;
    std::swap(mainThreadTasks_, emptyMainThreadTasks);
  }

  initialized_ = false;
  mainThreadId_ = std::thread::id();
  SDL_Log("TaskSystem shutdown");
}

bool TaskSystem::isMainThread() const {
  return initialized_ && std::this_thread::get_id() == mainThreadId_;
}

size_t TaskSystem::workerCount() const {
  return workers_.size();
}

uint32 TaskSystem::maxMainThreadCallbacksPerFrame() const {
  return config_.maxMainThreadCallbacksPerFrame;
}

void TaskSystem::drainMainThreadTasks(uint32 maxTasks) {
  if (!initialized_) {
    return;
  }

  uint32 taskBudget = maxTasks != 0 ? maxTasks : config_.maxMainThreadCallbacksPerFrame;
  uint32 processed = 0;

  while (true) {
    Function<void()> task;

    {
      std::lock_guard<std::mutex> lock(mainThreadMutex_);
      if (mainThreadTasks_.empty()) {
        break;
      }

      if (taskBudget != 0 && processed >= taskBudget) {
        break;
      }

      task = std::move(mainThreadTasks_.front());
      mainThreadTasks_.pop();
    }

    task();
    ++processed;
  }
}

TaskSystem::~TaskSystem() {
  shutdown();
}

void TaskSystem::enqueueWorkerTask(Function<void()> task) {
  if (!acceptingTasks_) {
    throw std::runtime_error("TaskSystem is not accepting new worker tasks");
  }

  {
    std::lock_guard<std::mutex> lock(workerMutex_);
    workerTasks_.push(std::move(task));
  }

  workerCondition_.notify_one();
}

void TaskSystem::enqueueMainThreadTask(Function<void()> task) {
  if (!acceptingTasks_) {
    return;
  }

  std::lock_guard<std::mutex> lock(mainThreadMutex_);
  mainThreadTasks_.push(std::move(task));
}

size_t TaskSystem::resolveWorkerCount(uint32 requestedCount) const {
  if (requestedCount > 0) {
    return requestedCount;
  }

#ifdef __SWITCH__
  return 1;
#else
  unsigned int hardwareCount = std::thread::hardware_concurrency();
  if (hardwareCount <= 1) {
    return 1;
  }

  return std::max(1u, std::min(hardwareCount - 1, 4u));
#endif
}

void TaskSystem::logUnhandledTaskException(std::exception_ptr error) {
  if (!error) {
    return;
  }

  try {
    std::rethrow_exception(error);
  } catch (const std::exception& ex) {
    SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "TaskSystem task failed: %s", ex.what());
  } catch (...) {
    SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "TaskSystem task failed with unknown exception");
  }
}

} // namespace frostbite2D
