// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
//
// Modifications Copyright (C) 2024-2026 Advanced Micro Devices, Inc. All rights reserved.
// Portions of this file consist of AI generated content.
#include <algorithm>
#include <array>
#include <climits>
#include <filesystem>
#include <functional>
#include <random>
#include <set>
#include <string>
#include <string_view>
#include <thread>
#include <unordered_map>

#include "../generators.h"
#include "../search.h"
#include "../tracing.h"
#include "model.h"
#include "model_package.h"
#include "tokenizer_tag_utils.h"
#include "gpt.h"
#include "decoder_only.h"
#include "speculative_decoding.h"
#include "whisper.h"
#include "parakeet.h"
#include "parakeet_processor.h"
#include "nemotron_speech.h"
#include "multi_modal.h"
#include "lfm2.h"
#include "marian.h"
#include "decoder_only_pipeline.h"
#include "qwen_vl_model.h"
#include "qwen2_5_vl_image_processor.h"
#include "videochat_flash_processor.h"
#include "mistral3_image_processor.h"
#include "../dml/interface.h"
#include "../openvino/interface.h"
#include "../qnn/interface.h"
#include "../ryzenai/interface.h"
#include "session_options.h"

namespace Generators {

namespace {

constexpr const char* kOrtSessionOptionsModelExternalInitializersFileFolderPath =
    "session.model_external_initializers_file_folder_path";
constexpr const char* kOrtSessionOptionEpContextFilePath = "ep.context_file_path";

// Session-option config keys whose values are file/folder path references. When a model is loaded
// from a package these may be sha256: shared-asset URIs or relative paths, so their values are
// resolved against the model root (Config::ResolvePath) before reaching ORT.
bool IsPathValuedSessionOption(std::string_view key) {
  return key == kOrtSessionOptionsModelExternalInitializersFileFolderPath ||
         key == kOrtSessionOptionEpContextFilePath;
}

}  // namespace

State::State(const GeneratorParams& params, const Model& model)
    : model_{model},
      params_{params.shared_from_this()},
      run_options_{OrtRunOptions::Create()},
      extra_outputs_{*this} {
  // Generate a random id for graph capture
  if (params_->use_graph_capture) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(1, INT_MAX);
    graph_id_value_ = dis(gen);
    graph_id_ = std::to_string(graph_id_value_);
  }
}

void State::DumpInputs() {
  if (g_log.enabled && g_log.model_input_values) {
    auto& stream = Log("model_input_values");
    stream << std::endl;
    DumpTensors(model_, stream, inputs_.data(), input_names_.data(), input_names_.size(), true);
  }

  if (g_log.enabled && g_log.model_output_shapes) {
    auto& stream = Log("model_output_shapes");
    stream << std::endl;
    DumpTensors(model_, stream, outputs_.data(), output_names_.data(), output_names_.size(), false);
  }
}

void State::DumpOutputs() {
  if (g_log.enabled && g_log.model_output_values) {
    auto& stream = Log("model_output_values");
    stream << std::endl;
    DumpTensors(model_, stream, outputs_.data(), output_names_.data(), output_names_.size(), true);
  }
}

void State::Run(OrtSession& session, bool graph_capture_this_run) {
  DurationTrace trace{"State::Run"};

  if (params_->use_graph_capture) {
    graph_capture_session_ = &session;
    if (graph_capture_this_run) {
      run_options_->AddConfigEntry("gpu_graph_id", graph_id_.c_str());
    } else {
      run_options_->AddConfigEntry("gpu_graph_id", "-1");
    }
  }

  if (first_run_) {
    extra_outputs_.Add(session.GetOutputNames());
    if (params_->use_multi_profile) {
      // Run the context phase profile for the first run
      run_options_->AddConfigEntry("nv_profile_index", "0");
    }
    first_run_ = false;
  } else {
    extra_outputs_.Update();
    if (params_->use_multi_profile) {
      run_options_->AddConfigEntry("nv_profile_index", "1");
    }
  }

  DumpInputs();

  if (!ep_dynamic_options_next_run_.empty()) {
    std::vector<const char*> keys;
    std::vector<const char*> values;
    for (auto& kv_pair : ep_dynamic_options_next_run_) {
      keys.push_back(kv_pair.first.c_str());
      values.push_back(kv_pair.second.c_str());
    }
    session.SetEpDynamicOptions(keys.data(), values.data(), ep_dynamic_options_next_run_.size());
    ep_dynamic_options_next_run_.clear();
  }

  if (model_.p_device_ && model_.p_device_->GetType() == DeviceType::NvTensorRtRtx) {
    run_options_->AddConfigEntry("disable_synchronize_execution_providers", "1");
  }

  session.Run(run_options_.get(), input_names_.data(), inputs_.data(), input_names_.size(),
              output_names_.data(), outputs_.data(), output_names_.size());

  extra_outputs_.RegisterOutputs();

  DumpOutputs();
}

void State::SetRunOption(const char* key, const char* value) {
  if (strcmp(key, "terminate_session") == 0) {
    if (strcmp(value, "0") == 0) {
      session_terminated_ = false;
      run_options_->UnsetTerminate();
    } else if (strcmp(value, "1") == 0) {
      session_terminated_ = true;
      run_options_->SetTerminate();
    } else {
      // Value not expected
      throw std::runtime_error(std::string("terminate_session key value unexpected: ") + value);
    }
    return;
  } else if (strcmp(key, "enable_profiling") == 0) {
#if ORT_API_VERSION >= 25
    if (strcmp(value, "0") == 0) {
      run_options_->DisableProfiling();
    } else {
      // Enable run-level profiling. The profiling output file is named: <prefix>_<timestamp>.json
      // Value "1" uses the default prefix; any other value is treated as a custom prefix.
      constexpr const char* default_profile_prefix = "onnxruntime_run_profile";
      const char* prefix = (strcmp(value, "1") == 0) ? default_profile_prefix : value;
      run_options_->EnableProfiling(fs::path(prefix).c_str());
    }
#else
    throw std::runtime_error("enable_profiling requires ONNX Runtime 1.25 or later");
#endif
    return;
  }
  run_options_->AddConfigEntry(key, value);
}

