// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
//
// Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved.

#include "generators.h"
#include <algorithm>
#include <cctype>
#include <cstring>
#include "models/streaming_processor.h"
#include "models/nemotron_speech.h"
#include "models/parakeet.h"
#include "sequences.h"
#include "models/env_utils.h"
#include "models/model.h"
#include "models/model_type.h"
#include "models/decoder_only.h"
#include "decoding_strategy.h"
#include "constrained_logits_processor.h"
#include "search.h"
#include "tracing.h"
#include "cpu/interface.h"
#include "cuda/interface.h"
#include "dml/interface.h"
#include "qnn/interface.h"
#include "webgpu/interface.h"
#include "openvino/interface.h"
#include "ryzenai/interface.h"
#include "engine/engine.h"

#if defined(_WIN32)
EXTERN_C IMAGE_DOS_HEADER __ImageBase;

std::string CurrentModulePath() {
  char path[MAX_PATH];
  GetModuleFileNameA((HINSTANCE)&__ImageBase, path, _countof(path));

  char absolute_path[MAX_PATH];
  char* name;
  GetFullPathNameA(path, _countof(path), absolute_path, &name);

  auto idx = std::distance(absolute_path, name);
  auto out_path = std::string(absolute_path);
  out_path.resize(idx);

  return out_path;
}

#include "dll_load_error.h"
#endif

void ThrowErrorIfSessionTerminated(bool is_session_terminated) {
  if (is_session_terminated)
    throw std::runtime_error("Session in Terminated state, exiting!");
}

