// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <memory>
#include <stdexcept>
#include <utility>
#include <cstdint>
#include <cstddef>
#include <limits>
#include "span.h"
#include "ort_genai_c.h"
#include "generators.h"
#include "models/model.h"
#include "constrained_logits_processor.h"
#include "runtime_settings.h"
#include "search.h"
#include "smartptrs.h"
#include "engine/engine.h"
#include "models/streaming_processor.h"
#include "models/nemotron_speech.h"
#include "models/parakeet.h"
#include "models/silero_vad.h"
#include "models/model_package.h"
#include "telemetry/model_telemetry.h"
#include "telemetry/telemetry.h"
#include "dll_load_error.h"

namespace Generators {

struct Result {
  explicit Result(const char* what) : what_{what} {}
  std::string what_;
};

}  // namespace Generators

// Allocate a null terminated C string from a std::string
const char* AllocOgaString(const std::string& string) {
  auto length = string.length() + 1;
  auto cstr_buffer = std::make_unique<char[]>(length);
  memcpy(cstr_buffer.get(), string.c_str(), length);
  return cstr_buffer.release();
}

// This type can't be created or copied by value, only by pointer. It's used for the definitions below to ensure nobody
// accidentally creates/copies one of the types that happens to have a default constructor.
struct OgaAbstract {
  OgaAbstract() = delete;
  OgaAbstract(const OgaAbstract&) = delete;
  void operator=(const OgaAbstract&) = delete;
};

// As the Oga* types are just typedefs, we can use them as the actual types in the C API.
// We still need to cast from internal types to the external ones, but these definitions ensure that the types are correct.
// But do not use reinterpret_cast!
//
// IMPORTANT: Objects handed out through these opaque handles are actually allocated as their
// internal base type (see `ReturnUnique<OgaX>(std::make_unique<BaseX>(...))`). On some ABIs
// (notably MSVC without `__declspec(empty_bases)`), the multiple-inheritance with `OgaAbstract`
// makes `sizeof(OgaX) > sizeof(BaseX)`. If the corresponding `OgaDestroyX` function does
// `delete p` on the `OgaX*`, C++17 sized-delete will pass `sizeof(OgaX)` to the allocator while
// the object was actually allocated with `sizeof(BaseX)` — a new-delete-type-mismatch that
// corrupts heap metadata. Every `OgaDestroyX` below MUST therefore delete through the base type,
// i.e. `delete static_cast<BaseX*>(p)`, to match what was allocated.
struct OgaAdapters : Generators::Adapters, OgaAbstract {};
struct OgaAudios : Generators::Audios, OgaAbstract {};
struct OgaConfig : Generators::Config, OgaAbstract {};
struct OgaGenerator : Generators::Generator, OgaAbstract {};
struct OgaGeneratorParams : Generators::GeneratorParams, OgaAbstract {};
struct OgaImages : Generators::Images, OgaAbstract {};
struct OgaModel : Generators::Model, OgaAbstract {};
struct OgaMultiModalProcessor : Generators::MultiModalProcessor, OgaAbstract {};
struct OgaNamedTensors : Generators::NamedTensors, OgaAbstract {};
struct OgaResult : Generators::Result, OgaAbstract {};
struct OgaRuntimeSettings : Generators::RuntimeSettings, OgaAbstract {};
struct OgaSequences : Generators::TokenSequences, OgaAbstract {};
struct OgaSpeculativeStats : Generators::SpeculativeStats, OgaAbstract {};
struct OgaStringArray : std::vector<std::string>, OgaAbstract {};
struct OgaTensor : Generators::Tensor, OgaAbstract {};
struct OgaTokenizer : Generators::Tokenizer, OgaAbstract {};
struct OgaTokenizerStream : Generators::TokenizerStream, OgaAbstract {};
struct OgaEngine : Generators::Engine, OgaAbstract {};
struct OgaRequest : Generators::Request, OgaAbstract {};
struct OgaStreamingProcessor : Generators::StreamingProcessor, OgaAbstract {};

// Helper function to return a shared pointer as a raw pointer. It won't compile if the types are wrong.
// Exposed types that are internally owned by shared_ptrs inherit from ExternalRefCounted. Then we
// manage external C API ownership through ExternalAddRef/ExternalRelease. This function to return
// a value to an external C API owner does the ExternalAddRef, and the OgaDestroy* method has the
// corresponding ExternalRelease.
template <typename T, typename U>
T* ReturnShared(std::shared_ptr<U>& p) {
  p->ExternalAddRef();
  return static_cast<T*>(p.get());
}

// Helper function to return a unique pointer as a raw pointer. It won't compile if the types are wrong.
template <typename T, typename U>
T* ReturnUnique(std::unique_ptr<U> p) {
  return static_cast<T*>(p.release());
}

extern "C" {

#define OGA_TRY try {
#define OGA_CATCH                                                                   \
  }                                                                                 \
  catch (const std::exception& e) {                                                 \
    return ReturnUnique<OgaResult>(std::make_unique<Generators::Result>(e.what())); \
  }

void OGA_API_CALL OgaShutdown() {
  if (!Generators::GenAiTelemetry::IsDestroyed()) {
    Generators::GenAiTelemetry::Instance().Shutdown();
  }
  Generators::Shutdown();
}

void OGA_API_CALL OgaSetTelemetryEnabled(bool enabled) {
  if (!Generators::GenAiTelemetry::IsDestroyed()) {
    Generators::GenAiTelemetry::Instance().SetEnabled(enabled);
  }
}

const char* OGA_API_CALL OgaResultGetError(const OgaResult* result) {
  return result->what_.c_str();
}