/*
 * Set all run options that are key-value pairs of strings.
 * Reference: https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_run_options_config_keys.h
 */
void State::SetRunOptions(const Config::RunOptions& config_run_options) {
  for (auto& config_entry : config_run_options) {
    SetRunOption(config_entry.first.c_str(), config_entry.second.c_str());
  }
}

OrtValue* State::GetInput(const char* name) {
  ThrowErrorIfSessionTerminated(session_terminated_);
  for (size_t i = 0; i < input_names_.size(); i++) {
    if (std::strcmp(input_names_[i], name) == 0) {
      return inputs_[i];
    }
  }
  return nullptr;
}

OrtValue* State::GetOutput(const char* name) {
  ThrowErrorIfSessionTerminated(session_terminated_);
  for (size_t i = 0; i < output_names_.size(); i++) {
    if (std::strcmp(output_names_[i], name) == 0) {
      return outputs_[i];
    }
  }
  return nullptr;
}

void State::ClearIO() {
  input_names_.clear();
  output_names_.clear();
  inputs_.clear();
  outputs_.clear();
}

void State::SetActiveAdapter(Adapters* adapters, const std::string& adapter_name) {
  if (!adapters_) {
    adapters_ = adapters->shared_from_this();
  } else if (adapters_.get() != adapters) {
    // Two different instances of Adapters are being used. The Generator state can only manage
    // active adapters from a single Adapters container.
    throw std::runtime_error("Generator state can only register a single Adapters container.");
  }

  run_options_->AddActiveLoraAdapter(*adapters_->AcquireAdapter(adapter_name));
  adapter_names_.push_back(adapter_name);
}

State::~State() {
  // Release captured graph resources in the EP.
  // Note: the EP contract is to no-op when the id was never captured (e.g., when
  // graph_capture_this_run was always false for this State). The call is wrapped
  // in try/catch because destructors must not throw -- a throw during unwinding
  // would call std::terminate.
#if ORT_API_VERSION >= 27
  if (graph_capture_session_ && graph_id_value_ > 0) {
    try {
      graph_capture_session_->ReleaseCapturedGraph(graph_id_value_);
    } catch (...) {
      // Best-effort cleanup; swallow to keep the destructor non-throwing.
      if (g_log.enabled && g_log.ort_lib) {
        Log("ort_lib") << "ReleaseCapturedGraph(id=" << graph_id_value_
                       << ") failed: unknown exception" << std::endl;
      }
    }
  }
#endif

  if (adapters_) {
    for (const auto& adapter_name : adapter_names_) {
      adapters_->ReleaseAdapter(adapter_name);
    }
  }
}

std::vector<int32_t> PadInputs(std::span<std::span<const int32_t>> sequences, int32_t pad_token_id) {
  bool pad_right_{true};

  size_t max_length = 0;
  for (auto& sequence : sequences)
    max_length = std::max(max_length, sequence.size());

  std::vector<int32_t> result(max_length * sequences.size());
  std::span<int32_t> result_span(result);

  // Copy and pad the sequences with pad_token_id
  for (size_t i = 0; i < sequences.size(); i++) {
    auto output_span = result_span.subspan(i * max_length, max_length);
    auto input_span = sequences[i];

    auto pad_count = max_length - input_span.size();
    if (pad_right_) {
      std::copy(input_span.begin(), input_span.end(), output_span.begin());
      std::fill(output_span.end() - pad_count, output_span.end(), pad_token_id);
    } else {
      std::fill(output_span.begin(), output_span.begin() + pad_count, pad_token_id);
      std::copy(input_span.begin(), input_span.end(), output_span.begin() + pad_count);
    }
  }

  return result;
}

void CheckResult(extError_t error) {
  if (error != kOrtxOK)
    throw std::runtime_error(OrtxGetLastErrorMessage());
}

TokenizerStream::TokenizerStream(const Tokenizer& tokenizer)
    : tokenizer_{tokenizer.shared_from_this()} {
  CheckResult(OrtxCreate(kOrtxKindDetokenizerCache, cache_.Address()));
}

const std::string& TokenizerStream::Decode(int32_t token) {
  const char* string;
  CheckResult(OrtxDetokenizeCached(tokenizer_->tokenizer_, cache_, token, &string));
  chunk_ = string;
  return chunk_;
}

Tokenizer::Tokenizer(Config& config) : bos_token_id_{config.model.bos_token_id},
                                       eos_token_id_{config.model.eos_token_id},
                                       pad_token_id_{config.model.pad_token_id},
                                       bot_token_id_{config.model.bot_token_id},
                                       eot_token_id_{config.model.eot_token_id},
                                       bor_token_id_{config.model.bor_token_id},
                                       eor_token_id_{config.model.eor_token_id} {
  // Default tokenizer options
  const char* keys[] = {"add_special_tokens", "skip_special_tokens"};
  const char* values[] = {"false", "true"};

  // Resolve tokenizer_dir (may be empty, relative, absolute, or a "sha256:" shared-asset reference).
  const fs::path tokenizer_dir = config.ResolvePath(config.model.tokenizer_dir);
  CheckResult(OrtxCreateTokenizerWithOptions(tokenizer_.Address(), tokenizer_dir.string().c_str(), keys, values, 2));

  // Resolve any unset bot/eot/bor/eor IDs via model-type fallback strings.
  // Resolve any unset bot/eot/bor/eor IDs via model-type fallback.
  if (!bot_token_id_) bot_token_id_ = ResolveFallbackTokenId(config.model.type, std::string(Config::Defaults::BotTokenIdName), *this);
  if (!eot_token_id_) eot_token_id_ = ResolveFallbackTokenId(config.model.type, std::string(Config::Defaults::EotTokenIdName), *this);
  if (!bor_token_id_) bor_token_id_ = ResolveFallbackTokenId(config.model.type, std::string(Config::Defaults::BorTokenIdName), *this);
  if (!eor_token_id_) eor_token_id_ = ResolveFallbackTokenId(config.model.type, std::string(Config::Defaults::EorTokenIdName), *this);
}