namespace Generators {

static bool _ = (Ort::InitApi(), false);

static OrtLoggingLevel GetDefaultOrtLoggingLevel() {
  bool ort_verbose_logging = false;
  GetEnv("ORTGENAI_ORT_VERBOSE_LOGGING", ort_verbose_logging);
  return ort_verbose_logging ? OrtLoggingLevel::ORT_LOGGING_LEVEL_VERBOSE : OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR;
}

OrtGlobals::OrtGlobals()
    : env_{OrtEnv::Create(GetDefaultOrtLoggingLevel())} {
  const char* keys[] = {"max_mem", "arena_extend_strategy", "initial_chunk_size_bytes", "max_dead_bytes_per_chunk"};
  const size_t values[] = {static_cast<size_t>(0), static_cast<size_t>(-1), static_cast<size_t>(-1), static_cast<size_t>(-1)};
  auto arena_config = OrtArenaCfg::Create(keys, values, 4);
  Ort::Allocator& allocator_cpu{Ort::Allocator::GetWithDefaultOptions()};
  env_->CreateAndRegisterAllocator(allocator_cpu.GetInfo(), *arena_config);

  // Init the CPU device (special case because it always exists, and its allocator is special).
  GetDeviceInterface(DeviceType::CPU)->InitOrt(*Ort::api, allocator_cpu);
}

// The single genai global. Lazily (re)created by GetOrtGlobals() and reset by Shutdown(), so that
// genai can be torn down and re-initialized in-process (e.g. a host recreating its wrapper).
static std::unique_ptr<OrtGlobals> g_ort_globals;
// Guards lazy (re)creation and teardown of g_ort_globals so concurrent first-use (or first-use
// after a shutdown) cannot double-construct the globals, and a concurrent Shutdown() cannot race
// creation. Re-init after shutdown must remain possible, so this is a plain mutex rather than
// std::call_once (which is process-once).
static std::mutex g_ort_globals_mutex;
// Set by the process-exit EnsureShutdown destructor to prevent GetOrtGlobals() from resurrecting
// the globals during static destruction.
static bool g_process_exiting = false;

// Ensure Shutdown() has been called before process exit
struct EnsureShutdown {
  ~EnsureShutdown() {
    // Set the process-exit flag under the mutex so GetOrtGlobals() (which reads it under the same
    // lock) can never race with this write. Capture whether teardown is needed while holding the
    // lock, then call Shutdown() outside it -- Shutdown() re-acquires the mutex, so holding it here
    // would self-deadlock (the mutex is non-recursive).
    bool needs_shutdown = false;
    {
      std::scoped_lock lock{g_ort_globals_mutex};
      g_process_exiting = true;
      needs_shutdown = static_cast<bool>(g_ort_globals);
    }

    if (needs_shutdown)
      Shutdown();
  }
};

std::unique_ptr<OrtGlobals>& GetOrtGlobals() {
  // Registered once; its destructor runs at process exit (before g_ort_globals is destroyed) and
  // performs the final Shutdown().
  static EnsureShutdown ensure_shutdown;
  // The OrtGlobals constructor does not re-enter GetOrtGlobals() (it bootstraps the CPU interface
  // via the OrtGlobals member accessor), so acquiring this lock here cannot deadlock. The lock is
  // released before the returned reference is dereferenced by callers, so it does not nest with
  // device_interfaces_mutex_.
  std::scoped_lock lock{g_ort_globals_mutex};
  if (!g_ort_globals && !g_process_exiting)
    g_ort_globals = std::make_unique<OrtGlobals>();

  return g_ort_globals;
}

// Used by Shutdown() to display the counts and types of any leaked objects.
// Leak reporting is gated to debug builds or when logging is enabled, since
// consuming applications often cannot guarantee model teardown before process
// exit, and there is little they can do in response at that point.
static bool ShouldReportLeaks() {
#if !defined(NDEBUG)
  return true;
#else
  return g_log.enabled;
#endif
}

template <typename... Types>
bool LeakTypeList<Types...>::Dump() {
  const bool report = ShouldReportLeaks();
  bool any_leaked = false;
  ((LeakChecked<Types>::Count() != 0
        ? (any_leaked = true,
           report ? (std::cerr << "OGA Error: " << LeakChecked<Types>::Count() << " instances of " << typeid(Types).name() << " were leaked." << std::endl, 0) : 0)
        : 0),
   ...);
  return any_leaked && report;
}

void Shutdown() {
  if (LeakTypes::Dump()) {
    std::cerr << "    Please see the documentation for the API being used to ensure proper cleanup." << std::endl;
  }

  // Reset g_ort_globals directly (rather than through GetOrtGlobals(), which would lazily construct
  // the globals just to immediately tear them down). If genai was never initialized there is nothing
  // to do. ~OrtGlobals clears the graph session cache and the device allocators, tears down the
  // device interfaces (including the RyzenAI EP shutdown), unloads the genai add-on libraries, and
  // finally releases the OrtEnv. All of that runs on both the runtime and the process-exit path; the
  // only difference at process exit is that the OrtEnv itself is leaked rather than destroyed, for
  // the reason described below.
  std::scoped_lock lock{g_ort_globals_mutex};

  if (g_ort_globals && g_process_exiting) {
    // At process exit (the EnsureShutdown static destructor sets g_process_exiting under this lock),
    // destroying the OrtEnv is unsafe: releasing it from within __cxa_finalize drives ORT's
    // Environment/LoggingManager destructor to lock a mutex whose backing static may already be
    // finalized, which throws "mutex lock failed: Invalid argument" -- and a throwing destructor at
    // teardown is uncaught, so the process aborts via std::terminate. Nothing else ~OrtGlobals does
    // (clearing the graph session cache, tearing down device interfaces including the RyzenAI EP
    // shutdown, unloading the CUDA add-on library, or releasing the trivial-session allocators)
    // touches that ORT static state, so only the env itself needs to be leaked: release it from the
    // unique_ptr before ~OrtGlobals runs below, so the rest of teardown still happens in full. The
    // released OrtEnv is simply never destructed -- it stays validly constructed in memory for as
    // long as the remaining teardown code (or anything else) needs it, and the OS reclaims it when
    // the process exits. The default OrtEnv owns no thread pools (genai never asks for global ones),
    // so leaking it keeps no threads alive. g_process_exiting is also set when the genai module is
    // unloaded at runtime (dlclose / FreeLibrary), where destroying the env would in fact be safe;
    // the cost there is one leaked OrtEnv per load/unload cycle, which is preferable to aborting.
    (void)g_ort_globals->env_.release();
  }

  g_ort_globals.reset();
}

OrtEnv& GetOrtEnv() {
  return *GetOrtGlobals()->env_;
}

// Fallback to copy between two separate device buffers by going through CPU memory (slow unless we're the CPU device)
void CopyThroughCpu(DeviceBuffer& dest, size_t begin_dest, DeviceBuffer& source, size_t begin_source, size_t size_in_bytes) {
  source.CopyDeviceToCpu();
  auto source_span = std::span<const uint8_t>(source.p_cpu_ + begin_source, size_in_bytes);
  // If we're overwriting the entire destination
  if (dest.size_in_bytes_ == size_in_bytes)
    dest.AllocateCpu();
  else
    dest.CopyDeviceToCpu();  // Overwriting part of destination, so copy over initial contents first
  std::copy(source_span.begin(), source_span.end(), dest.p_cpu_ + begin_dest);
  dest.CopyCpuToDevice();
}

struct GenaiInterfaceImpl : GenaiInterface {
#if _WIN32
  void* HeapAllocate(size_t size) override { return std::malloc(size); }
  void HeapFree(void* p) override { std::free(p); }
#endif

  void CopyThroughCpu(DeviceBuffer& dest, size_t begin_dest, DeviceBuffer& source, size_t begin_source, size_t size_in_bytes) override {
    return Generators::CopyThroughCpu(dest, begin_dest, source, begin_source, size_in_bytes);
  }

  Generators::LogItems& GetLogItems() override { return g_log; }
  std::ostream& operator_leftshift(std::ostream& stream, Generators::SGR sgr_code) override { return stream << sgr_code; }
  std::ostream& Log(std::string_view label, std::string_view text = {}) override { return Log(label, text); }

  void DumpSpan(std::ostream& stream, std::span<const float> values) override { return Generators::DumpSpan(stream, values); }
  void DumpSpan(std::ostream& stream, std::span<const int> values) override { return Generators::DumpSpan(stream, values); }

