/*
 * Copyright 2026 Rive
 */

#include "rive/async/work_pool.hpp"
#ifdef WITH_RIVE_SCRIPTNET
#include "rive/scriptnet/net.hpp"
#endif
#ifdef WITH_RIVE_SCRIPTING_WASM
#include "rive/wasm/wasm_jobs.hpp"
#include "rive/wasm/wasm_scripting_vm.hpp"
#endif
#include <algorithm>
#include <mutex>

namespace rive
{

std::atomic<uint64_t> WorkPool::s_nextOwnerId{1};

uint64_t WorkPool::nextOwnerId()
{
  return s_nextOwnerId.fetch_add(1, std::memory_order_relaxed);
}

// ============================================================================
// No-threading implementation
// ============================================================================

#ifndef WITH_RIVE_THREADING

WorkPool::WorkPool() {}

WorkPool::~WorkPool()
{
  // Deliver onCancel for tasks whose cancel flag was already set
  // (e.g. by cancelAllForOwner) but never polled. Tasks that were
  // never cancelled are silently dropped — calling virtual callbacks
  // during destruction is unsafe if dependent state (e.g. Lua VM)
  // has already been torn down.
  for (auto& task : m_workQueue)
  {
    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
    }
  }
  m_workQueue.clear();
}

uint64_t WorkPool::submit(rcp<WorkTask> task)
{
  if (!task)
    return 0;
  task->setStatus(WorkStatus::Pending);
  task->setSubmitGeneration(m_nextHandle);
  m_workQueue.push_back(std::move(task));
  return m_nextHandle++;
}

uint32_t WorkPool::pollCompletedWork(uint32_t maxCallbacks)
{
  uint32_t processed = 0;
  // Only what was submitted before this poll: a task submitted while
  // another completes runs on a later poll, as it would behind worker
  // threads, so a chain of decodes spreads over polls rather than landing in
  // one. Stamped by submit order, so a completion that polls again neither
  // drains this loop's queue from under it nor lets it run later work.
  const uint64_t submittedBefore = m_nextHandle;

  while (processed < maxCallbacks && !m_workQueue.empty() &&
         m_workQueue.front()->submitGeneration() < submittedBefore)
  {
    auto task = std::move(m_workQueue.front());
    m_workQueue.pop_front();

    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
      processed++;
      continue;
    }

    task->setStatus(WorkStatus::Running);
    bool success = task->execute();

    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
    }
    else if (success)
    {
      task->setStatus(WorkStatus::Completed);
      task->onComplete();
    }
    else
    {
      task->setStatus(WorkStatus::Failed);
      task->onError(task->errorMessage());
    }
    processed++;
  }
  return processed;
}

bool WorkPool::hasPendingWork() const { return !m_workQueue.empty(); }

bool WorkPool::hasPendingWorkForOwner(uint64_t ownerId) const
{
  // Tasks leave the queue as they are delivered, so it holds exactly the
  // undelivered ones.
  return std::any_of(
      m_workQueue.begin(),
      m_workQueue.end(),
      [ownerId](const rcp<WorkTask>& task) {
        return task->ownerId() == ownerId;
      });
}

void WorkPool::cancelAllForOwner(uint64_t ownerId)
{
  // Only set the cancel flag here. onCancel() is delivered once from
  // pollCompletedWork() when the task is dequeued, avoiding double-cancel.
  for (auto& task : m_workQueue)
  {
    if (task->ownerId() == ownerId)
      task->cancel();
  }
}

// ============================================================================
// Threaded implementation
// ============================================================================

#else // WITH_RIVE_THREADING

WorkPool::WorkPool()
{
  unsigned int n = std::max(std::thread::hardware_concurrency(), 1u);
  // Cap at a reasonable number for image decode.
  n = std::min(n, 4u);
  for (unsigned int i = 0; i < n; i++)
  {
    m_threads.emplace_back(&WorkPool::workerLoop, this);
  }
}