int32_t Tokenizer::GetBotTokenId() const {
  if (!bot_token_id_) throw std::runtime_error("bot_token_id is not defined for this model");
  return *bot_token_id_;
}

int32_t Tokenizer::GetEotTokenId() const {
  if (!eot_token_id_) throw std::runtime_error("eot_token_id is not defined for this model");
  return *eot_token_id_;
}

int32_t Tokenizer::GetBorTokenId() const {
  if (!bor_token_id_) throw std::runtime_error("bor_token_id is not defined for this model");
  return *bor_token_id_;
}

int32_t Tokenizer::GetEorTokenId() const {
  if (!eor_token_id_) throw std::runtime_error("eor_token_id is not defined for this model");
  return *eor_token_id_;
}

std::unique_ptr<TokenizerStream> Tokenizer::CreateStream() const {
  return std::make_unique<TokenizerStream>(*this);
}

void Tokenizer::UpdateOptions(const char* const* keys, const char* const* values, size_t num_options) {
  // Tap into ORT Extensions API
  CheckResult(OrtxUpdateTokenizerOptions(tokenizer_, const_cast<const char**>(keys), const_cast<const char**>(values), num_options));
}

std::vector<int32_t> Tokenizer::Encode(const char* text) const {
  OrtxPtr<OrtxTokenId2DArray> ids;
  CheckResult(OrtxTokenize(tokenizer_, &text, 1, ids.Address()));

  const extTokenId_t* tokens;
  size_t count;
  CheckResult(OrtxTokenId2DArrayGetItem(ids, 0, &tokens, &count));
  return {tokens, tokens + count};
}

std::string Tokenizer::Decode(std::span<const int32_t> tokens) const {
  OrtxPtr<OrtxStringArray> ortx_string_array;
  CheckResult(OrtxDetokenize1D(tokenizer_, reinterpret_cast<const uint32_t*>(tokens.data()), tokens.size(), ortx_string_array.Address()));

  const char* string;
  CheckResult(OrtxStringArrayGetItem(ortx_string_array, 0, &string));
  return string;
}

std::string Tokenizer::ApplyChatTemplate(const char* template_str, const char* messages, const char* tools, bool add_generation_prompt) const {
  ort_extensions::OrtxObjectPtr<OrtxTensorResult> templated_text;
  CheckResult(OrtxApplyChatTemplate(tokenizer_, template_str, messages, tools, templated_text.ToBeAssigned(), add_generation_prompt, false /*tokenize*/));

  ort_extensions::OrtxObjectPtr<OrtxTensor> tensor;
  CheckResult(OrtxTensorResultGetAt(templated_text.get(), 0, tensor.ToBeAssigned()));

  const char* text_ptr{};
  CheckResult(OrtxGetTensorData(tensor.get(), reinterpret_cast<const void**>(&text_ptr), nullptr, nullptr));

  return text_ptr;
}

std::vector<int32_t> Tokenizer::EncodeBatch(std::span<const std::string> strings) const {
  std::vector<std::vector<int32_t>> sequences;
  std::vector<std::span<const int32_t>> span_sequences;
  for (size_t i = 0; i < strings.size(); i++) {
    sequences.emplace_back(Encode(strings[i].c_str()));
    span_sequences.emplace_back(sequences.back());
  }

  return PadInputs(span_sequences, pad_token_id_);
}

std::shared_ptr<Tensor> Tokenizer::EncodeBatch(std::span<const char*> strings) const {
  if (strings.empty()) {
    throw std::runtime_error("EncodeBatch: input strings must not be empty");
  }
  for (size_t i = 0; i < strings.size(); i++) {
    if (strings[i] == nullptr) {
      throw std::runtime_error("EncodeBatch: input string at index " + std::to_string(i) + " must not be null");
    }
  }

  std::vector<std::vector<int32_t>> sequences;
  std::vector<std::span<const int32_t>> span_sequences;
  for (size_t i = 0; i < strings.size(); i++) {
    sequences.emplace_back(Encode(strings[i]));
    span_sequences.emplace_back(sequences.back());
  }

  auto encoded = PadInputs(span_sequences, pad_token_id_);  // TODO: Pad directly into tensor vs copying?

  auto shape = std::array<int64_t, 2>{static_cast<int64_t>(strings.size()), static_cast<int64_t>(encoded.size() / strings.size())};
  auto ort_tensor_ = OrtValue::CreateTensor<int32_t>(Ort::Allocator::GetWithDefaultOptions(), shape);
  auto tensor = std::make_shared<Tensor>(std::move(ort_tensor_));
  std::copy(encoded.begin(), encoded.end(), tensor->GetMutableData<int32_t>());

  return tensor;
}

std::vector<std::string> Tokenizer::DecodeBatch(std::span<const int32_t> sequences, size_t count) const {
  if (sequences.size() % count != 0)
    throw std::runtime_error("DecodeBatch: sequences must be evenly divisible by the count");
  size_t sequence_length = sequences.size() / count;
  std::vector<std::string> strings;
  for (size_t i = 0; i < count; i++)
    strings.emplace_back(Decode(sequences.subspan(sequence_length * i, sequence_length)));
  return strings;
}