  void Sequences_AfterAppendNextTokens(Sequences* p_this, DeviceSpan<int32_t> next_tokens, size_t batch_beam_size) override { return p_this->AfterAppendNextTokens(next_tokens, batch_beam_size); }
  void Sequences_RewindTo(Sequences* p_this, size_t new_length) override { return p_this->RewindTo(new_length); }
} g_genai;

#if defined(_WIN32)
struct LibraryHandle {
  LibraryHandle(const char* filename) {
    auto path = CurrentModulePath() + filename;
    if (!fs::path(path).exists())
      path = filename;
    handle_ = LoadLibrary(path.c_str());
    if (!handle_)
      throw std::runtime_error(std::string("Failed to load library: ") + DetermineLoadLibraryError(filename));
  };

  ~LibraryHandle() { FreeLibrary(handle_); }

  FARPROC __stdcall GetSymbol(const char* name) { return ::GetProcAddress(handle_, name); }

  operator HANDLE() { return handle_; }

 private:
  HMODULE handle_{};
};
#elif defined(__linux__) && !defined(__ANDROID__)
struct LibraryHandle {
  LibraryHandle(const char* filename) {
    // Hosts such as Foundry Local may preload the add-on from a managed bundle outside this module's directory.
    handle_ = dlopen(filename, RTLD_NOW | RTLD_LOCAL | RTLD_NOLOAD);
    if (!handle_) {
      auto path = Ort::GetCurrentModuleDir() + "/" + filename;
      handle_ = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL);
    }
    if (!handle_)
      throw std::runtime_error(std::string("Failed to load library: ") + dlerror());  // dlerror() includes the path
  }
  ~LibraryHandle() {
    dlclose(handle_);
  }

  void* GetSymbol(const char* name) { return ::dlsym(handle_, name); }

  operator void*() { return handle_; }

 private:
  void* handle_{};
};
#else
struct LibraryHandle {
  LibraryHandle(const char* filename) {}
  ~LibraryHandle() {}

  void* GetSymbol(const char* name) { return nullptr; }

  operator bool() { return false; }
};
#endif

OrtGlobals::~OrtGlobals() {
  // Teardown is the reverse of construction: destroy everything that uses env-owned resources
  // (the trivial-session allocators and registered EP libraries) BEFORE the env that owns them.
  // Ending all allocator usage before the env is destroyed keeps this correct even if a future
  // interface or buffer dereferences a cached allocator during its own destruction — we do not
  // rely on "nothing happens to deref it today". Shutdown is single-threaded, so no locks.

  // 1. Sessions that hold env / EP state directly. Models (which own these indirectly) are already
  //    gone per the lifetime contract; drop the genai-side session caches here.
  graph_session_cache_.sessions_.clear();

  // 2. Device interfaces + the CUDA add-on library. Interfaces cache the trivial-session allocator
  //    (via device_allocators_), and the CUDA add-on interface lives inside cuda_library_. Destroy
  //    them here so all allocator usage ends before the env. device_interfaces_ is only a
  //    non-owning index (cleared first); owned_interfaces_ and cuda_library_ are the real owners.
  device_interfaces_.clear();
  owned_interfaces_.clear();
  cuda_library_.reset();

  // 3. The trivial-session env-derived allocators, now unreferenced by any interface. Within each
  //    entry session_ is declared before allocator_, so ~allocator_ runs first.
  for (auto& a : device_allocators_) a = {};

  // 4. Finally the env. If genai held the last reference, ORT destroys the environment here,
  //    unregistering / unloading any still-registered EP libraries — by now nothing references them.
  env_.reset();
}

DeviceInterface* OrtGlobals::LoadCudaInterface(DeviceType type) {
  assert(type == DeviceType::NvTensorRtRtx || type == DeviceType::CUDA);
  try {
    if (!cuda_library_) {
#if defined(_WIN32)
      cuda_library_ = std::make_unique<LibraryHandle>("onnxruntime-genai-cuda.dll");
#elif defined(__linux__) && !defined(__ANDROID__)
      cuda_library_ = std::make_unique<LibraryHandle>("libonnxruntime-genai-cuda.so");
#else
      cuda_library_ = std::make_unique<LibraryHandle>("");
#endif
    }
    if (!*cuda_library_)
      throw std::runtime_error("Shared library load failure (see first error)");

    Generators::DeviceInterface* GetInterface(GenaiInterface * p_genai, const char* deviceType, const OrtApi* ort_api);
    auto get_interface = reinterpret_cast<decltype(&GetInterface)>(cuda_library_->GetSymbol("GetInterface"));
    if (!get_interface)
      throw std::runtime_error("CUDA add-on library does not export GetInterface");

    return get_interface(&g_genai, to_string(type).c_str(), Ort::api);
  } catch (const std::exception& e) {
    throw std::runtime_error("Cuda interface not available: " + std::string(e.what()));
  }
}

