// Copyright 2026 The Action Engine Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "thread/select.h"

#include <algorithm>
#include <atomic>
#include <cstddef>
#include <cstdint>

#include <absl/log/check.h>

#include "thread/boost_primitives.h"
#include "thread/fiber_diagnostics.h"
#include "thread/util.h"

namespace thread {
internal::CaseStateArray MakeShuffledCaseStateArray(size_t num_cases) {
  internal::CaseStateArray case_states(num_cases);

  // Use inside-out Fisher-Yates shuffle to combine initialization and
  // permutation.
  if (num_cases > 0) {
    case_states[0].index = 0;
  }
  for (size_t i = 1; i < num_cases; ++i) {
    const size_t swap = Rand32() % (i + 1);
    case_states[i].index = case_states[swap].index;
    case_states[swap].index = static_cast<int>(i);
  }

  for (size_t i = 0; i < num_cases; ++i) {
    case_states[i].next =
        &case_states[i];  // Points to itself, not on any list.
  }

  return case_states;
}

int SelectUntil(absl::Time deadline, const CaseArray& cases) {
  CHECK(!cases.empty()) << "No cases provided";
  internal::Selector selector;

  internal::CaseStateArray case_states =
      MakeShuffledCaseStateArray(cases.size());

  const bool need_to_block = deadline != absl::InfinitePast();
  bool ready = false;
  int registered_case_states = 0;
  for (size_t case_state_idx = 0; case_state_idx < cases.size();
       ++case_state_idx) {
    internal::CaseInSelectClause* case_state = &case_states[case_state_idx];

    case_state->case_ptr = &cases[case_state->index];
    case_state->prev = nullptr;  // Not on any list so far.
    case_state->selector = &selector;

    ++registered_case_states;

    if (case_state->Handle(/*enqueue=*/need_to_block)) {
      ready = true;
      break;
    }
  }

  // Either the deadline is in the past, or at least one case was ready
  // (therefore there is one picked)
  if (!need_to_block) {
    // No need to Unregister() any cases since we passed enqueue=false
    // to each Handle() above.
    thread::MutexLock lock(&selector.mu);
    return ready ? selector.picked_case_index : -1;
  }

  bool expired = false;
  int selected_case = internal::Selector::kNonePicked;
  {
    thread::MutexLock lock(&selector.mu);
    if (!ready) {
      // Annotated here rather than in CondVar, so a report names the
      // selectables instead of the selector's internal condition variable.
      THREAD_WAIT_SCOPE(wait, WaitKind::kSelect, &selector,
                        DeadlineNanos(deadline));
      const void* selectables[kMaxRecordedSelectables];
      const size_t recorded = std::min(cases.size(), kMaxRecordedSelectables);
      for (size_t index = 0; index < recorded; ++index) {
        selectables[index] = cases[index].selectable;
      }
      wait.RecordSelectables(selectables, recorded);

      const bool expirable = deadline != absl::InfiniteFuture();
      expired = !selector.WaitForPickUntil(deadline);
      DCHECK(expirable || !expired);
      if (expired) {
        // Deadline expiry. Ensure nothing is picked from this point.
        selector.picked_case_index = static_cast<int>(cases.size());
      }
    }
    selected_case = selector.picked_case_index;
  }

  // Unregister from all events with which we are registered. We know that
  // all registered case states were Handle()d with enqueue=true.
  for (int i = 0; i < registered_case_states; ++i) {
    if (internal::CaseInSelectClause* case_state = &case_states[i];
        case_state->index != selected_case) {
      // sel.picked was unregistered by the notifier.
      case_state->GetCase()->selectable->Unregister(case_state);
    }
  }

  return expired ? -1 : selected_case;
}
}  // namespace thread
