// Copyright 2021 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/task/sequence_manager/wake_up_queue.h"

#include <optional>

#include "base/task/sequence_manager/associated_thread_id.h"
#include "base/task/sequence_manager/sequence_manager_impl.h"
#include "base/task/sequence_manager/task_queue_impl.h"
#include "base/threading/thread_checker.h"

namespace base::sequence_manager::internal {

WakeUpQueue::WakeUpQueue(
    scoped_refptr<const internal::AssociatedThreadId> associated_thread)
    : associated_thread_(std::move(associated_thread)) {}

WakeUpQueue::~WakeUpQueue() {
  DCHECK_CALLED_ON_VALID_THREAD(associated_thread_->thread_checker);
}

void WakeUpQueue::RemoveAllCanceledDelayedTasksFromFront(LazyNow* lazy_now) {
  // Repeatedly trim the front of the top queue until it stabilizes. This is
  // needed because a different queue can become the top one once you remove the
  // canceled tasks.
  while (!wake_up_queue_.empty()) {
    auto* top_queue = wake_up_queue_.top().queue;

    // If no tasks are removed from the top queue, then it means the top queue
    // cannot change anymore.
    if (!top_queue->RemoveAllCanceledDelayedTasksFromFront(lazy_now)) {
      break;
    }
  }
}

// TODO(kraynov): https://crbug.com/857101 Consider making an interface
// for SequenceManagerImpl which will expose SetNextDelayedDoWork and
// MaybeScheduleImmediateWork methods to make the functions below pure-virtual.

void WakeUpQueue::SetNextWakeUpForQueue(internal::TaskQueueImpl* queue,
                                        LazyNow* lazy_now,
                                        std::optional<WakeUp> wake_up) {
  DCHECK_CALLED_ON_VALID_THREAD(associated_thread_->thread_checker);
  DCHECK_EQ(queue->wake_up_queue(), this);
  DCHECK(queue->IsQueueEnabled() || !wake_up);

  std::optional<WakeUp> previous_wake_up = GetNextDelayedWakeUp();

  if (wake_up) {
    // Insert a new wake-up into the heap.
    if (queue->heap_handle().IsValid()) {
      // O(log n)
      wake_up_queue_.Replace(queue->heap_handle(), {wake_up.value(), queue});
    } else {
      // O(log n)
      wake_up_queue_.insert({wake_up.value(), queue});
    }
  } else {
    // Remove a wake-up from heap if present.
    if (queue->heap_handle().IsValid()) {
      wake_up_queue_.erase(queue->heap_handle());
    }
  }

  std::optional<WakeUp> new_wake_up = GetNextDelayedWakeUp();

  if (new_wake_up != previous_wake_up) {
    OnNextWakeUpChanged(lazy_now, GetNextDelayedWakeUp());
  }
}

void WakeUpQueue::MoveReadyDelayedTasksToWorkQueues(
    LazyNow* lazy_now,
    EnqueueOrder enqueue_order) {
  DCHECK_CALLED_ON_VALID_THREAD(associated_thread_->thread_checker);
  bool update_needed = false;
  while (!wake_up_queue_.empty() &&
         wake_up_queue_.top().wake_up.earliest_time() <= lazy_now->Now()) {
    internal::TaskQueueImpl* queue = wake_up_queue_.top().queue;
    // OnWakeUp() is expected to update the next wake-up for this queue with
    // SetNextWakeUpForQueue(), thus allowing us to make progress.
    queue->OnWakeUp(lazy_now, enqueue_order);
    update_needed = true;
  }

  if (!update_needed || wake_up_queue_.empty()) {
    return;
  }
  // If any queue was notified, possibly update following queues. This ensures
  // the wake up is up to date, which is necessary because calling OnWakeUp() on
  // a throttled queue may affect state that is shared between other related
  // throttled queues. The wake up for an affected queue might be pushed back
  // and needs to be updated. This is done lazily only once the related queue
  // becomes the next one to wake up, since that wake up can't be moved up.
  // `wake_up_queue_` is non-empty here, per the condition above.
  internal::TaskQueueImpl* queue = wake_up_queue_.top().queue;
  queue->UpdateWakeUp(lazy_now);
  while (!wake_up_queue_.empty()) {
    internal::TaskQueueImpl* old_queue =
        std::exchange(queue, wake_up_queue_.top().queue);
    if (old_queue == queue) {
      break;
    }
    queue->UpdateWakeUp(lazy_now);
  }
}

std::optional<WakeUp> WakeUpQueue::GetNextDelayedWakeUp() const {
  DCHECK_CALLED_ON_VALID_THREAD(associated_thread_->thread_checker);
  if (wake_up_queue_.empty()) {
    return std::nullopt;
  }
  WakeUp wake_up = wake_up_queue_.top().wake_up;
  return wake_up;
}

Value::Dict WakeUpQueue::AsValue(TimeTicks now) const {
  Value::Dict state;
  state.Set("name", GetName());
  // TODO(crbug.com/40228085): Make base::Value able to store an int64_t and
  // remove this cast.
  state.Set("registered_delay_count", checked_cast<int>(wake_up_queue_.size()));
  if (!wake_up_queue_.empty()) {
    TimeDelta delay = wake_up_queue_.top().wake_up.time - now;
    state.Set("next_delay_ms", delay.InMillisecondsF());
  }
  return state;
}

DefaultWakeUpQueue::DefaultWakeUpQueue(
    scoped_refptr<internal::AssociatedThreadId> associated_thread,
    internal::SequenceManagerImpl* sequence_manager)
    : WakeUpQueue(std::move(associated_thread)),
      sequence_manager_(sequence_manager) {}

DefaultWakeUpQueue::~DefaultWakeUpQueue() = default;

void DefaultWakeUpQueue::OnNextWakeUpChanged(LazyNow* lazy_now,
                                             std::optional<WakeUp> wake_up) {
  sequence_manager_->SetNextWakeUp(lazy_now, wake_up);
}

void DefaultWakeUpQueue::UnregisterQueue(internal::TaskQueueImpl* queue) {
  DCHECK_EQ(queue->wake_up_queue(), this);
  LazyNow lazy_now(sequence_manager_->main_thread_clock());
  SetNextWakeUpForQueue(queue, &lazy_now, std::nullopt);
}

const char* DefaultWakeUpQueue::GetName() const {
  return "DefaultWakeUpQueue";
}

NonWakingWakeUpQueue::NonWakingWakeUpQueue(
    scoped_refptr<internal::AssociatedThreadId> associated_thread)
    : WakeUpQueue(std::move(associated_thread)) {}

NonWakingWakeUpQueue::~NonWakingWakeUpQueue() = default;

void NonWakingWakeUpQueue::OnNextWakeUpChanged(LazyNow* lazy_now,
                                               std::optional<WakeUp> wake_up) {}

const char* NonWakingWakeUpQueue::GetName() const {
  return "NonWakingWakeUpQueue";
}

void NonWakingWakeUpQueue::UnregisterQueue(internal::TaskQueueImpl* queue) {
  DCHECK_EQ(queue->wake_up_queue(), this);
  SetNextWakeUpForQueue(queue, nullptr, std::nullopt);
}

}  // namespace base::sequence_manager::internal