#include <ATen/Functions.h>
#include <ATen/Tensor.h>
#include <ATen/Utils.h>
#include <ATen/xpu/XPUGeneratorImpl.h>
#include <ATen/xpu/XPUGraph.h>
#include <ATen/xpu/XPUGraphsUtils.h>
#include <c10/core/StreamGuard.h>
#include <c10/util/CallOnce.h>
#include <c10/xpu/XPUFunctions.h>

constexpr uint64_t PHILOX_ROUND_SIZE = 4;

namespace at {
namespace xpu::detail {
namespace {

/*
 * Currently, there is one generator pool containing XPU generator per device.
 * Each generator is lazily initialized the first time generator is
 * requested for a device.
 */
DeviceIndex num_gpus = -1;
std::deque<c10::once_flag> xpu_gens_init_flag;
std::vector<Generator> default_gens_xpu;

void initXPUGenVector() {
  static bool init_flag [[maybe_unused]] = []() {
    num_gpus = device_count();
    xpu_gens_init_flag.resize(num_gpus);
    default_gens_xpu.resize(num_gpus);
    return true;
  }();
}

} // anonymous namespace

// Get the default generator with a random seed for a specific xpu device.
const Generator& getDefaultXPUGenerator(DeviceIndex device) {
  initXPUGenVector();
  if (device == -1) {
    device = c10::xpu::current_device();
  }
  check_device_index(device);
  c10::call_once(xpu_gens_init_flag[device], [&]() {
    default_gens_xpu[device] = make_generator<XPUGeneratorImpl>(device);
    default_gens_xpu[device].seed();
  });
  return default_gens_xpu[device];
}

// Create a generator with a fixed seed for a specific xpu device.
Generator createXPUGenerator(DeviceIndex device) {
  initXPUGenVector();
  if (device == -1) {
    device = c10::xpu::current_device();
  }
  check_device_index(device);
  auto gen = make_generator<XPUGeneratorImpl>(device);
  auto xpu_gen = check_generator<XPUGeneratorImpl>(gen);
  xpu_gen->set_current_seed(default_rng_seed_val);
  xpu_gen->set_philox_offset_per_thread(0);
  return gen;
}

} // namespace xpu::detail

// Creates a clone of this XPU Generator State.
c10::intrusive_ptr<XPUGeneratorState> XPUGeneratorState::clone() {
  return make_intrusive<XPUGeneratorState>(
      seed_, philox_offset_per_thread_, offset_intragraph_);
}

// Function to increase the internal offset based on the specified increment.
void XPUGeneratorState::increase(uint64_t increment) {
  increment = ((increment + PHILOX_ROUND_SIZE - 1) / PHILOX_ROUND_SIZE) *
      PHILOX_ROUND_SIZE;
  if (at::xpu::currentStreamCaptureStatus() !=
      at::xpu::CaptureStatus::Executing) {
    TORCH_INTERNAL_ASSERT(
        capturing_,
        "Attempt to increase offset for a XPU generator not in capture mode.");
    TORCH_INTERNAL_ASSERT(
        offset_intragraph_ % 4 == 0, "RNG offset must be a multiple of 4.");
    TORCH_INTERNAL_ASSERT(
        offset_intragraph_ <= std::numeric_limits<uint32_t>::max() - increment,
        "Increment causes overflow in the offset value.");
    offset_intragraph_ += increment;
  } else {
    TORCH_INTERNAL_ASSERT(
        !capturing_,
        "Offset increment outside graph capture encountered unexpectedly.");
    TORCH_INTERNAL_ASSERT(
        philox_offset_per_thread_ % 4 == 0,
        "RNG offset must be a multiple of 4.");
    philox_offset_per_thread_ += increment;
  }
}

// State can be used by multiple graph
void XPUGeneratorState::register_graph(xpu::XPUGraphImpl* graph) {
  // Ensures that the RNG state is not currently being captured.
  at::xpu::assertNotCapturing(
      "Cannot register the state during capturing stage.");

  if (registered_graphs_.empty()) {
    auto options = at::TensorOptions().device(at::kXPU).dtype(at::kLong);
    seed_extragraph_ = at::empty({1}, options);
    offset_extragraph_ = at::empty({1}, options);
  }

  if (registered_graphs_.find(graph) == registered_graphs_.end()) {
    registered_graphs_.insert(graph);
  }
}

void XPUGeneratorState::unregister_graph(xpu::XPUGraphImpl* graph) {
  TORCH_CHECK(
      registered_graphs_.find(graph) != registered_graphs_.end(),
      "The graph should be registered to the state");
  registered_graphs_.erase(graph);

  if (registered_graphs_.empty()) {
    seed_extragraph_.reset();
    offset_extragraph_.reset();
  }
}

void XPUGeneratorState::capture_prologue() {
  capturing_ = true;
  offset_intragraph_ = 0;
  seed_extragraph_.fill_(int64_t(seed_));
  offset_extragraph_.fill_(int64_t(0));
}

uint64_t XPUGeneratorState::capture_epilogue() {
  capturing_ = false;
  return offset_intragraph_;
}

void XPUGeneratorState::replay_prologue(uint64_t wholegraph_increment) {
  // Ensures the generator is not in capturing mode.
  at::xpu::assertNotCapturing(
      "Cannot prepare for replay during capturing stage.");
  if (wholegraph_increment) {
    seed_extragraph_.fill_(int64_t(seed_));
    offset_extragraph_.fill_(int64_t(philox_offset_per_thread_));
    increase(wholegraph_increment);
  }
}

XPUGeneratorImpl::XPUGeneratorImpl(DeviceIndex device_index)
    : GeneratorImpl{
          Device(DeviceType::XPU, device_index),
          DispatchKeySet(c10::DispatchKey::XPU)} {
  at::xpu::assertNotCapturing("Cannot construct a new XPUGeneratorImpl");
  state_ = make_intrusive<XPUGeneratorState>();
}

XPUGeneratorImpl::XPUGeneratorImpl(
    DeviceIndex device_index,
    intrusive_ptr<XPUGeneratorState> state)
    : GeneratorImpl{Device(DeviceType::XPU, device_index), DispatchKeySet(c10::DispatchKey::XPU)},
      state_(std::move(state)) {}

void XPUGeneratorImpl::set_current_seed(uint64_t seed) {
  if (C10_LIKELY(
          at::xpu::currentStreamCaptureStatus() ==
          at::xpu::CaptureStatus::Executing)) {
    state_->seed_ = seed;
    state_->philox_offset_per_thread_ = 0;
  } else {
    TORCH_CHECK(
        state_->seed_ == seed,
        "XPUGeneratorImpl::set_current_seed can be called during stream capture only if new seed is the same as the original seed.");
  }
}

void XPUGeneratorImpl::set_offset(uint64_t offset) {
  at::xpu::assertNotCapturing("Cannot call XPUGeneratorImpl::set_offset");
  set_philox_offset_per_thread(offset);
}

uint64_t XPUGeneratorImpl::get_offset() const {
  at::xpu::assertNotCapturing("Cannot call XPUGeneratorImpl::get_offset");
  return state_->philox_offset_per_thread_;
}

uint64_t XPUGeneratorImpl::current_seed() const {
  return state_->seed_;
}

uint64_t XPUGeneratorImpl::seed() {
  at::xpu::assertNotCapturing("Cannot call XPUGeneratorImpl::seed");
  auto random = c10::detail::getNonDeterministicRandom(true);
  this->set_current_seed(random);
  return random;
}

c10::intrusive_ptr<c10::TensorImpl> XPUGeneratorImpl::get_state() const {
  // The RNG state comprises the seed, and an offset used for Philox.
  constexpr size_t seed_size = sizeof(uint64_t);
  constexpr size_t offset_size = sizeof(uint64_t);
  constexpr size_t total_size = seed_size + offset_size;

  // The internal state is returned as a CPU byte tensor.
  auto state_tensor = at::detail::empty_cpu(
      {static_cast<int64_t>(total_size)},
      ScalarType::Byte,
      std::nullopt,
      std::nullopt,
      std::nullopt,
      std::nullopt);
  auto rng_state = state_tensor.data_ptr<uint8_t>();
  auto current_seed = this->current_seed();
  auto offset = this->philox_offset_per_thread();
  memcpy(rng_state, &current_seed, seed_size);
  memcpy(rng_state + seed_size, &offset, offset_size);

  return state_tensor.getIntrusivePtr();
}

void XPUGeneratorImpl::set_state(const c10::TensorImpl& new_state) {
  constexpr size_t seed_size = sizeof(uint64_t);
  constexpr size_t offset_size = sizeof(uint64_t);
  constexpr size_t total_size = seed_size + offset_size;

  at::detail::check_rng_state(new_state);

  bool no_philox_seed = false;
  auto new_state_size = new_state.numel();
  if (new_state_size == total_size - offset_size) {
    no_philox_seed = true;
  } else {
    TORCH_CHECK(new_state_size == total_size, "RNG state is wrong size");
  }

  uint64_t input_seed = 0;
  auto new_rng_state = new_state.data_dtype_initialized<uint8_t>();
  memcpy(&input_seed, new_rng_state, seed_size);
  this->set_current_seed(input_seed);
  uint64_t philox_offset = 0;
  if (!no_philox_seed) {
    memcpy(&philox_offset, new_rng_state + seed_size, offset_size);
  }
  this->set_philox_offset_per_thread(philox_offset);
}

void XPUGeneratorImpl::graphsafe_set_state(
    const c10::intrusive_ptr<GeneratorImpl>& gen) {
  c10::intrusive_ptr<XPUGeneratorImpl> xpu_gen =
      dynamic_intrusive_pointer_cast<XPUGeneratorImpl>(gen);
  TORCH_CHECK(xpu_gen, "Expected a XPU Generator");
  state_ = xpu_gen->state_;
}

c10::intrusive_ptr<c10::GeneratorImpl> XPUGeneratorImpl::graphsafe_get_state()
    const {
  auto gen = make_intrusive<XPUGeneratorImpl>(device().index(), state_);
  return gen;
}

void XPUGeneratorImpl::set_philox_offset_per_thread(uint64_t offset) {
  TORCH_CHECK(offset % 4 == 0, "offset must be a multiple of 4");
  if (C10_LIKELY(
          at::xpu::currentStreamCaptureStatus() ==
          at::xpu::CaptureStatus::Executing)) {
    state_->philox_offset_per_thread_ = offset;
  } else {
    state_->offset_intragraph_ = offset;
  }
}

uint64_t XPUGeneratorImpl::philox_offset_per_thread() const {
  if (C10_LIKELY(
          at::xpu::currentStreamCaptureStatus() ==
          at::xpu::CaptureStatus::Executing)) {
    return state_->philox_offset_per_thread_;
  } else {
    return state_->offset_intragraph_;
  }
}

void XPUGeneratorImpl::register_graph(xpu::XPUGraphImpl* graph) {
  graph->register_generator_state(state_);
  state_->register_graph(graph);
}

void XPUGeneratorImpl::unregister_graph(xpu::XPUGraphImpl* graph) {
  state_->unregister_graph(graph);
}

// 1, During graph capture, constructs a PhiloxXpuState
//    [extragraph seed ptr, extragraph offset ptr, intragraph offset] on host
// 2, Before each replay, the extragraph seed and offset tensors will be updated
//    extragraph offset = philox_offset_per_thread_ + intragraph offset
// 3, During replay, kernel will compute final offset = *extragraph offset ptr +
// intragraph offset
PhiloxXpuState XPUGeneratorImpl::philox_xpu_state(uint64_t increment) {
  if (at::xpu::currentStreamCaptureStatus() !=
      at::xpu::CaptureStatus::Executing) {
    uint32_t offset = state_->offset_intragraph_;
    state_->increase(increment);
    return PhiloxXpuState(
        state_->seed_extragraph_.data_ptr<int64_t>(),
        state_->offset_extragraph_.data_ptr<int64_t>(),
        offset);
  } else {
    uint64_t offset = state_->philox_offset_per_thread_;
    state_->increase(increment);
    return PhiloxXpuState(state_->seed_, offset);
  }
}

std::pair<uint64_t, uint64_t> XPUGeneratorImpl::philox_engine_inputs(
    uint64_t increment) {
  at::xpu::assertNotCapturing(
      "Refactor this op to use XPUGeneratorImpl::philox_xpu_state. Cannot call XPUGeneratorImpl::philox_engine_inputs");
  uint64_t offset = state_->philox_offset_per_thread_;
  state_->increase(increment);
  return std::make_pair(state_->seed_, offset);
}

DeviceType XPUGeneratorImpl::device_type() {
  return DeviceType::XPU;
}

std::shared_ptr<XPUGeneratorImpl> XPUGeneratorImpl::clone() const {
  return std::shared_ptr<XPUGeneratorImpl>(this->clone_impl());
}

XPUGeneratorImpl* XPUGeneratorImpl::clone_impl() const {
  at::xpu::assertNotCapturing("Cannot call XPUGeneratorImpl::clone_impl");
  auto gen = new XPUGeneratorImpl(this->device().index(), state_->clone());
  return gen;
}

} // namespace at