DeviceInterface* OrtGlobals::GetDeviceInterface(DeviceType type) {
  std::scoped_lock lock{device_interfaces_mutex_};

#if USE_DML
  // DML is deliberately NOT cached in device_interfaces_. Its interface (g_dml_device) is created
  // lazily by dml/session_options.cpp (it needs the LUID / device_index from the model's provider
  // options) and destroyed per-Model in Model::~Model via CloseDmlInterface() to release DML's
  // background-thread hardware resources promptly. Caching the pointer here would dangle after a
  // DML model is freed (a later DML model would get the stale pointer), so always fetch the current
  // instance instead.
  if (type == DeviceType::DML)
    return GetDmlInterface();
#endif

  auto& slot = device_interfaces_[type];
  if (slot)
    return slot;

  switch (type) {
    case DeviceType::CUDA:
    case DeviceType::NvTensorRtRtx:
      slot = LoadCudaInterface(type);  // Non-owning; the interface is owned by cuda_library_.
      break;
    case DeviceType::WEBGPU:
      owned_interfaces_.push_back(CreateWebGPUInterface());
      slot = owned_interfaces_.back().get();
      break;
    case DeviceType::QnnHtp:
    case DeviceType::QnnGpu:
      owned_interfaces_.push_back(CreateQNNInterface(type));
      slot = owned_interfaces_.back().get();
      break;
    case DeviceType::OpenVINO:
      owned_interfaces_.push_back(CreateOpenVINOInterface());
      slot = owned_interfaces_.back().get();
      break;
    case DeviceType::RyzenAI:
      owned_interfaces_.push_back(CreateRyzenAIInterface(*env_));
      slot = owned_interfaces_.back().get();
      break;
    case DeviceType::CPU:
    default:
      owned_interfaces_.push_back(CreateCpuInterface());
      slot = owned_interfaces_.back().get();
      break;
  }
  return slot;
}

std::string to_string(DeviceType device_type) {
  switch (device_type) {
    case DeviceType::CPU:
      return "CPU";
    case DeviceType::CUDA:
      return "CUDA";
    case DeviceType::DML:
      return "DirectML";
    case DeviceType::WEBGPU:
      return "WebGPU";
    case DeviceType::QnnHtp:
    case DeviceType::QnnGpu:
      return "QnnWithSharedMemory";
    case DeviceType::OpenVINO:
      return "OpenVINO";
    case DeviceType::NvTensorRtRtx:
      return "NvTensorRtRtx";
    case DeviceType::RyzenAI:
      return "RyzenAI";
    default:
      throw std::runtime_error("Unknown device type");
  }
}

DeviceInterface* GetDeviceInterface(DeviceType type) {
  return GetOrtGlobals()->GetDeviceInterface(type);
}

GeneratorParams::GeneratorParams(const Config& config)
    : config{config},
      p_device{GetDeviceInterface(DeviceType::CPU)} {
}

GeneratorParams::GeneratorParams(const Model& model)
    : model_{model.shared_from_this()},
      config{*model_->config_.get()},
      use_graph_capture{IsGraphCaptureEnabled(model.config_->model.decoder.session_options)},
      use_multi_profile{IsMultiProfileEnabled(model.config_->model.decoder.session_options)},
      p_device{model.p_device_scoring_} {
  if (use_graph_capture) {
    max_batch_size = 1;  // set it to 1 by default
  }
}

void GeneratorParams::SetGuidance(std::string_view type, std::string_view data, bool enable_ff_tokens = false) {
  guidance_type = type;
  guidance_data = data;
  guidance_ff_tokens_enabled = enable_ff_tokens;
}

bool GeneratorParams::IsPastPresentShareBufferEnabled(const std::string& model_type) const {
  // past_present_share_buffer is only actually enabled when:
  // 1. The config option is set to true, AND
  // 2. Either num_beams == 1 OR the model is Whisper
  return search.past_present_share_buffer &&
         (search.num_beams == 1 || model_type == "whisper");
}

double GeneratorParams::GetSearchNumber(std::string_view name) const {
  if (name == "batch_size") {
    return static_cast<double>(search.batch_size);
  } else if (name == "chunk_size") {
    return static_cast<double>(search.chunk_size.value_or(0));
  } else if (name == "diversity_penalty") {
    return search.diversity_penalty;
  } else if (name == "length_penalty") {
    return search.length_penalty;
  } else if (name == "max_length") {
    return static_cast<double>(search.max_length);
  } else if (name == "min_length") {
    return static_cast<double>(search.min_length);
  } else if (name == "no_repeat_ngram_size") {
    return static_cast<double>(search.no_repeat_ngram_size);
  } else if (name == "num_beams") {
    return static_cast<double>(search.num_beams);
  } else if (name == "num_return_sequences") {
    return static_cast<double>(search.num_return_sequences);
  } else if (name == "random_seed") {
    return static_cast<double>(search.random_seed);
  } else if (name == "repetition_penalty") {
    return search.repetition_penalty;
  } else if (name == "temperature") {
    return search.temperature;
  } else if (name == "top_k") {
    return static_cast<double>(search.top_k);
  } else if (name == "top_p") {
    return search.top_p;
  } else {
    throw std::runtime_error(std::string(name) + " is an invalid name for GetSearchNumber.");
  }
}