int32_t Tokenizer::TokenToTokenId(const char* token) const {
  extTokenId_t token_id;
  CheckResult(OrtxConvertTokenToId(tokenizer_, token, &token_id));
  return token_id;
}

// Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
// the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
// has been destroyed. Without this, we will crash in the OnnxRuntime BFCArena code when deleting tensors due to the
// arena already being destroyed.
void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) {
  // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
  // OpenVINO delegates to the CPU device allocator
  auto type = device.GetType();
  if (type == DeviceType::CPU || type == DeviceType::OpenVINO)
    return;

  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
  if (allocator.allocator_)
    return;

  // Allocator lifetime is tied to the execution provider lifetime, which is typically per Session.
  // We create a global dummy Session using a trivial model with the required EP in order to get the allocator so we can
  // re-use it for all models.
  // This ensures memory allocated on-device for model inputs/outputs is valid for the lifetime of GenAI.

  // Names for the device types used by 'SetProviderSessionOptions'
  static const char* device_type_names[] = {"CPU (Not used, see above)", "cuda", "DML", "WebGPU", "QNN", "QNN", "OpenVINO (Not used, see above)", "NvTensorRtRtx", "RyzenAI"};
  static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));

  // Create an OrtSessionOptions and set the options to use the DeviceType we're using here
  auto session_options = OrtSessionOptions::Create();
  std::vector<Config::ProviderOptions> provider_options_list;
  const char* provider_name = device_type_names[static_cast<int>(type)];
  Config::ProviderOptions init_session_provider_options{provider_name, {}};

  // Look up the user-supplied provider options entry for this provider (if any),
  // then let the EP shape the trivial-model init session options. Most EPs use
  // the default no-op; WebGPU forwards global/singleton options and QNN injects
  // the QnnHtpShared allocator gating option.
  const auto& user_provider_options_list = config.model.decoder.session_options.provider_options;
  const auto user_provider_options_it = std::find_if(
      user_provider_options_list.begin(), user_provider_options_list.end(),
      [provider_name](const Config::ProviderOptions& po) { return po.name == provider_name; });
  const Config::ProviderOptions* user_provider_options =
      user_provider_options_it != user_provider_options_list.end() ? &*user_provider_options_it : nullptr;
  if (user_provider_options)
    init_session_provider_options.device_filtering_options = user_provider_options->device_filtering_options;

  device.ShapeInitSessionProviderOptions(init_session_provider_options, user_provider_options);

  provider_options_list.emplace_back(std::move(init_session_provider_options));
  const std::vector<std::string> providers{device_type_names[static_cast<int>(type)]};
  SetProviderSessionOptions(*session_options, providers, provider_options_list, true, config);
  session_options->SetLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);  // Errors only here, as warnings are not useful to the user

  const auto trivial_model = GetTrivialModel();
  allocator.session_ = OrtSession::Create(GetOrtEnv(), trivial_model.data(), trivial_model.size(), session_options.get());

  try {
    auto memory_info = device.GetMemoryInfo();
    allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
  } catch (const Ort::Exception& e) {
    throw std::runtime_error("Failed to create allocator for " + to_string(type) + ": " + std::string(e.what()));
  }
  if (!allocator.allocator_) {
    allocator = {};  // Reset everything just to be safe
    throw std::runtime_error("Unexpected failure to create device memory allocator for " + to_string(type));
  }
  device.InitOrt(*Ort::api, *allocator.allocator_);
}

void SessionInfo::Add(OrtSession& session) {
  auto input_names = session.GetInputNames();
  for (size_t i = 0; i < input_names.size(); i++) {
    auto type_info = session.GetInputTypeInfo(i);
    auto input_type = type_info->GetTensorTypeAndShapeInfo().GetElementType();
    auto found_input = inputs_.find(input_names[i]);
    if (found_input != inputs_.end() && found_input->second->GetTensorTypeAndShapeInfo().GetElementType() != input_type)
      throw std::runtime_error("Model input type mismatch: " + input_names[i] + " expected " +
                               std::to_string(found_input->second->GetTensorTypeAndShapeInfo().GetElementType()) +
                               " got " + std::to_string(input_type));
    inputs_.emplace(std::make_pair(std::move(input_names[i]), std::move(type_info)));
  }

  auto output_names = session.GetOutputNames();
  for (size_t i = 0; i < output_names.size(); i++) {
    auto type_info = session.GetOutputTypeInfo(i);
    outputs_.emplace(std::make_pair(std::move(output_names[i]), std::move(type_info)));
  }
}

bool SessionInfo::HasInput(const std::string& name) const {
  return inputs_.find(name) != inputs_.end();
}

bool SessionInfo::HasOutput(const std::string& name) const {
  return outputs_.find(name) != outputs_.end();
}

ONNXTensorElementDataType SessionInfo::GetInputDataType(const std::string& name) const {
  auto result = inputs_.find(name);
  if (result == inputs_.end())
    throw std::runtime_error("Model input was not found: " + name);
  return result->second->GetTensorTypeAndShapeInfo().GetElementType();
}

ONNXTensorElementDataType SessionInfo::GetOutputDataType(const std::string& name) const {
  auto result = outputs_.find(name);
  if (result == outputs_.end())
    throw std::runtime_error("Model output was not found: " + name);
  return result->second->GetTensorTypeAndShapeInfo().GetElementType();
}

std::vector<std::string> SessionInfo::GetInputNames() const {
  std::vector<std::string> names;
  names.reserve(inputs_.size());
  for (const auto& input : inputs_)
    names.push_back(input.first);
  return names;
}

std::vector<int64_t> SessionInfo::GetInputShape(const std::string& name) const {
  auto type_info = inputs_.find(name);
  if (type_info == inputs_.end())
    throw std::runtime_error("Model input was not found: " + name);
  return type_info->second->GetTensorTypeAndShapeInfo().GetShape();
}