WorkPool::~WorkPool()
{
  {
    std::lock_guard<std::mutex> lock(m_queueMutex);
    m_shutdown = true;
  }
  m_haveWork.notify_all();
  for (auto& t : m_threads)
  {
    if (t.joinable())
      t.join();
  }
  // Deliver onCancel for tasks whose cancel flag was already set but
  // never polled. Tasks that were never cancelled are silently dropped —
  // calling virtual callbacks during destruction is unsafe if dependent
  // state (e.g. Lua VM) has already been torn down.
  for (auto& task : m_completedQueue)
  {
    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
    }
  }
  m_completedQueue.clear();
  for (auto& task : m_workQueue)
  {
    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
    }
  }
  m_workQueue.clear();
}

uint64_t WorkPool::submit(rcp<WorkTask> task)
{
  if (!task)
    return 0;
  uint64_t handle;
  {
    std::lock_guard<std::mutex> lock(m_queueMutex);
    task->setStatus(WorkStatus::Pending);
    task->setSubmitGeneration(m_cancelGeneration);
    handle = m_nextHandle++;
    m_undeliveredByOwner[task->ownerId()]++;
    m_workQueue.push_back(std::move(task));
  }
  m_haveWork.notify_one();
  return handle;
}

void WorkPool::workerLoop()
{
  while (true)
  {
    rcp<WorkTask> task;
    {
      std::unique_lock<std::mutex> lock(m_queueMutex);
      m_haveWork.wait(lock, [this] {
        return m_shutdown || !m_workQueue.empty();
      });
      if (m_shutdown && m_workQueue.empty())
        return;
      task = std::move(m_workQueue.front());
      m_workQueue.pop_front();
    }

    m_inFlightCount.fetch_add(1, std::memory_order_relaxed);

    if (task->isCancelled())
    {
      task->setStatus(WorkStatus::Cancelled);
      std::lock_guard<std::mutex> lock(m_completedMutex);
      m_completedQueue.push_back(std::move(task));
      m_inFlightCount.fetch_sub(1, std::memory_order_relaxed);
      continue;
    }

    task->setStatus(WorkStatus::Running);
    bool success = task->execute();

    if (task->isCancelled())
      task->setStatus(WorkStatus::Cancelled);
    else if (success)
      task->setStatus(WorkStatus::Completed);
    else
      task->setStatus(WorkStatus::Failed);

    {
      std::lock_guard<std::mutex> lock(m_completedMutex);
      m_completedQueue.push_back(std::move(task));
    }
    m_inFlightCount.fetch_sub(1, std::memory_order_relaxed);
  }
}

uint32_t WorkPool::pollCompletedWork(uint32_t maxCallbacks)
{
  uint32_t processed = 0;
  // Only what had completed before this poll, as the single threaded pool
  // delivers only what was queued before it.
  // Workers can finish several decodes a frame, each held twice in the wasm
  // heap until the boundary: budget this by decoded bytes, not task count.
  size_t completed;
  {
    std::lock_guard<std::mutex> lock(m_completedMutex);
    completed = m_completedQueue.size();
  }

  while (processed < maxCallbacks && processed < completed)
  {
    rcp<WorkTask> task;
    {
      std::lock_guard<std::mutex> lock(m_completedMutex);
      if (m_completedQueue.empty())
        break;
      task = std::move(m_completedQueue.front());
      m_completedQueue.pop_front();
    }

    bool ownerCancelled = false;
    {
      std::lock_guard<std::mutex> lock(m_queueMutex);
      auto undelivered = m_undeliveredByOwner.find(task->ownerId());
      if (undelivered != m_undeliveredByOwner.end() &&
          --undelivered->second == 0)
      {
        m_undeliveredByOwner.erase(undelivered);
      }
      auto it = m_cancelledOwners.find(task->ownerId());
      if (it != m_cancelledOwners.end())
      {
        // Task is owner-cancelled only if it was submitted
        // before the cancel call (submit gen < cancel gen).
        ownerCancelled = task->submitGeneration() < it->second;
      }
    }

    if (!task->isCancelled() && !ownerCancelled)
    {
      if (task->status() == WorkStatus::Completed)
        task->onComplete();
      else if (task->status() == WorkStatus::Failed)
        task->onError(task->errorMessage());
    }
    else
    {
      task->setStatus(WorkStatus::Cancelled);
      task->onCancel();
    }
    processed++;
  }
  return processed;
}