bool GeneratorParams::GetSearchBool(std::string_view name) const {
  if (name == "do_sample") {
    return search.do_sample;
  } else if (name == "early_stopping") {
    return search.early_stopping;
  } else if (name == "past_present_share_buffer") {
    return search.past_present_share_buffer;
  } else {
    throw std::runtime_error(std::string(name) + " is an invalid name for GetSearchBool.");
  }
}

void GeneratorParams::SetSpeculativeNumber(std::string_view name, double value) {
  Generators::SetSpeculativeNumber(speculative, name, value);
}

double GeneratorParams::GetSpeculativeNumber(std::string_view name) const {
  if (name == "max_draft_tokens")
    return static_cast<double>(speculative.max_draft_tokens);
  throw std::runtime_error(std::string(name) + " is an invalid name for GetSpeculativeNumber.");
}

std::unique_ptr<Generator> CreateGenerator(const Model& model, const GeneratorParams& params) {
  return std::make_unique<Generator>(model, params);
}

std::unique_ptr<Search> CreateSearch(const GeneratorParams& params) {
  if (params.search.num_beams > 1)
    return params.p_device->CreateBeam(params);
  return params.p_device->CreateGreedy(params);
}

Generator::Generator(const Model& model, const GeneratorParams& params)
    : model_{model.shared_from_this()},
      generation_telemetry_{model.telemetry_session_id_, ModelType::IsTransducer(model.config_->model.type)} {
  // RNNT and TDT models don't use the traditional search/logits pipeline,
  // so skip the standard validations and just create the state.
  if (ModelType::IsTransducer(model.config_->model.type)) {
    state_ = model.CreateState({}, params);
    transducer_state_ = dynamic_cast<TransducerState*>(state_.get());
    strategy_ = MakeDecodingStrategy(*this);
    LogGeneratorCreate(params);
    return;
  }

  if (params.search.max_length == 0)
    throw std::runtime_error("search max_length is 0");
  if (params.search.max_length > model.config_->model.context_length)
    throw std::runtime_error("max_length (" + std::to_string(params.search.max_length) + ") cannot be greater than model context_length (" + std::to_string(model.config_->model.context_length) + ")");

  constexpr int kMaxBatchSize = 32;
  constexpr int kMaxNumBeams = 32;
  constexpr int kMaxNumBeamsCuda = 32;

  if (params.search.batch_size < 1 || params.search.batch_size > kMaxBatchSize)
    throw std::runtime_error("batch_size (" + std::to_string(params.search.batch_size) + ") must be in [1, " + std::to_string(kMaxBatchSize) + "]");

  const int max_num_beams = (params.search.num_beams > 1 &&
                             (params.p_device->GetType() == DeviceType::CUDA || params.p_device->GetType() == DeviceType::NvTensorRtRtx))
                                ? kMaxNumBeamsCuda
                                : kMaxNumBeams;
  if (params.search.num_beams < 1 || params.search.num_beams > max_num_beams)
    throw std::runtime_error("num_beams (" + std::to_string(params.search.num_beams) + ") must be in [1, " + std::to_string(max_num_beams) + "]");
  if (params.search.num_return_sequences < 1 || params.search.num_return_sequences > params.search.num_beams)
    throw std::runtime_error("num_return_sequences (" + std::to_string(params.search.num_return_sequences) + ") must be in [1, " + std::to_string(params.search.num_beams) + "]");
  if (params.config.model.vocab_size < 1)
    throw std::runtime_error("vocab_size must be 1 or greater, is " + std::to_string(params.config.model.vocab_size));
  // Beam search selects the top 2*num_beams (beam, token) candidates out of
  // num_beams*vocab_size entries in BeamSearch_Cpu::SelectTop, which requires
  // num_beams*vocab_size >= 2*num_beams, i.e. vocab_size >= 2. A smaller
  // vocabulary would drive an out-of-bounds partial_sort.
  if (params.search.num_beams > 1 && params.config.model.vocab_size < 2)
    throw std::runtime_error("vocab_size (" + std::to_string(params.config.model.vocab_size) + ") must be 2 or greater when using beam search (num_beams=" + std::to_string(params.search.num_beams) + ")");

  // eos_token_id values are used directly as indices into the per-token score
  // row (of size vocab_size), e.g. in Search::ApplyMinLength. An out-of-range
  // value would cause an out-of-bounds write, so reject it here.
  for (auto eos_token_id : params.config.model.eos_token_id) {
    if (eos_token_id < 0 || eos_token_id >= params.config.model.vocab_size)
      throw std::runtime_error("eos_token_id (" + std::to_string(eos_token_id) + ") must be in range [0, " + std::to_string(params.config.model.vocab_size) + ") (vocab_size)");
  }

  search_ = CreateSearch(params);
  state_ = model.CreateState(search_->GetSequenceLengths(), params);    // Search sequence lengths set when creating state
  guidance_logits_processor_ = CreateGuidanceLogitsProcessor(*state_);  // Could be nullptr if use_guidance (constrained decoding) is not used

  InitializePhi3RopeThreshold(params);
  InitializeSamplingMethod(params);
  strategy_ = MakeDecodingStrategy(*this);
  LogGeneratorCreate(params);
}