std::vector<int64_t> SessionInfo::GetOutputShape(const std::string& name) const {
  auto type_info = outputs_.find(name);
  if (type_info == outputs_.end())
    throw std::runtime_error("Model output was not found: " + name);
  return type_info->second->GetTensorTypeAndShapeInfo().GetShape();
}

std::vector<const char*> SessionInfo::GetInputSymbolicShape(const std::string& name) const {
  auto type_info = inputs_.find(name);
  if (type_info == inputs_.end())
    throw std::runtime_error("Model input was not found: " + name);
  return type_info->second->GetTensorTypeAndShapeInfo().GetSymbolicDimensions();
}

std::vector<const char*> SessionInfo::GetOutputSymbolicShape(const std::string& name) const {
  auto type_info = outputs_.find(name);
  if (type_info == outputs_.end())
    throw std::runtime_error("Model output was not found: " + name);
  return type_info->second->GetTensorTypeAndShapeInfo().GetSymbolicDimensions();
}

Model::Model(std::unique_ptr<Config> config) : config_{std::move(config)} {
  CreateSessionOptions();
  EnsureDeviceOrtInit(*p_device_, *config_);

  // Only CUDA, TRT-RTX, RyzenAI and DML does every input on the device
  // For WebGPU, use device memory only if graph capture is enabled, otherwise use CPU
  if (p_device_->GetType() == DeviceType::CUDA || p_device_->GetType() == DeviceType::DML || p_device_->GetType() == DeviceType::NvTensorRtRtx ||
      p_device_->GetType() == DeviceType::RyzenAI ||
      (p_device_->GetType() == DeviceType::WEBGPU && IsGraphCaptureEnabled(config_->model.decoder.session_options)))
    p_device_inputs_ = p_device_;
  else
    p_device_inputs_ = GetDeviceInterface(DeviceType::CPU);

  // Search and sampling are performed on the CPU for all device types,
  // except for CUDA and NvTensorRtRtx, where this is performed on the device.
  if (p_device_->GetType() == DeviceType::CUDA ||
      p_device_->GetType() == DeviceType::NvTensorRtRtx)
    p_device_scoring_ = p_device_;
  else
    p_device_scoring_ = GetDeviceInterface(DeviceType::CPU);

  // The kvcache is always allocated in device memory
  p_device_kvcache_ = p_device_;
}

Model::~Model() {
#if USE_DML
  if (p_device_->GetType() == DeviceType::DML) {
    auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(DeviceType::DML)];
    allocator.session_.reset();
    allocator.allocator_.reset();
    session_options_.reset();
    // DML objects are globally scoped and launch background threads that retain hardware resources.
    // These threads persist beyond the lifetime of a Model, preventing proper cleanup and potentially causing deadlocks.
    // To avoid blocking driver threads, we explicitly destroy DML objects when the Model is destroyed.
    // They will be recreated as needed when a new Model is initialized.
    CloseDmlInterface();
  }
#endif
}