OgaResult* OGA_API_CALL OgaSetLogBool(const char* name, bool value) {
  OGA_TRY
  Generators::SetLogBool(name, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaSetLogString(const char* name, const char* value) {
  OGA_TRY
  // Turn nullptr into an empty std::string (nullptr directly will crash the std::string constructor)
  Generators::SetLogString(name, value ? value : std::string{});
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaSetLogCallback(void (*callback)(const char* string, size_t length)) {
  OGA_TRY
  Generators::SetLogCallback(callback);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateSequences(OgaSequences** out) {
  OGA_TRY
  *out = ReturnUnique<OgaSequences>(std::make_unique<Generators::TokenSequences>());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaAppendTokenSequence(const int32_t* token_ptr, size_t token_cnt, OgaSequences* sequence) {
  OGA_TRY
  std::vector<int32_t> tmp(token_cnt);
  for (size_t i = 0; i < token_cnt; i++) {
    tmp[i] = token_ptr[i];
  }
  sequence->emplace_back(std::move(tmp));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaAppendTokenToSequence(int32_t token, OgaSequences* sequences, size_t sequence_index) {
  OGA_TRY
  if (sequence_index > sequences->size()) {
    throw std::runtime_error("sequence index out of bounds");
  }
  if (sequence_index == sequences->size()) {
    sequences->emplace_back();
  }

  sequences->at(sequence_index).push_back(token);

  return nullptr;
  OGA_CATCH
}

size_t OGA_API_CALL OgaSequencesCount(const OgaSequences* p) {
  return p->size();
}

size_t OGA_API_CALL OgaSequencesGetSequenceCount(const OgaSequences* p, size_t sequence) {
  if (sequence >= p->size()) {
    return 0;
  }
  return (*p)[sequence].size();
}

const int32_t* OGA_API_CALL OgaSequencesGetSequenceData(const OgaSequences* p, size_t sequence) {
  if (sequence >= p->size()) {
    return nullptr;
  }
  return (*p)[sequence].data();
}

OgaResult* OGA_API_CALL OgaLoadImage(const char* image_path, OgaImages** images) {
  OGA_TRY
  const std::vector<const char*> image_paths_vector{image_path};
  *images = ReturnUnique<OgaImages>(Generators::LoadImages(image_paths_vector));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaLoadImages(const OgaStringArray* image_paths, OgaImages** images) {
  OGA_TRY
  std::vector<const char*> image_paths_vector_c;
  for (const auto& image_path : *image_paths)
    image_paths_vector_c.push_back(image_path.c_str());
  *images = ReturnUnique<OgaImages>(Generators::LoadImages(image_paths_vector_c));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaLoadImagesFromBuffers(const void** image_data, const size_t* image_data_sizes, size_t count, OgaImages** images) {
  OGA_TRY
  *images = ReturnUnique<OgaImages>(Generators::LoadImagesFromBuffers(std::span<const void*>(image_data, count), std::span<const size_t>(image_data_sizes, count)));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaLoadAudio(const char* audio_path, OgaAudios** audios) {
  OGA_TRY
  const std::vector<const char*> audio_paths_vector{audio_path};
  *audios = ReturnUnique<OgaAudios>(Generators::LoadAudios(audio_paths_vector));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaLoadAudios(const OgaStringArray* audio_paths, OgaAudios** audios) {
  OGA_TRY
  std::vector<const char*> audio_paths_vector_c;
  for (const auto& audio_path : *audio_paths)
    audio_paths_vector_c.push_back(audio_path.c_str());
  *audios = ReturnUnique<OgaAudios>(Generators::LoadAudios(audio_paths_vector_c));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaLoadAudiosFromBuffers(const void** audio_data, const size_t* audio_data_sizes, size_t count, OgaAudios** audios) {
  OGA_TRY
  *audios = ReturnUnique<OgaAudios>(Generators::LoadAudiosFromBuffers(std::span<const void*>(audio_data, count), std::span<const size_t>(audio_data_sizes, count)));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateRuntimeSettings(OgaRuntimeSettings** out) {
  OGA_TRY
  *out = ReturnUnique<OgaRuntimeSettings>(Generators::CreateRuntimeSettings());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateModelWithRuntimeSettings(const char* config_path, const OgaRuntimeSettings* settings, OgaModel** out) {
  OGA_TRY
  auto model = Generators::CreateModelWithTelemetry(
      [&] { return Generators::CreateModel(Generators::GetOrtEnv(), config_path, settings); });
  *out = ReturnShared<OgaModel>(model);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateConfig(const char* config_path, OgaConfig** out) {
  OGA_TRY
  auto config = Generators::CreateConfig(Generators::GetOrtEnv(), config_path);
  *out = ReturnUnique<OgaConfig>(std::move(config));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateConfigFromPackageEp(const char* config_path, const char* ep, OgaConfig** out) {
  OGA_TRY
  if (!Generators::IsModelPackage(fs::path{config_path})) {
    throw std::runtime_error(std::string("\"") + config_path +
                             "\" is not a model package. Use OgaCreateConfig for a flat model directory.");
  }
  auto config = Generators::CreateConfig(Generators::GetOrtEnv(), config_path, ep);
  *out = ReturnUnique<OgaConfig>(std::move(config));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigClearProviders(OgaConfig* config) {
  OGA_TRY
  Generators::ClearProviders(*config);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigAppendProvider(OgaConfig* config, const char* provider) {
  OGA_TRY
  Generators::SetProviderOption(*config, provider, {}, {});
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigSetProviderOption(OgaConfig* config, const char* provider, const char* key, const char* value) {
  OGA_TRY
  Generators::SetProviderOption(*config, provider, key, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigOverlay(OgaConfig* config, const char* json) {
  OGA_TRY
  Generators::OverlayConfig(*config, json);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigAddModelData(OgaConfig* config, const char* model_filename, const void* model_data, size_t model_data_length) {
  OGA_TRY
  if (model_data == nullptr || model_data_length == 0) {
    throw std::runtime_error("Expected a valid model data pointer and length. Received nullptr or zero length.");
  }

  const auto emplaced = config->model_data_spans_.emplace(model_filename, std::span<const std::byte>(static_cast<const std::byte*>(model_data), model_data_length));
  if (!emplaced.second) {
    throw std::runtime_error("Model data for '" + std::string(model_filename) +
                             "' was already added previously. "
                             "If you want to replace it, please remove it first.");
  }

  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigRemoveModelData(OgaConfig* config, const char* model_filename) {
  OGA_TRY
  auto it = config->model_data_spans_.find(model_filename);
  if (it == config->model_data_spans_.end()) {
    throw std::runtime_error("Model data for '" + std::string(model_filename) + "' was not found.");
  }
  config->model_data_spans_.erase(it);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigSetDecoderProviderOptionsHardwareDeviceType(OgaConfig* config, const char* provider, const char* hardware_device_type) {
  OGA_TRY
  Generators::SetDecoderProviderOptionsHardwareDeviceType(*config, provider, hardware_device_type);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigSetDecoderProviderOptionsHardwareDeviceId(OgaConfig* config, const char* provider, uint32_t hardware_device_id) {
  OGA_TRY
  Generators::SetDecoderProviderOptionsHardwareDeviceId(*config, provider, hardware_device_id);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigSetDecoderProviderOptionsHardwareVendorId(OgaConfig* config, const char* provider, uint32_t hardware_vendor_id) {
  OGA_TRY
  Generators::SetDecoderProviderOptionsHardwareVendorId(*config, provider, hardware_vendor_id);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigClearDecoderProviderOptionsHardwareDeviceType(OgaConfig* config, const char* provider) {
  OGA_TRY
  Generators::ClearDecoderProviderOptionsHardwareDeviceType(*config, provider);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigClearDecoderProviderOptionsHardwareDeviceId(OgaConfig* config, const char* provider) {
  OGA_TRY
  Generators::ClearDecoderProviderOptionsHardwareDeviceId(*config, provider);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaConfigClearDecoderProviderOptionsHardwareVendorId(OgaConfig* config, const char* provider) {
  OGA_TRY
  Generators::ClearDecoderProviderOptionsHardwareVendorId(*config, provider);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateModelFromConfig(const OgaConfig* config, OgaModel** out) {
  OGA_TRY
  auto model = Generators::CreateModelWithTelemetry([&] {
    auto config_copy = std::make_unique<Generators::Config>(*config);
    return Generators::CreateModel(Generators::GetOrtEnv(), std::move(config_copy));
  });
  *out = ReturnShared<OgaModel>(model);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateModel(const char* config_path, OgaModel** out) {
  return OgaCreateModelWithRuntimeSettings(config_path, nullptr, out);
}

OgaResult* OGA_API_CALL OgaModelGetType(const OgaModel* model, const char** out) {
  OGA_TRY
  *out = AllocOgaString(model->config_->model.type.c_str());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaModelGetDeviceType(const OgaModel* model, const char** out) {
  OGA_TRY
  *out = AllocOgaString(to_string(model->p_device_->GetType()));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateGeneratorParams(const OgaModel* model, OgaGeneratorParams** out) {
  OGA_TRY
  auto params = std::make_shared<Generators::GeneratorParams>(*model);
  *out = ReturnShared<OgaGeneratorParams>(params);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaRuntimeSettingsSetHandle(OgaRuntimeSettings* settings, const char* handle_name, void* handle) {
  OGA_TRY
  settings->handles_[handle_name] = handle;
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsSetSearchNumber(OgaGeneratorParams* params, const char* name, double value) {
  OGA_TRY
  Generators::SetSearchNumber(params->search, name, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsSetSearchBool(OgaGeneratorParams* params, const char* name, bool value) {
  OGA_TRY
  Generators::SetSearchBool(params->search, name, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsSetGuidance(OgaGeneratorParams* params, const char* type, const char* data, bool enable_ff_tokens) {
  OGA_TRY
  params->SetGuidance(type, data, enable_ff_tokens);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsGetSearchNumber(const OgaGeneratorParams* params, const char* name, double* value) {
  OGA_TRY
  *value = params->GetSearchNumber(name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsGetSearchBool(const OgaGeneratorParams* params, const char* name, bool* value) {
  OGA_TRY
  *value = params->GetSearchBool(name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsSetSpeculativeNumber(OgaGeneratorParams* params, const char* name, double value) {
  OGA_TRY
  params->SetSpeculativeNumber(name, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGeneratorParamsGetSpeculativeNumber(const OgaGeneratorParams* params, const char* name, double* value) {
  OGA_TRY
  *value = params->GetSpeculativeNumber(name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaCreateGenerator(const OgaModel* model, const OgaGeneratorParams* params, OgaGenerator** out) {
  OGA_TRY
  *out = ReturnUnique<OgaGenerator>(CreateGenerator(*model, *params));
  return nullptr;
  OGA_CATCH
}

bool OGA_API_CALL OgaGenerator_IsDone(OgaGenerator* generator) {
  return generator->IsDone();
}

bool OGA_API_CALL OgaGenerator_IsSessionTerminated(const OgaGenerator* generator) {
  return generator->IsSessionTerminated();
}

OgaResult* OGA_API_CALL OgaGenerator_SetModelInput(OgaGenerator* generator, const char* name, OgaTensor* tensor) {
  OGA_TRY
  generator->extra_inputs_.push_back({std::string{name}, tensor->shared_from_this()});
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_SetInputs(OgaGenerator* generator, const OgaNamedTensors* p_named_tensors) {
  OGA_TRY
  generator->SetInputs(*p_named_tensors);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_AppendTokenSequences(OgaGenerator* generator, const OgaSequences* sequences) {
  OGA_TRY

  if (sequences->empty()) {
    throw std::runtime_error("input sequences are empty");
  } else if (sequences->size() != generator->state_->params_->search.batch_size) {
    throw std::runtime_error("input sequences count does not match batch size");
  }
  std::vector<std::span<const int32_t>> span_sequences;
  for (size_t i = 0; i < sequences->size(); i++) {
    span_sequences.emplace_back((*sequences)[i]);
  }

  auto input_ids = Generators::PadInputs(span_sequences, generator->model_->config_->model.pad_token_id);
  generator->AppendTokens(input_ids);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_AppendTokens(OgaGenerator* generator, const int32_t* input_ids, size_t input_ids_count) {
  OGA_TRY
  generator->AppendTokens(Generators::cpu_span<const int32_t>(input_ids, input_ids_count));
  return nullptr;
  OGA_CATCH
}

size_t OGA_API_CALL OgaGenerator_TokenCount(const OgaGenerator* generator) {
  return generator->TokenCount();
}

OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator) {
  OGA_TRY
  generator->GenerateNextToken();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_GetNextTokens(const OgaGenerator* generator, const int32_t** out, size_t* out_count) {
  OGA_TRY
  // Transducer models (RNNT, TDT) bypass the standard search; return tokens
  // emitted by the most recent StepToken() call.
  if (auto* transducer = dynamic_cast<Generators::TransducerState*>(generator->state_.get())) {
    auto tokens = transducer->GetStepTokens();
    *out = tokens.data();
    *out_count = tokens.size();
    return nullptr;
  }
  auto tokens = generator->search_->GetNextTokens().CopyDeviceToCpu();
  *out = tokens.data();
  *out_count = tokens.size();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_RewindTo(OgaGenerator* generator, size_t new_length) {
  OGA_TRY
  generator->RewindToLength(new_length);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_SetRuntimeOption(OgaGenerator* generator, const char* key, const char* value) {
  OGA_TRY
  generator->SetRuntimeOption(key, value);
  return nullptr;
  OGA_CATCH
}

namespace {

/**
 * \brief Returns a copy of the model output identified by the given name as an OgaTensor on CPU. The buffer is owned by returned OgaTensor
 *       and will be released when the OgaTensor is destroyed
 * \param[in] oga_generator The generator to run the GetInput or GetOutput method on the name provided and the out pointer to store the output.
 * \param[in] name The name of the tensor.
 * \param[in] is_input Whether the tensor name is for an input or not.
 * \param[out] out The returned OgaTensor.
 * \return OgaResult containing the error message if the computation failed.
 */
OgaResult* OGA_API_CALL OgaGenerator_GetInputOutput(const OgaGenerator* oga_generator, const char* name, bool is_input, OgaTensor** out) {
  OGA_TRY
  auto& generator = *reinterpret_cast<const Generators::Generator*>(oga_generator);
  auto* ortvalue = is_input ? generator.state_->GetInput(name) : generator.state_->GetOutput(name);
  if (ortvalue == nullptr)
    throw std::runtime_error(std::string("Generator ") + (is_input ? "input '" : "output '") + name + "' was not found.");

  auto type_info = ortvalue->GetTensorTypeAndShapeInfo();
  auto ortvalue_clone = OrtValue::CreateTensor(generator.model_->allocator_cpu_, type_info->GetShape(), type_info->GetElementType());

  // Copy data to ortvalue_clone
  bool is_cpu = ortvalue->GetTensorMemoryInfo().GetDeviceType() == OrtMemoryInfoDeviceType_CPU;
  auto tensor_span = Generators::ByteWrapTensor(is_cpu ? *Generators::GetDeviceInterface(Generators::DeviceType::CPU) : *generator.model_->p_device_, *ortvalue);
  auto copy_span = Generators::ByteWrapTensor(*Generators::GetDeviceInterface(Generators::DeviceType::CPU), *ortvalue_clone);
  copy_span.CopyFrom(tensor_span);

  auto tensor = std::make_shared<Generators::Tensor>(std::move(ortvalue_clone));
  *out = ReturnShared<OgaTensor>(tensor);
  return nullptr;
  OGA_CATCH
}

}  // namespace

OgaResult* OGA_API_CALL OgaGenerator_GetInput(const OgaGenerator* generator, const char* name, OgaTensor** out) {
  return OgaGenerator_GetInputOutput(generator, name, true, out);
}

OgaResult* OGA_API_CALL OgaGenerator_GetOutput(const OgaGenerator* generator, const char* name, OgaTensor** out) {
  return OgaGenerator_GetInputOutput(generator, name, false, out);
}

OgaResult* OGA_API_CALL OgaGenerator_GetLogits(OgaGenerator* generator, OgaTensor** out) {
  OGA_TRY
  auto logits_span = generator->GetLogits();
  const std::array<int64_t, 3> shape{generator->state_->params_->BatchBeamSize(), 1, generator->model_->config_->model.vocab_size};
  std::span<const float> cpu_logits_span = logits_span.CopyDeviceToCpu();
  const size_t expected_size = static_cast<size_t>(shape[0]) * static_cast<size_t>(shape[2]);
  if (cpu_logits_span.size() != expected_size)
    throw std::runtime_error("Generator logits have " + std::to_string(cpu_logits_span.size()) +
                             " elements but the public logits shape requires " +
                             std::to_string(expected_size) + ".");

  // Copy logits to cpu tensor
  std::unique_ptr<OrtValue> ortvalue_clone = OrtValue::CreateTensor<float>(generator->model_->allocator_cpu_, shape);
  auto clone_span = std::span<float>(ortvalue_clone->GetTensorMutableData<float>(), cpu_logits_span.size());
  Generators::copy(cpu_logits_span, clone_span);
  auto tensor = std::make_shared<Generators::Tensor>(std::move(ortvalue_clone));
  *out = ReturnShared<OgaTensor>(tensor);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaGenerator_SetLogits(OgaGenerator* generator, OgaTensor* tensor) {
  OGA_TRY
  generator->PrepareForSetLogits();
  auto logits = generator->search_->GetLogits();
  if (!generator->computed_logits_ && logits.size() != 0) {
    throw std::runtime_error("logits are not computed yet. Please call GenerateNextToken or AppendTokens before calling SetLogits.");
  }
  size_t element_count = tensor->GetElementCount();
  auto new_logits_span = std::span<const float>(tensor->GetData<float>(), element_count);
  if (logits.size() == 0) {
    logits = generator->model_->p_device_inputs_->Allocate<float>(element_count);
    generator->SetLogits(logits);
  } else if (new_logits_span.size() != logits.size()) {
    throw std::runtime_error("Generator::SetLogits passed an array of size " +
                             std::to_string(new_logits_span.size()) + " but should be size " + std::to_string(logits.size()));
  }
  Generators::copy(new_logits_span, logits.CpuSpan());
  logits.CopyCpuToDevice();
  generator->computed_logits_ = true;
  return nullptr;
  OGA_CATCH
}

size_t OGA_API_CALL OgaGenerator_GetSequenceCount(const OgaGenerator* generator, size_t index) {
  // Transducer models (RNNT, TDT) bypass the standard search; the orchestrator
  // state owns the accumulated transcript directly.
  if (auto* transducer = dynamic_cast<Generators::TransducerState*>(generator->state_.get())) {
    return transducer->GetAllTokens().size();
  }
  return generator->GetSequence(static_cast<int>(index)).size();
}

const int32_t* OGA_API_CALL OgaGenerator_GetSequenceData(const OgaGenerator* generator, size_t index) {
  if (auto* transducer = dynamic_cast<Generators::TransducerState*>(generator->state_.get())) {
    return transducer->GetAllTokens().data();
  }
  return generator->GetSequence(static_cast<int>(index)).CopyDeviceToCpu().data();
}

OgaResult* OGA_API_CALL OgaGenerator_GetSpeculativeStats(const OgaGenerator* generator, OgaSpeculativeStats** out) {
  OGA_TRY
  if (!out)
    throw std::invalid_argument("out must not be null.");
  *out = ReturnUnique<OgaSpeculativeStats>(
      std::make_unique<Generators::SpeculativeStats>(generator->GetSpeculativeStats()));
  return nullptr;
  OGA_CATCH
}

void OGA_API_CALL OgaDestroySpeculativeStats(OgaSpeculativeStats* stats) {
  delete static_cast<Generators::SpeculativeStats*>(stats);
}

OgaResult* OGA_API_CALL OgaSpeculativeStatsGetCount(
    const OgaSpeculativeStats* stats, const char* name, uint64_t* value) {
  OGA_TRY
  if (!stats || !name || !value)
    throw std::invalid_argument("stats, name, and value must not be null.");

  const std::string_view key{name};
  if (key == "rounds")
    *value = stats->rounds;
  else if (key == "completed_rounds")
    *value = stats->completed_rounds;
  else if (key == "interrupted_rounds")
    *value = stats->interrupted_rounds;
  else if (key == "active_rounds")
    *value = stats->active_rounds;
  else if (key == "draft_tokens_proposed")
    *value = stats->draft_tokens_proposed;
  else if (key == "draft_tokens_evaluated")
    *value = stats->draft_tokens_evaluated;
  else if (key == "draft_tokens_accepted")
    *value = stats->draft_tokens_accepted;
  else if (key == "correction_tokens")
    *value = stats->correction_tokens;
  else if (key == "bonus_tokens")
    *value = stats->bonus_tokens;
  else if (key == "tokens_queued")
    *value = stats->tokens_queued;
  else if (key == "tokens_emitted")
    *value = stats->tokens_emitted;
  else if (key == "tokens_discarded")
    *value = stats->tokens_discarded;
  else if (key == "tokens_buffered")
    *value = stats->tokens_buffered;
  else if (key == "draft_forward_passes")
    *value = stats->draft_forward_passes;
  else if (key == "target_forward_passes")
    *value = stats->target_forward_passes;
  else
    throw std::runtime_error(std::string(name) + " is an invalid name for OgaSpeculativeStatsGetCount.");
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaSpeculativeStatsGetNumber(
    const OgaSpeculativeStats* stats, const char* name, double* value) {
  OGA_TRY
  if (!stats || !name || !value)
    throw std::invalid_argument("stats, name, and value must not be null.");

  const std::string_view key{name};
  if (key == "total_draft_ms")
    *value = stats->total_draft_ms;
  else if (key == "total_target_ms")
    *value = stats->total_target_ms;
  else if (key == "total_reconciliation_ms")
    *value = stats->total_reconciliation_ms;
  else if (key == "avg_draft_ms_per_token")
    *value = stats->avg_draft_ms_per_token;
  else if (key == "acceptance_rate")
    *value = stats->acceptance_rate;
  else if (key == "avg_draft_tokens_per_round")
    *value = stats->avg_draft_tokens_per_round;
  else if (key == "mean_emitted_tokens_per_round")
    *value = stats->mean_emitted_tokens_per_round;
  else if (key == "expected_tokens_per_round")
    *value = stats->expected_tokens_per_round;
  else if (key == "avg_target_ms_per_round")
    *value = stats->avg_target_ms_per_round;
  else if (key == "target_baseline_ms_per_token")
    *value = stats->target_baseline_ms_per_token;
  else if (key == "target_overhead_ratio")
    *value = stats->target_overhead_ratio;
  else if (key == "estimated_speedup")
    *value = stats->estimated_speedup;
  else if (key == "observed_speedup")
    *value = stats->observed_speedup;
  else
    throw std::runtime_error(std::string(name) + " is an invalid name for OgaSpeculativeStatsGetNumber.");
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaSpeculativeStatsGetBool(
    const OgaSpeculativeStats* stats, const char* name, bool* value) {
  OGA_TRY
  if (!stats || !name || !value)
    throw std::invalid_argument("stats, name, and value must not be null.");

  if (std::string_view{name} == "formula_supported")
    *value = stats->formula_supported != 0;
  else
    throw std::runtime_error(std::string(name) + " is an invalid name for OgaSpeculativeStatsGetBool.");
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateTokenizer(const OgaModel* model, OgaTokenizer** out) {
  OGA_TRY
  auto tokenizer = model->CreateTokenizer();
  *out = ReturnShared<OgaTokenizer>(tokenizer);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaUpdateTokenizerOptions(
    OgaTokenizer* tokenizer,
    const char* const* keys,
    const char* const* values,
    size_t num_options) {
  OGA_TRY

  if (!tokenizer)
    throw std::runtime_error("Tokenizer pointer is null");

  tokenizer->UpdateOptions(keys, values, num_options);

  return nullptr;

  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetBosTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetBosTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetEosTokenIds(const OgaTokenizer* tokenizer, const int32_t** eos_token_ids, size_t* token_count) {
  OGA_TRY
  auto& eos_ids = tokenizer->GetEosTokenIds();
  *eos_token_ids = eos_ids.data();
  *token_count = eos_ids.size();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetPadTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetPadTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetBotTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetBotTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetEotTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetEotTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetBorTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetBorTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerGetEorTokenId(const OgaTokenizer* tokenizer, int32_t* out) {
  OGA_TRY
  *out = tokenizer->GetEorTokenId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerEncode(const OgaTokenizer* tokenizer, const char* str, OgaSequences* sequences) {
  OGA_TRY
  sequences->emplace_back(tokenizer->Encode(str));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerEncodeBatch(const OgaTokenizer* tokenizer, const char** strings, size_t count, OgaTensor** out) {
  OGA_TRY
  if (count > 0 && strings == nullptr)
    throw std::runtime_error("EncodeBatch: strings pointer must not be null when count > 0");
  auto tensor = tokenizer->EncodeBatch(std::span<const char*>(strings, count));
  *out = ReturnShared<OgaTensor>(tensor);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerToTokenId(const OgaTokenizer* tokenizer, const char* str, int32_t* token_id) {
  OGA_TRY
  *token_id = tokenizer->TokenToTokenId(str);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerDecode(const OgaTokenizer* tokenizer, const int32_t* tokens, size_t token_count, const char** out_string) {
  OGA_TRY
  *out_string = AllocOgaString(tokenizer->Decode({tokens, token_count}));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerApplyChatTemplate(const OgaTokenizer* tokenizer, const char* template_str, const char* messages, const char* tools, bool add_generation_prompt, const char** out_string) {
  OGA_TRY
  *out_string = AllocOgaString(tokenizer->ApplyChatTemplate(template_str, messages, tools, add_generation_prompt));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerDecodeBatch(const OgaTokenizer* tokenizer, const OgaTensor* tensor, OgaStringArray** out) {
  OGA_TRY
  auto shape = tensor->GetShape();
  if (shape.size() != 2)
    throw std::runtime_error("Expected a 2D tensor");
  auto strings = tokenizer->DecodeBatch(std::span<const int32_t>{tensor->GetData<int32_t>(), tensor->GetElementCount()}, shape[0]);
  *out = ReturnUnique<OgaStringArray>(std::make_unique<std::vector<std::string>>(std::move(strings)));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorDecode(const OgaMultiModalProcessor* processor, const int32_t* tokens, size_t token_count, const char** out_string) {
  OGA_TRY
  *out_string = AllocOgaString(processor->tokenizer_->Decode({tokens, token_count}));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateTokenizerStream(const OgaTokenizer* p, OgaTokenizerStream** out) {
  OGA_TRY
  *out = ReturnUnique<OgaTokenizerStream>(p->CreateStream());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateTokenizerStreamFromProcessor(const OgaMultiModalProcessor* p, OgaTokenizerStream** out) {
  OGA_TRY
  *out = ReturnUnique<OgaTokenizerStream>(p->tokenizer_->CreateStream());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTokenizerStreamDecode(OgaTokenizerStream* p, int32_t token, const char** out) {
  OGA_TRY
  *out = p->Decode(token).c_str();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateTensorFromBuffer(void* data, const int64_t* shape_dims, size_t shape_dims_count, OgaElementType element_type, OgaTensor** out) {
  OGA_TRY
  auto p_memory_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  auto ort_element_type = static_cast<ONNXTensorElementDataType>(element_type);
  size_t byte_count = Ort::SizeOf(ort_element_type);
  if (shape_dims_count > 0 && shape_dims == nullptr)
    throw std::runtime_error("shape_dims must not be null when shape_dims_count is non-zero");
  auto shape = std::span<const int64_t>{shape_dims, shape_dims_count};
  for (size_t i = 0; i < shape_dims_count; i++) {
    if (shape_dims[i] < 0)
      throw std::runtime_error("shape dimension must be non-negative");
    if (static_cast<uint64_t>(shape_dims[i]) > std::numeric_limits<size_t>::max())
      throw std::runtime_error("shape dimension exceeds size_t range");
    const size_t dim = static_cast<size_t>(shape_dims[i]);
    if (dim != 0 && byte_count > std::numeric_limits<size_t>::max() / dim)
      throw std::runtime_error("tensor byte count overflow");
    byte_count *= dim;
  }
  std::unique_ptr<OrtValue> ort_tensor;
  if (data)
    ort_tensor = OrtValue::CreateTensor(*p_memory_info, data, byte_count, shape, ort_element_type);
  else
    ort_tensor = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), shape, ort_element_type);
  auto tensor = std::make_shared<Generators::Tensor>(std::move(ort_tensor));
  *out = ReturnShared<OgaTensor>(tensor);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTensorGetType(OgaTensor* tensor, OgaElementType* out) {
  OGA_TRY
  *out = static_cast<OgaElementType>(tensor->GetType());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTensorGetShapeRank(OgaTensor* tensor, size_t* out) {
  OGA_TRY
  *out = tensor->GetShape().size();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTensorGetShape(OgaTensor* tensor, int64_t* shape_dims, size_t rank) {
  OGA_TRY
  auto shape = tensor->GetShape();
  if (rank != shape.size())
    throw std::runtime_error("shape_dims_count doesn't match result of OgaTensorGetShapeRank");
  std::copy(shape.begin(), shape.end(), shape_dims);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaTensorGetData(OgaTensor* tensor, void** out) {
  OGA_TRY
  *out = tensor->GetMutableRawData();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateNamedTensors(OgaNamedTensors** out) {
  OGA_TRY
  *out = ReturnUnique<OgaNamedTensors>(std::make_unique<Generators::NamedTensors>());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaNamedTensorsGet(OgaNamedTensors* named_tensors, const char* name, OgaTensor** out) {
  OGA_TRY
  auto iter = named_tensors->find(name);
  if (iter == named_tensors->end())
    *out = nullptr;
  else
    *out = ReturnShared<OgaTensor>(iter->second);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaNamedTensorsSet(OgaNamedTensors* named_tensors, const char* name, OgaTensor* tensor) {
  OGA_TRY(*named_tensors)
  [name] = tensor->shared_from_this();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaNamedTensorsDelete(OgaNamedTensors* named_tensors, const char* name) {
  OGA_TRY
  named_tensors->erase(name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaNamedTensorsCount(const OgaNamedTensors* named_tensors, size_t* out) {
  OGA_TRY
  *out = named_tensors->size();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaNamedTensorsGetNames(const OgaNamedTensors* named_tensors, OgaStringArray** out) {
  OGA_TRY
  auto names = std::make_unique<std::vector<std::string>>();
  for (const auto& pair : *named_tensors)
    names->push_back(pair.first);
  *out = ReturnUnique<OgaStringArray>(std::move(names));
  return nullptr;
  OGA_CATCH
}

OGA_EXPORT OgaResult* OGA_API_CALL OgaSetCurrentGpuDeviceId(int device_id) {
  OGA_TRY
  Ort::SetCurrentGpuDeviceId(device_id);
  return nullptr;
  OGA_CATCH
}

OGA_EXPORT OgaResult* OGA_API_CALL OgaGetCurrentGpuDeviceId(int* device_id) {
  OGA_TRY
  *device_id = Ort::GetCurrentGpuDeviceId();
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateMultiModalProcessor(const OgaModel* model, OgaMultiModalProcessor** out) {
  OGA_TRY
  auto processor = model->CreateMultiModalProcessor();
  *out = ReturnShared<OgaMultiModalProcessor>(processor);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessImages(const OgaMultiModalProcessor* processor, const char* prompt, const OgaImages* images, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Image processor is not available for this model.");

  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompt, images, nullptr));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessImagesAndPrompts(const OgaMultiModalProcessor* processor, const OgaStringArray* prompts, const OgaImages* images, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Image processor is not available for this model.");

  std::vector<const char*> prompts_;
  for (const auto& prompt : *prompts)
    prompts_.push_back(prompt.c_str());
  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompts_, images, nullptr));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessAudios(const OgaMultiModalProcessor* processor, const char* prompt, const OgaAudios* audios, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Audio processor is not available for this model.");

  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompt, nullptr, audios));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessAudiosAndPrompts(const OgaMultiModalProcessor* processor, const OgaStringArray* prompts, const OgaAudios* audios, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Audio processor is not available for this model.");

  std::vector<const char*> prompts_;
  for (const auto& prompt : *prompts)
    prompts_.push_back(prompt.c_str());
  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompts_, nullptr, audios));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessImagesAndAudios(const OgaMultiModalProcessor* processor, const char* prompt, const OgaImages* images,
                                                           const OgaAudios* audios, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Processor is not available for this model.");

  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompt, images, audios));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaProcessorProcessImagesAndAudiosAndPrompts(const OgaMultiModalProcessor* processor, const OgaStringArray* prompts,
                                                                     const OgaImages* images, const OgaAudios* audios, OgaNamedTensors** input_tensors) {
  OGA_TRY
  if (!processor->processor_)
    throw std::runtime_error("Processor is not available for this model.");

  std::vector<const char*> prompts_;
  for (const auto& prompt : *prompts)
    prompts_.push_back(prompt.c_str());
  *input_tensors = ReturnUnique<OgaNamedTensors>(processor->Process(prompts_, images, audios));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateStringArray(OgaStringArray** out) {
  OGA_TRY
  *out = ReturnUnique<OgaStringArray>(std::make_unique<std::vector<std::string>>());
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaCreateStringArrayFromStrings(const char* const* strs, size_t count, OgaStringArray** out) {
  OGA_TRY
  auto string_array = std::make_unique<std::vector<std::string>>();
  for (size_t i = 0; i < count; i++)
    string_array->push_back(strs[i]);
  *out = ReturnUnique<OgaStringArray>(std::move(string_array));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaStringArrayAddString(OgaStringArray* string_array, const char* str) {
  OGA_TRY
  string_array->push_back(str);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaStringArrayGetCount(const OgaStringArray* string_array, size_t* out) {
  OGA_TRY
  *out = string_array->size();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaStringArrayGetString(const OgaStringArray* string_array, size_t index, const char** out) {
  OGA_TRY
  *out = string_array->at(index).c_str();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaCreateAdapters(const OgaModel* model, OgaAdapters** out) {
  OGA_TRY
  auto adapters = std::make_shared<Generators::Adapters>(model);
  *out = ReturnShared<OgaAdapters>(adapters);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaLoadAdapter(OgaAdapters* adapters, const char* adapter_file_path, const char* adapter_name) {
  OGA_TRY
  adapters->LoadAdapter(adapter_file_path, adapter_name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaUnloadAdapter(OgaAdapters* adapters, const char* adapter_name) {
  OGA_TRY
  adapters->UnloadAdapter(adapter_name);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaSetActiveAdapter(OgaGenerator* generator, OgaAdapters* adapters, const char* adapter_name) {
  OGA_TRY
  generator->state_->SetActiveAdapter(adapters, adapter_name);
  generator->LogAdapterActivated();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaCreateEngine(OgaModel* model, OgaEngine** out) {
  OGA_TRY
  auto engine = std::make_shared<Generators::Engine>(model->shared_from_this());
  *out = ReturnShared<OgaEngine>(engine);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaEngineStep(OgaEngine* engine, OgaRequest** request) {
  OGA_TRY
  auto ready_request = engine->Step();
  *request = ready_request ? ReturnShared<OgaRequest>(ready_request) : nullptr;
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaEngineHasPendingRequests(OgaEngine* engine, bool* out) {
  OGA_TRY
  *out = engine->HasPendingRequests();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaEngineAddRequest(OgaEngine* engine, OgaRequest* request) {
  OGA_TRY
  engine->AddRequest(request->shared_from_this());
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaEngineRemoveRequest(OgaEngine* engine, OgaRequest* request) {
  OGA_TRY
  engine->RemoveRequest(request->shared_from_this());
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaCreateRequest(OgaGeneratorParams* params, OgaRequest** out) {
  OGA_TRY
  auto request = std::make_shared<Generators::Request>(params->shared_from_this());
  *out = ReturnShared<OgaRequest>(request);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaRequestAddTokens(OgaRequest* request, const OgaSequences* tokens) {
  OGA_TRY
  if (tokens->size() > 1) {
    throw std::runtime_error("Request can only be created with a single sequence");
  }
  request->AddTokens((*tokens)[0]);
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaRequestHasUnseenTokens(const OgaRequest* request, bool* out) {
  OGA_TRY
  *out = request->HasUnseenTokens();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaRequestGetUnseenToken(OgaRequest* request, int32_t* token) {
  OGA_TRY
  *token = request->UnseenToken();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaRequestIsDone(const OgaRequest* request, bool* out) {
  OGA_TRY
  *out = request->IsDone();
  return nullptr;
  OGA_CATCH
}

OgaResult* OgaRequestSetOpaqueData(OgaRequest* request, void* data) {
  OGA_TRY
  request->SetOpaqueData(data);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaRequestGetOpaqueData(OgaRequest* request, void** data) {
  OGA_TRY
  *data = request->GetOpaqueData();
  return nullptr;
  OGA_CATCH
}

// NOTE: Every `delete` below casts the Oga* handle back to the internal base type that was
// actually allocated by the matching OgaCreate*/ReturnUnique call. Deleting through the Oga*
// pointer directly would trigger a new-delete-type-mismatch (heap corruption) on ABIs where
// sizeof(OgaX) > sizeof(BaseX) — e.g. MSVC does not apply empty-base optimization to the
// `OgaAbstract` secondary base across multiple inheritance by default, so the C++17 sized-delete
// would hand a mismatched allocation size to the allocator.
void OGA_API_CALL OgaDestroyStringArray(OgaStringArray* string_array) { delete static_cast<std::vector<std::string>*>(string_array); }
void OGA_API_CALL OgaDestroyResult(OgaResult* p) { delete static_cast<Generators::Result*>(p); }
void OGA_API_CALL OgaDestroyString(const char* p) { delete[] p; }
void OGA_API_CALL OgaDestroySequences(OgaSequences* p) { delete static_cast<Generators::TokenSequences*>(p); }
void OGA_API_CALL OgaDestroyConfig(OgaConfig* p) { delete static_cast<Generators::Config*>(p); }
void OGA_API_CALL OgaDestroyModel(OgaModel* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyGeneratorParams(OgaGeneratorParams* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyGenerator(OgaGenerator* p) { delete static_cast<Generators::Generator*>(p); }
void OGA_API_CALL OgaDestroyTokenizer(OgaTokenizer* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyTokenizerStream(OgaTokenizerStream* p) { delete static_cast<Generators::TokenizerStream*>(p); }
void OGA_API_CALL OgaDestroyTensor(OgaTensor* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyMultiModalProcessor(OgaMultiModalProcessor* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyImages(OgaImages* p) { delete static_cast<Generators::Images*>(p); }
void OGA_API_CALL OgaDestroyAudios(OgaAudios* p) { delete static_cast<Generators::Audios*>(p); }
void OGA_API_CALL OgaDestroyNamedTensors(OgaNamedTensors* p) { delete static_cast<Generators::NamedTensors*>(p); }
void OGA_API_CALL OgaDestroyAdapters(OgaAdapters* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyRuntimeSettings(OgaRuntimeSettings* p) { delete static_cast<Generators::RuntimeSettings*>(p); }
void OGA_API_CALL OgaDestroyEngine(OgaEngine* p) { p->ExternalRelease(); }
void OGA_API_CALL OgaDestroyRequest(OgaRequest* p) { p->ExternalRelease(); }

#if defined(_MSC_VER)
#pragma warning(push)
// C4297: this extern "C" entry point is assumed non-throwing under /EHc, but it intentionally
// throws to report a registration failure to the caller -- exactly as the ORT registration call it
// wraps already does (that throw is simply not visible to the compiler because it is in another
// translation unit). The exception propagates and is caught by callers normally.
#pragma warning(disable : 4297)
#endif
void OGA_API_CALL OgaRegisterExecutionProviderLibrary(const char* registration_name, const char* library_path) {
  // Pre-flight the provider library before registering it. A failed ORT registration (e.g. a
  // hardware EP whose native runtime cannot initialize because the GPU was removed or disabled) can
  // leave ONNX Runtime with a half-registered library that crashes at environment teardown. By
  // verifying the library and its native dependencies load first, a broken EP fails cleanly here
  // with a descriptive error instead of destabilizing the process (see CanLoadLibrary).
  if (!CanLoadLibrary(library_path))
    throw std::runtime_error("Failed to register execution provider library '" + std::string(registration_name) +
                             "'. " + DetermineLoadLibraryError(library_path) +
                             " The required hardware or driver may be unavailable (for example, a GPU that has been "
                             "removed or disabled).");

  Ort::RegisterExecutionProviderLibrary(&(Generators::GetOrtEnv()), registration_name, fs::path(library_path).c_str());
}
#if defined(_MSC_VER)
#pragma warning(pop)
#endif

void OGA_API_CALL OgaUnregisterExecutionProviderLibrary(const char* registration_name) {
  Ort::UnregisterExecutionProviderLibrary(&(Generators::GetOrtEnv()), registration_name);
}

OgaResult* OGA_API_CALL OgaCreateStreamingProcessor(OgaModel* model, OgaStreamingProcessor** out) {
  OGA_TRY
  auto processor = Generators::CreateStreamingProcessor(*model);
  *out = ReturnUnique<OgaStreamingProcessor>(std::move(processor));
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaStreamingProcessorProcess(OgaStreamingProcessor* processor, const float* audio_data, size_t num_samples, OgaNamedTensors** out) {
  OGA_TRY
  auto result = processor->Process(audio_data, num_samples);
  if (result) {
    *out = ReturnUnique<OgaNamedTensors>(std::move(result));
  } else {
    *out = nullptr;
  }
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaStreamingProcessorFlush(OgaStreamingProcessor* processor, OgaNamedTensors** out) {
  OGA_TRY
  auto result = processor->Flush();
  if (result) {
    *out = ReturnUnique<OgaNamedTensors>(std::move(result));
  } else {
    *out = nullptr;
  }
  return nullptr;
  OGA_CATCH
}

void OGA_API_CALL OgaDestroyStreamingProcessor(OgaStreamingProcessor* p) { delete static_cast<Generators::StreamingProcessor*>(p); }

OgaResult* OGA_API_CALL OgaStreamingProcessorSetOption(OgaStreamingProcessor* processor, const char* key, const char* value) {
  OGA_TRY
  processor->SetOption(key, value);
  return nullptr;
  OGA_CATCH
}

OgaResult* OGA_API_CALL OgaStreamingProcessorGetOption(const OgaStreamingProcessor* processor, const char* key, const char** value) {
  OGA_TRY
  auto result = processor->GetOption(key);
  *value = AllocOgaString(result);
  return nullptr;
  OGA_CATCH
}

}  // extern "C"