void Generator::LogGeneratorCreate(const GeneratorParams& params) {
  generation_telemetry_.LogGeneratorCreate(
      params.search.batch_size,
      params.search.num_beams,
      params.search.max_length,
      params.search.top_k,
      params.search.top_p,
      params.search.temperature,
      params.search.do_sample,
      params.use_graph_capture,
      !params.guidance_type.empty());
}

void Generator::InitializePhi3RopeThreshold(const GeneratorParams& params) {
  // Phi3 ROPE factor rewind threshold: 4097 for phi3/phimoe, 8193 for phi3small, 0 otherwise
  // TODO: Extend to support batch size > 1, num beams > 1, and multimodal models
  const auto& model_type = model_->config_->model.type;
  if (params.BatchBeamSize() == 1 && model_->p_device_->SupportsPhi3RopeRewind(*model_->config_)) {
    if (model_type == "phi3" || model_type == "phimoe")
      phi3_rope_threshold_ = 4097;
    else if (model_type == "phi3small")
      phi3_rope_threshold_ = 8193;
  }
}

bool Generator::IsGreedySampling() const {
  return sampling_method_ == SamplingMethod::kGreedy;
}

void Generator::InitializeSamplingMethod(const GeneratorParams& params) {
  const auto& search = params.search;
  if (!search.do_sample || search.top_k == 1 || search.temperature == 0) {
    sampling_method_ = SamplingMethod::kGreedy;
  } else {
    if (search.num_beams != 1)
      throw std::runtime_error("TopK and TopP cannot be used with a beam search");
    if (search.top_p < 0.0f || search.top_p > 1.0f)
      throw std::runtime_error("top_p must be between 0.0 and 1.0");
    if (search.top_k < 0)
      throw std::runtime_error("top_k must be 0 or greater");
    if (search.top_k > params.config.model.vocab_size)
      throw std::runtime_error("top_k (" + std::to_string(search.top_k) + ") must be less than or equal to vocab_size (" + std::to_string(params.config.model.vocab_size) + ")");
    if (search.top_p > 0.0f && search.top_p < 1.0f && search.top_k > 1) {
      sampling_method_ = SamplingMethod::kTopKTopP;
    } else if (search.top_k > 1) {
      sampling_method_ = SamplingMethod::kTopK;
    } else {
      sampling_method_ = SamplingMethod::kTopP;
    }
  }
}

Generator::~Generator() = default;

DeviceSpan<int32_t> Generator::AllocateInputIdsOnDevice(cpu_span<const int32_t> input_ids) {
  size_t padded_input_ids_size = input_ids.size();
  if (model_->config_->model.decoder.sliding_window.has_value()) {
    // If the model has a sliding window, pad the input_ids to the next multiple of the window size
    // so that the input_ids can be divided into window size chunks.
    const auto window_size = model_->config_->model.decoder.sliding_window->window_size;

    if (model_->config_->model.decoder.sliding_window->slide_inputs) {
      padded_input_ids_size = ((input_ids.size() + window_size - 1) / window_size) * window_size;
    }
  }

  auto input_ids_device = state_->params_->p_device->Allocate<int32_t>(padded_input_ids_size);
  auto cpu_span = input_ids_device.CpuSpan();
  auto padding_begin = cpu_span.begin();
  auto data_end = cpu_span.end();
  if (model_->config_->model.decoder.sliding_window.has_value() && model_->config_->model.decoder.sliding_window->alignment == "left") {
    padding_begin = cpu_span.begin() + input_ids.size();
    data_end = padding_begin;
  }
  std::fill_n(padding_begin, padded_input_ids_size - input_ids.size(), model_->config_->model.pad_token_id);
  std::copy_backward(input_ids.begin(), input_ids.end(), data_end);
  input_ids_device.CopyCpuToDevice();
  return input_ids_device;
}