void Model::CreateSessionOptionsFromConfig(const Config::SessionOptions& config_session_options,
                                           OrtSessionOptions& session_options,
                                           bool is_primary_session_options,
                                           bool disable_graph_capture) {
  // Default to a limit of 16 threads to optimize performance
  constexpr int min_thread_nums = 1;
  constexpr int max_thread_nums = 16;
  int num_of_cores = std::max(min_thread_nums, static_cast<int>(std::thread::hardware_concurrency() / 2));
  session_options.SetIntraOpNumThreads(std::min(num_of_cores, max_thread_nums));

  if (config_session_options.intra_op_num_threads.has_value()) {
    session_options.SetIntraOpNumThreads(config_session_options.intra_op_num_threads.value());
  }

  if (config_session_options.inter_op_num_threads.has_value()) {
    session_options.SetInterOpNumThreads(config_session_options.inter_op_num_threads.value());
  }

  if (config_session_options.enable_cpu_mem_arena.has_value()) {
    if (config_session_options.enable_cpu_mem_arena.value())
      session_options.EnableCpuMemArena();
    else
      session_options.DisableCpuMemArena();
  }

  if (config_session_options.enable_mem_pattern.has_value()) {
    if (config_session_options.enable_mem_pattern.value())
      session_options.EnableMemPattern();
    else
      session_options.DisableMemPattern();
  }

  if (config_session_options.log_id.has_value()) {
    session_options.SetLogId(config_session_options.log_id.value().c_str());
  }

  if (config_session_options.log_severity_level.has_value()) {
    session_options.SetLogSeverityLevel(config_session_options.log_severity_level.value());
  }

  if (config_session_options.log_verbosity_level.has_value()) {
    session_options.SetLogVerbosityLevel(config_session_options.log_verbosity_level.value());
  }

  if (config_session_options.enable_profiling.has_value()) {
    fs::path profile_file_prefix{config_session_options.enable_profiling.value()};
    session_options.EnableProfiling(profile_file_prefix.c_str());
  }

  /*
   * Set all session options that are key-value pairs of strings.
   * Reference: https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
   */
  for (auto& config_entry : config_session_options.config_entries) {
    if (!config_entry.second.empty() && IsPathValuedSessionOption(config_entry.first)) {
      const std::string resolved = config_->ResolvePath(config_entry.second).string();
      session_options.AddConfigEntry(config_entry.first.c_str(), resolved.c_str());
    } else {
      session_options.AddConfigEntry(config_entry.first.c_str(), config_entry.second.c_str());
    }
  }

  // Register custom ops libraries only if explicitly configured
  if (config_session_options.custom_ops_library.has_value()) {
    // Reference: https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h
    constexpr const char* const library_path_metadata_key_name = "library_path";

    std::string custom_library_file_prefix = config_session_options.custom_ops_library.value();

    std::filesystem::path custom_library_path{custom_library_file_prefix};

    // Reject path traversal components regardless of absolute/relative
    if (custom_library_file_prefix.find("..") != std::string::npos) {
      throw std::runtime_error("custom_ops_library must not contain path traversal (..): " + custom_library_file_prefix);
    }

    // Build the set of allowed directories. The resolved library path must fall
    // within one of these to prevent loading arbitrary libraries from disk.
    std::vector<std::filesystem::path> allowed_dirs;

    // 1. GenAI model folder
    std::error_code ec;
    auto model_dir = std::filesystem::canonical(std::filesystem::path(config_->config_path.string()), ec);
    if (!ec) allowed_dirs.push_back(model_dir);

    // 2. EP library directories
    {
      size_t num_devices = 0;
      const OrtEpDevice* const* device_ptrs = nullptr;
      Ort::GetEpDevices(&GetOrtEnv(), &device_ptrs, &num_devices);

      for (size_t i = 0; i < num_devices; ++i) {
        const OrtKeyValuePairs* keyvals = Ort::GetEpDeviceMetadata(device_ptrs[i]);
        size_t num_entries = 0;
        const char* const* keys = nullptr;
        const char* const* values = nullptr;
        Ort::GetKeyValuePairs(keyvals, &keys, &values, &num_entries);

        for (size_t kvi = 0; kvi < num_entries; kvi++) {
          if (std::string(keys[kvi]) == library_path_metadata_key_name) {
            auto ep_dir = std::filesystem::canonical(std::filesystem::path(values[kvi]).parent_path(), ec);
            if (!ec) allowed_dirs.push_back(ep_dir);
          }
        }
      }
    }

    // 3. Current working directory
    {
      auto cwd = std::filesystem::current_path(ec);
      if (!ec) {
        auto cwd_canonical = std::filesystem::canonical(cwd, ec);
        if (!ec) allowed_dirs.push_back(cwd_canonical);
      }
    }

    // Resolve the library path. For relative paths, search allowed directories in order.
    // For absolute paths, use as-is and validate against the allowlist.
    bool resolved = false;
    if (custom_library_path.is_relative()) {
      for (const auto& dir : allowed_dirs) {
        std::filesystem::path candidate = dir / custom_library_path;
        if (std::filesystem::exists(candidate, ec)) {
          custom_library_file_prefix = std::filesystem::canonical(candidate).string();
          resolved = true;
          break;
        }
      }
    } else {
      // Absolute path — canonicalize and validate below
      auto canonical_path = std::filesystem::canonical(custom_library_path, ec);
      if (!ec && std::filesystem::exists(canonical_path, ec)) {
        custom_library_file_prefix = canonical_path.string();
        resolved = true;
      }
    }

    // Validate that the resolved path is under one of the allowed directories
    if (resolved) {
      std::filesystem::path resolved_canonical =
          std::filesystem::canonical(std::filesystem::path(custom_library_file_prefix), ec);
      if (ec) {
        throw std::runtime_error("Failed to canonicalize custom_ops_library path: " + custom_library_file_prefix);
      }
      bool in_allowed_dir = false;
      auto resolved_str = resolved_canonical.string();
      for (const auto& dir : allowed_dirs) {
        auto dir_str = dir.string();
        // Check that resolved path starts with the allowed directory prefix
        if (resolved_str.size() >= dir_str.size() &&
            resolved_str.compare(0, dir_str.size(), dir_str) == 0 &&
            (resolved_str.size() == dir_str.size() ||
             resolved_str[dir_str.size()] == std::filesystem::path::preferred_separator)) {
          in_allowed_dir = true;
          break;
        }
      }
      if (!in_allowed_dir) {
        throw std::runtime_error(
            "custom_ops_library path is not within an allowed directory: " + custom_library_file_prefix);
      }
    }

    // Convert to fs::path for proper wide string handling on Windows
    fs::path custom_ops_lib_path(custom_library_file_prefix);
    session_options.RegisterCustomOpsLibrary(custom_ops_lib_path.c_str());
  }

  if (config_session_options.graph_optimization_level.has_value()) {
    session_options.SetGraphOptimizationLevel(config_session_options.graph_optimization_level.value());
  }

  auto session_device = SetProviderSessionOptions(session_options, config_session_options.providers,
                                                  config_session_options.provider_options, is_primary_session_options,
                                                  *config_, disable_graph_capture);

  if (!p_device_) {
    p_device_ = session_device;
  } else if (session_device != nullptr && session_device->GetType() != p_device_->GetType()) {
    throw std::runtime_error("Running a model with multiple providers is not supported. Encountered " +
                             to_string(session_device->GetType()) + " and " + to_string(p_device_->GetType()));
  }
}

void Model::CreateSessionOptions() {
  session_options_ = OrtSessionOptions::Create();

  CreateSessionOptionsFromConfig(config_->model.decoder.session_options, *session_options_, true);

  for (auto& pipeline_model : config_->model.decoder.pipeline) {
    if (pipeline_model.session_options.has_value()) {
      auto emplaced = pipeline_session_options_.emplace(pipeline_model.model_id, OrtSessionOptions::Create());
      CreateSessionOptionsFromConfig(*pipeline_model.session_options, *emplaced.first->second, false);
    }
  }

  // Fallback to CPU if no provider specific interface was set
  if (!p_device_)
    p_device_ = GetDeviceInterface(DeviceType::CPU);
}

OrtSessionOptions* Model::GetSessionOptions(const std::string& model_id) const {
  auto session_options = pipeline_session_options_.find(model_id);
  // Use the pipeline model session options id config defined it.
  if (session_options != pipeline_session_options_.end())
    return session_options->second.get();

  // Else fallback to the main session options.
  return session_options_.get();
}

