// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "execution_context.h"
#include "request.h"

namespace Generators {

struct DecoderIO;

struct BatchedSamplingPlan {
  void Reserve(size_t capacity) {
    requests.reserve(capacity);
    logits.reserve(capacity);
    params.reserve(capacity);
    states.reserve(capacity);
  }

  void Clear() {
    requests.clear();
    logits.clear();
    params.clear();
    states.clear();
  }

  std::vector<Request*> requests;
  std::vector<DeviceSpan<float>> logits;
  std::vector<BatchedSamplingParams> params;
  std::vector<BatchedSamplerState*> states;
};

struct ScheduledRequests {
  ScheduledRequests(std::vector<std::shared_ptr<Request>> requests,
                    std::shared_ptr<Model> model,
                    BatchedSampler* batched_sampler,
                    BatchedSamplingPlan* sampling_plan);

  ExecutionContext& CreateExecutionContext();

  std::shared_ptr<GeneratorParams> Params();

  auto begin() const {
    return requests_.begin();
  }

  auto end() const {
    return requests_.end();
  }

  size_t size() const {
    return requests_.size();
  }

  explicit operator bool() const { return !requests_.empty(); };

  auto operator[](size_t idx) const {
    assert(idx < requests_.size());
    return requests_[idx];
  }

  const std::vector<std::shared_ptr<Request>>& Requests() const { return requests_; }

  void AddDecoderState(std::unique_ptr<DecoderIO> decoder_state);

  std::vector<DeviceSpan<float>> ProcessLogits();

  void GenerateNextTokens();
  void BeginTransaction();
  void GenerateNextTokensForTransaction(
      const StepPlan& plan,
      std::vector<RequestStepResult>& results);
  void RestoreStateForTransaction();
  void CommitStateForTransaction();

 private:
  bool PrepareBatchedSamplingPlan(bool require_transaction_support);
  bool TryGenerateNextTokensBatched(std::vector<DeviceSpan<float>>& logits);

  std::vector<std::shared_ptr<Request>> requests_;
  std::shared_ptr<Model> model_;
  std::unique_ptr<DecoderIO> decoder_state_;
  std::unique_ptr<ExecutionContext> execution_context_;
  std::shared_ptr<GeneratorParams> params_;
  BatchedSampler* batched_sampler_{};
  BatchedSamplingPlan* sampling_plan_{};
  size_t transaction_checkpoint_count_{};
  bool transaction_uses_batched_sampler_{};
  bool sampler_checkpoint_active_{};
};

}  // namespace Generators