void Generator::AppendTokens(cpu_span<const int32_t> input_ids) {
  DurationTrace trace{"Generator::AppendTokens"};

  ThrowErrorIfSessionTerminated(state_->session_terminated_);
  if (input_ids.size() == 0)
    throw std::runtime_error("input_ids is empty");
  if ((input_ids.size() / state_->params_->search.batch_size) + search_->GetSequenceLength() > state_->params_->search.max_length)
    throw std::runtime_error("input_ids size (" + std::to_string(input_ids.size()) + ") + current sequence length (" + std::to_string(search_->GetSequenceLength()) + ") exceeds max length (" + std::to_string(state_->params_->search.max_length) + ")");
  if (search_->GetSequenceLength() != 0 && state_->params_->search.batch_size > 1)
    throw std::runtime_error("AppendTokens can only be called once for batch_size > 1. To call AppendTokens again, use RewindToLength(0)");

  // Some models fallback to CPU for the attention operator (for example, some decoder-pipeline NPU models).
  // Continuous decoding is supported for this case as the kv cache for such models is always on CPU.
  constexpr std::array<DeviceType, 6> devices_supporting_continuous_decoding{
      DeviceType::CPU,
      DeviceType::CUDA,
      DeviceType::WEBGPU,
      DeviceType::OpenVINO,
      DeviceType::NvTensorRtRtx,
      DeviceType::RyzenAI};

  if (search_->GetSequenceLength() != 0 &&
      std::none_of(devices_supporting_continuous_decoding.begin(), devices_supporting_continuous_decoding.end(),
                   [this](DeviceType device_type) { return device_type == state_->model_.p_device_kvcache_->GetType(); }))
    // Support for continuous decoding should be based on the type of device used for KV cache
    throw std::runtime_error("Continuous decoding is not supported on the selected device type (" + to_string(state_->model_.p_device_kvcache_->GetType()) +
                             "). Please recreate the generator instance to avoid using continuous decoding.");

  std::string_view append_input_modality = transducer_state_ ? "audio" : "text";
  if (generation_telemetry_.BeginAppend()) {
    bool used_vision = false, used_audio = false;
    for (const auto& extra : extra_inputs_) {
      std::string name = extra.name;
      std::transform(name.begin(), name.end(), name.begin(),
                     [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
      if (name.find("image") != std::string::npos || name.find("pixel") != std::string::npos ||
          name.find("vision") != std::string::npos)
        used_vision = true;
      else if (name.find("audio") != std::string::npos)
        used_audio = true;
    }
    append_input_modality = (used_vision && used_audio) ? "multimodal"
                            : used_vision               ? "vision"
                            : used_audio                ? "audio"
                                                        : "text";
  }

  // Set any extra inputs (those defined in extra_inputs and those defined in the PresetExtraInputs registry)
  if (set_extra_inputs_) {
    state_->SetExtraInputs(extra_inputs_);
    set_extra_inputs_ = false;
  }

  // Continuous decoding - let the decoding strategy realign any deferred per-round state with the
  // committed sequence before we append on top (nothing for standard decoding + initial
  // prefill; speculative -> reconcile two inner KV caches).
  strategy_->PrepareForAppend(*this);

  auto input_ids_device = AllocateInputIdsOnDevice(input_ids);
  search_->AppendTokens(input_ids_device);
  computed_logits_ = false;
  ComputeLogits(input_ids_device);

  generation_telemetry_.CompleteAppend(
      input_ids.size(), state_->params_->search.num_beams, append_input_modality);
}

void Generator::SetInputs(const NamedTensors& named_tensors) {
  if (ModelType::IsLLM(model_->config_->model.type) || ModelType::IsPipe(model_->config_->model.type)) {
    throw std::runtime_error("Please use generator.AppendTokens for " + model_->config_->model.type + ". SetInputs is not supported for this model type.");
  }

  cpu_span<int32_t> input_ids;
  for (const auto& [name, tensor] : named_tensors) {
    if (name == Config::Defaults::InputIdsName) {
      input_ids = cpu_span<int32_t>(tensor->ort_tensor_->GetTensorMutableData<int32_t>(),
                                    tensor->ort_tensor_->GetTensorTypeAndShapeInfo()->GetElementCount());
    } else {
      // If the nominal name is found in the map, use the graph name.
      // Else, use the nominal name as the graph name.
      [[maybe_unused]] const auto [graph_name, found] = model_->config_->GetGraphName(name);
      extra_inputs_.push_back({graph_name, tensor});
    }
  }

  // Set any extra inputs (those defined in extra_inputs and those defined in the PresetExtraInputs registry)
  if (set_extra_inputs_) {
    state_->SetExtraInputs(extra_inputs_);
    set_extra_inputs_ = false;
  }

  // Append tokens and run ComputeLogits after setting all other possible inputs
  if (input_ids.size() > 0) {
    AppendTokens(input_ids);
  }
}

void Generator::ComputeLogits(DeviceSpan<int32_t> next_tokens) {
  if (computed_logits_)
    throw std::runtime_error("ComputeLogits called again without calling AppendTokens or GenerateNextToken first");

  if (guidance_logits_processor_ && last_action_ == Action::generated) {
    auto next_tokens_span = next_tokens.CopyDeviceToCpu();
    guidance_logits_processor_->CommitTokens(next_tokens_span);
  }

  auto logits = state_->Run(search_->GetSequenceLength(), next_tokens, search_->GetNextIndices());
  if (g_log.enabled && g_log.model_logits) {
    auto& stream = Log("model_logits");
    DumpValues(stream, Ort::TypeToTensorType<float>, logits.CopyDeviceToCpu().data(), logits.size());
    stream << std::endl;
  }
  SetLogits(logits);

  if (guidance_logits_processor_ && last_action_ == Action::generated) {
    auto ff_tokens = guidance_logits_processor_->GetFFTokens(0);
    if (!ff_tokens.empty()) {
      // process fast-forward tokens
      std::span<int32_t> forced_tokens_span{ff_tokens};
      auto forced_tokens = AllocateInputIdsOnDevice(forced_tokens_span);
      search_->AppendTokens(forced_tokens);

      std::span<int32_t> new_next_token_span{ff_tokens};
      auto new_next_token = AllocateInputIdsOnDevice(new_next_token_span);
      logits = state_->Run(search_->GetSequenceLength(), new_next_token, search_->GetNextIndices());
      if (g_log.enabled && g_log.model_logits) {
        auto& stream_ = Log("model_logits");
        DumpValues(stream_, Ort::TypeToTensorType<float>, logits.CopyDeviceToCpu().data(), logits.size());
        stream_ << std::endl;
      }
      SetLogits(logits);
      generation_telemetry_.OnTokenGenerated(static_cast<int64_t>(ff_tokens.size()));
    }
  }

  last_action_ = Action::standard;
  computed_logits_ = true;
}

void Generator::SetRuntimeOption(const char* key, const char* value) {
  // Nemotron speech models support per-generator "lang_id" override so that
  // a single loaded model can serve generators in different languages.
  if (ModelType::IsRNNT(model_->config_->model.type) && key != nullptr && std::strcmp(key, "lang_id") == 0) {
    if (value == nullptr) {
      throw std::runtime_error("SetRuntimeOption(\"lang_id\"): value must not be null.");
    }
    size_t pos = 0;
    int lang_id = std::stoi(value, &pos);
    if (pos != std::strlen(value)) {
      throw std::runtime_error(std::string("SetRuntimeOption(\"lang_id\"): invalid integer value \"") + value + "\".");
    }
    static_cast<NemotronSpeechState*>(state_.get())->SetLangId(lang_id);
    return;
  }
  state_->SetRunOption(key, value);
}

size_t Generator::TokenCount() const {
  if (transducer_state_)
    return transducer_state_->TokenCount();
  return static_cast<size_t>(search_->GetSequenceLength());
}

bool Generator::IsDone() {
  ThrowErrorIfSessionTerminated(state_->session_terminated_);

  if (transducer_state_) {
    // Pending mel input means we haven't started processing this chunk yet
    if (!extra_inputs_.empty()) return false;
    return transducer_state_->IsChunkDone();
  }

  if (computed_logits_) {
    return false;
  }

  bool is_done = search_->IsDone();
  if (is_done) {
    state_->Finalize(search_->GetSequenceLength());
    if (guidance_logits_processor_) {
      guidance_logits_processor_->Reset();
      last_action_ = Action::standard;
    }
  }

  return is_done;
}

bool Generator::IsSessionTerminated() const {
  return state_->session_terminated_;
}

void Generator::SetLogits(DeviceSpan<float> logits) {
  search_->SetLogits(logits);
  computed_logits_ = true;
}

void Generator::PrepareForSetLogits() {
  strategy_->PrepareForSetLogits(*this);
}

void Generator::GenerateNextToken() {
  DurationTrace trace{"Generator::GenerateNextToken"};
  ThrowErrorIfSessionTerminated(state_->session_terminated_);

  if (!transducer_state_ &&
      search_->GetSequenceLength() >= state_->params_->search.max_length)
    throw std::runtime_error(
        "GenerateNextToken called with sequence length already at max_length (" +
        std::to_string(state_->params_->search.max_length) + ")");

  strategy_->Step(*this);

  const int64_t active_token_count = transducer_state_
                                         ? static_cast<int64_t>(transducer_state_->GetStepTokens().size())
                                         : static_cast<int64_t>(search_->params_->BatchBeamSize());
  generation_telemetry_.OnTokenGenerated(active_token_count);
}

SpeculativeStats Generator::GetSpeculativeStats() const {
  return strategy_->GetStats();
}

void Generator::RewindToLength(size_t new_length) {
  if (model_->config_->model.type == "whisper" || model_->config_->model.type == "phi3v" || model_->config_->model.type == "decoder-pipeline" || model_->config_->model.type == "lfm2")
    throw std::runtime_error("RewindTo is currently not supported for " + model_->config_->model.type + ".");
  const size_t current_length = search_->GetSequenceLength();
  if (new_length > current_length)
    throw std::runtime_error("Cannot rewind to a length greater than the current sequence length");
  if (new_length == current_length)
    return;
  size_t batch_size = search_->params_->search.batch_size;
  if (batch_size > 1 && new_length != 0)
    throw std::runtime_error("RewindToLength must be called with new_length=0 when batch_size > 1");
  if (search_->params_->search.num_beams > 1)
    throw std::runtime_error("RewindToLength is not supported with beam search");
  const int64_t rewound_token_count =
      static_cast<int64_t>(current_length - new_length) *
      static_cast<int64_t>(search_->params_->BatchBeamSize());
  search_->RewindTo(new_length);
  state_->RewindTo(new_length);
  if (guidance_logits_processor_) {
    guidance_logits_processor_->Reset();
  }
  strategy_->Reset();
  computed_logits_ = false;
  last_action_ = Action::rewound;
  generation_telemetry_.OnRewind(rewound_token_count);
}

DeviceSpan<float> Generator::GetLogits() {
  DeviceSpan<float> strategy_logits;
  if (strategy_->TryGetExternalLogits(*this, strategy_logits)) {
    return strategy_logits;
  }
  if (!computed_logits_) {
    ComputeLogits(search_->GetNextTokens());
  }
  return search_->GetLogits();
}

DeviceSpan<int32_t> Generator::GetSequence(size_t index) const {
  return search_->GetSequence(index);
}

}  // namespace Generators