bool WorkPool::hasPendingWork() const
{
  if (m_inFlightCount.load(std::memory_order_relaxed) > 0)
    return true;
  {
    std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_queueMutex));
    if (!m_workQueue.empty())
      return true;
  }
  {
    std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_completedMutex));
    if (!m_completedQueue.empty())
      return true;
  }
  return false;
}

bool WorkPool::hasPendingWorkForOwner(uint64_t ownerId) const
{
  std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_queueMutex));
  return m_undeliveredByOwner.count(ownerId) != 0;
}

void WorkPool::cancelAllForOwner(uint64_t ownerId)
{
  // Only set the cancel flag under the lock. Do NOT call the virtual
  // onCancel() here — that risks deadlocks (virtual code under mutex)
  // and double-cancel (pollCompletedWork also delivers onCancel).
  // The single onCancel() delivery happens from pollCompletedWork()
  // on the main thread when the task is dequeued.
  {
    std::lock_guard<std::mutex> lock(m_queueMutex);
    // Bump the cancel generation so that tasks submitted after this
    // call are not treated as cancelled. In-flight worker tasks
    // submitted before this point have submitGeneration < cancelGen
    // and will be caught by pollCompletedWork.
    ++m_cancelGeneration;
    m_cancelledOwners[ownerId] = m_cancelGeneration;
    for (auto& task : m_workQueue)
    {
      if (task->ownerId() == ownerId)
        task->cancel();
    }
  }
  {
    std::lock_guard<std::mutex> lock(m_completedMutex);
    for (auto& task : m_completedQueue)
    {
      if (task->ownerId() == ownerId)
        task->cancel();
    }
  }
}

#endif // WITH_RIVE_THREADING

// ============================================================================
// Global singleton + polling (shared by both threading modes)
// ============================================================================

static rcp<WorkPool>& globalWorkPoolStorage()
{
  static rcp<WorkPool> s_workPool;
  return s_workPool;
}

rcp<WorkPool>& getGlobalWorkPool()
{
  // Thread-safe lazy initialization via std::call_once.
  static std::once_flag s_initFlag;
  auto& pool = globalWorkPoolStorage();
  std::call_once(s_initFlag, [&pool] { pool = make_rcp<WorkPool>(); });
  return pool;
}

rcp<WorkPool>& getGlobalWorkPoolIfExists() { return globalWorkPoolStorage(); }

uint32_t rive_pollAsyncWork(uint32_t maxCallbacks)
{
  uint32_t processed = 0;
#ifdef WITH_RIVE_SCRIPTNET
  // Fetches poll once, first, with half the budget, so decodes cannot
  // starve them and outcomes queued during delivery wait a frame.
  processed = scriptnet::poll((maxCallbacks + 1) / 2);
#endif
  auto& pool = getGlobalWorkPoolIfExists();
  if (pool && pool->hasPendingWork())
  {
    processed += pool->pollCompletedWork(maxCallbacks - processed);
  }
#ifdef WITH_RIVE_SCRIPTING_WASM
  // A script advancing an artboard is mid call, so lanes wait for the frame.
  if (!WasmScriptingVM::inModuleCallOnThread())
  {
    WasmJobs::runQueuedInline();
    WasmScriptingVM::deliverAllLaneOutput();
  }
  WasmScriptingVM::deliverAllHeldOutcomes();
#endif
  return processed;
}

bool rive_hasPendingAsyncWork()
{
  auto& pool = getGlobalWorkPoolIfExists();
  if (pool && pool->hasPendingWork())
  {
    return true;
  }
#ifdef WITH_RIVE_SCRIPTNET
  return scriptnet::hasPendingWork();
#else
  return false;
#endif
}

} // namespace rive