std::unique_ptr<OrtSession> Model::CreateSession(OrtEnv& ort_env, const std::string& model_filename, OrtSessionOptions* session_options) {
  if (auto model_data_it = config_->model_data_spans_.find(model_filename);
      model_data_it != config_->model_data_spans_.end()) {
    // If model data was provided, load the model from memory
    if (model_data_it->second.empty()) {
      throw std::runtime_error("Failed to load model data from memory for " + model_filename);
    }
    // For models loaded from memory that reference external data files, ORT has no model
    // directory to resolve them against. Default the external-initializers folder to the config
    // directory, but only when the config did not already set it (a genai_config session option,
    // resolved in CreateSessionOptionsFromConfig, takes precedence).
    if (!session_options->HasConfigEntry(kOrtSessionOptionsModelExternalInitializersFileFolderPath)) {
      const fs::path external_initializers_path = fs::absolute(config_->config_path);
      session_options->AddConfigEntry(kOrtSessionOptionsModelExternalInitializersFileFolderPath,
                                      external_initializers_path.string().c_str());
    }
    return OrtSession::Create(ort_env, model_data_it->second.data(), model_data_it->second.size(), session_options);
  }

  // Otherwise, load the model from the file system
  return OrtSession::Create(ort_env, (config_->config_path / fs::path(model_filename)).c_str(), session_options);
}

std::shared_ptr<Tokenizer> Model::CreateTokenizer() const {
  return std::make_shared<Tokenizer>(*config_);
}

std::shared_ptr<MultiModalProcessor> Model::CreateMultiModalProcessor() const {
  return std::make_shared<MultiModalProcessor>(*config_, session_info_);
}

bool Model::IsPruned() const {
  const auto& logits_name = config_->model.decoder.outputs.logits;
  if (!session_info_.HasOutput(logits_name))
    return false;
  const auto logits_shape = session_info_.GetOutputShape(logits_name);
  return logits_shape[1] == 1;
}

std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path, const RuntimeSettings* settings /*= nullptr*/) {
  std::string config_overlay;
  if (settings) {
    config_overlay = settings->GenerateConfigOverlay();
  }
  auto config = CreateConfig(ort_env, config_path, /*ep=*/nullptr, config_overlay);
  return CreateModel(ort_env, std::move(config));
}

std::unique_ptr<Config> CreateConfig(OrtEnv& ort_env, const char* config_path, const char* ep,
                                     std::string_view json_overlay) {
  const std::string ep_str = (ep != nullptr) ? std::string{ep} : std::string{};
  const fs::path path{config_path};

  if (IsModelPackage(path)) {
#if ORT_GENAI_HAS_MODEL_PACKAGE
    if (!json_overlay.empty()) {
      throw std::runtime_error(
          "Loading a model package with runtime settings is not supported.");
    }
    auto load = OpenAndSelectVariant(ort_env, path, ep_str);
    auto config = std::make_unique<Config>(load.variant_dir, std::string_view{});
    config->package_root = load.package_root;
    // Delegate genai_config path resolution to ORT's package resolver, capturing the package
    // context to keep it alive for the lifetime of the config.
    auto package_context = load.context;
    config->package_resolver = [package_context](const fs::path& base_dir,
                                                 std::string_view value) -> fs::path {
      return fs::path{package_context->ResolveStringRef(base_dir.string(), std::string{value},
                                                        /*must_exist=*/false)};
    };
    return config;
#else
    throw std::runtime_error(
        "Model package support requires building onnxruntime-genai against "
        "ONNX Runtime with the OrtModelPackageApi experimental functions.");
#endif
  }

  if (!ep_str.empty()) {
    throw std::runtime_error(
        "An execution provider was specified, but \"" + path.string() +
        "\" is not a model package. Drop the ep argument to load a flat directory.");
  }
  return std::make_unique<Config>(path, json_overlay);
}

std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, std::unique_ptr<Config> config) {
  if (config->model.draft)
    return std::make_shared<SpeculativeDecodingModel>(std::move(config), ort_env);
  // Check if it's a pipeline model by checking if decoder.pipeline is configured
  if ((config->model.type == "fara" || config->model.type == "qwen2_5_vl" || config->model.type == "qwen3_vl") && !config->model.decoder.pipeline.empty())
    return std::make_shared<Qwen2_5_VL_PipelineModel>(std::move(config), ort_env);
  if (config->model.type == "lfm2")
    return std::make_shared<LFM2_Model>(std::move(config), ort_env);
  if (config->model.type == "gpt2")
    return std::make_shared<Gpt_Model>(std::move(config), ort_env);
  if (ModelType::IsLLM(config->model.type))
    return std::make_shared<DecoderOnly_Model>(std::move(config), ort_env);
  if (ModelType::IsRNNT(config->model.type))
    return std::make_shared<NemotronSpeechModel>(std::move(config), ort_env);
  if (ModelType::IsTDT(config->model.type))
    return std::make_shared<ParakeetTdtModel>(std::move(config), ort_env);
  if (ModelType::IsALM(config->model.type))
    return std::make_shared<WhisperModel>(std::move(config), ort_env);
  if (ModelType::IsVLM(config->model.type))
    return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, false);
  if (ModelType::IsPipe(config->model.type))
    return std::make_shared<DecoderOnlyPipelineModel>(std::move(config), ort_env);
  if (ModelType::IsMMM(config->model.type)) {
    // Auto-detect speech support: require both the speech ONNX model filename
    // and the preprocessing config to be present. If only one is set, throw
    // a clear error so misconfigurations don't silently disable audio.
    bool has_speech_model = !config->model.speech.filename.empty();
    bool has_speech_config = !config->model.speech.config_filename.empty();
    if (has_speech_model && !has_speech_config) {
      throw std::runtime_error(
          "speech.filename is set but speech.config_filename is missing. "
          "Both are required for audio support.");
    }
    if (!has_speech_model && has_speech_config) {
      throw std::runtime_error(
          "speech.config_filename is set but speech.filename is missing. "
          "Both are required for audio support.");
    }
    return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, has_speech_model);
  }
  if (config->model.type == "marian-ssru")
    return std::make_shared<MarianModel>(std::move(config), ort_env);

  throw std::runtime_error("Unsupported model_type in config.json: " + config->model.type);
}

std::shared_ptr<GeneratorParams> CreateGeneratorParams(const Model& model) {
  return std::make_shared<GeneratorParams>(model);
}

// Used by benchmarking tests only, should not be used normally
std::shared_ptr<GeneratorParams> CreateGeneratorParams(const Config& config) {
  return std::make_shared<GeneratorParams>(config);
}

void Cast(OrtValue& input, std::unique_ptr<OrtValue>& output, DeviceInterface& device, ONNXTensorElementDataType output_type) {
  auto input_info = input.GetTensorTypeAndShapeInfo();
  auto shape = input_info->GetShape();

  if (output) {
    auto output_info = output->GetTensorTypeAndShapeInfo();
    if (shape != output_info->GetShape() || output_type != output_info->GetElementType())
      output = nullptr;
  }
  if (!output)
    output = OrtValue::CreateTensor(device.GetAllocator(), shape, output_type);

  auto input_type = input_info->GetElementType();
  auto element_count = input_info->GetElementCount();
  if (element_count != output->GetTensorTypeAndShapeInfo()->GetElementCount())
    throw std::runtime_error("Cast: input and output element count mismatch");

  void* input_data = input.GetTensorMutableRawData();
  void* output_data = output->GetTensorMutableRawData();
  if (!device.Cast(input_data, output_data, input_type, output_type, element_count)) {
    auto input_span = ByteWrapTensor(device, input);
    auto output_span = ByteWrapTensor(device, *output);
    input_data = input_span.CopyDeviceToCpu().data();
    output_data = output_span.CopyDeviceToCpu().data();
    GetDeviceInterface(DeviceType::CPU)->Cast(input_data, output_data, input_type, output_type, element_count);
    output_span.CopyCpuToDevice();
  }
}

std::unique_ptr<OrtValue> Model::ExpandInputs(std::unique_ptr<OrtValue>& input, int num_beams) const {
  // Input shape (batch_size, sequence_length). The input is required with data type T.
  // Output shape (batch_size * num_beams, sequence_length)

  // When num_beams == 1, we don't need to expand the input, but the expand has a side effect of copying from
  // CPU memory to device memory, so we can skip if the p_device_inputs_ is the CPU device
  if (num_beams == 1 && p_device_inputs_ == GetDeviceInterface(DeviceType::CPU))
    return std::move(input);

  auto input_type_info = input->GetTensorTypeAndShapeInfo();
  auto element_type = input_type_info->GetElementType();
  auto input_shape = input_type_info->GetShape();
  const int64_t batch_size = input_shape[0];
  const int64_t data_size_bytes = input_type_info->GetElementCount() * Ort::SizeOf(element_type) / batch_size;

  input_shape[0] *= num_beams;

  auto input_span = ByteWrapTensor(*GetDeviceInterface(DeviceType::CPU), *input);
  auto expanded = OrtValue::CreateTensor(p_device_inputs_->GetAllocator(), input_shape, element_type);
  auto expanded_span = ByteWrapTensor(*p_device_inputs_, *expanded);

  // Detect fast & simple copy case
  if (num_beams == 1) {
    expanded_span.CopyFrom(input_span);
  } else {
    // TODO (RyanHill): To avoid cuda uninitialized memory warnings, we should copy input_span to device memory first
    for (int i = 0; i < batch_size; i++) {
      for (int j = 0; j < num_beams; j++) {
        expanded_span.subspan((i * num_beams + j) * data_size_bytes, data_size_bytes).CopyFrom(input_span.subspan(i * data_size_bytes, data_size_bytes));
      }
    }
  }
  return expanded;
}

MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& session_info)
    : tokenizer_{std::make_shared<Tokenizer>(config)},
      processor_factory_{
          {"phi3v", Processor::Create<PhiImageProcessor>},
          {"whisper", Processor::Create<WhisperProcessor>},
          {"parakeet_tdt", Processor::Create<ParakeetTdtProcessor>},
          {"phi4mm", Processor::Create<PhiMultiModalProcessor>},
          {"gemma3", Processor::Create<GemmaImageProcessor>},
          {"gemma4", Processor::Create<Gemma4MultiModalProcessor>},
          {"mistral3", Processor::Create<Mistral3ImageProcessor>},
          {"fara", Processor::Create<QwenImageProcessor>},
          {"qwen2_5_vl", Processor::Create<QwenImageProcessor>},
          {"qwen3_vl", Processor::Create<QwenImageProcessor>},
          {"qwen3_5", Processor::Create<QwenImageProcessor>},
          {"qwen3_5_moe", Processor::Create<QwenImageProcessor>},
          {"videochat_flash_qwen", Processor::Create<VideoChatFlashProcessor>}} {
  auto processor = processor_factory_.find(config.model.type);
  if (processor != processor_factory_.end()) {
    processor_ = processor->second(config, session_info);
  } else {
    throw std::runtime_error("MultiModalProcessor cannot be created. " + config.model.type + " is not a registered multi-modal model type.");
  }
}

std::unique_ptr<NamedTensors> MultiModalProcessor::Process(const std::string& prompt, const Images* images, const Audios* audios) const {
  Payload payload{prompt, {}, images, audios};
  return processor_->Process(*tokenizer_, payload);
}

std::unique_ptr<NamedTensors> MultiModalProcessor::Process(std::span<const char*> prompts, const Images* images, const Audios* audios) const {
  Payload payload{"", prompts, images, audios};
  return processor_->Process(*tokenizer_, payload);
}

}  // namespace Generators
