// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#define OGA_USE_SPAN 1
#include "../models/onnxruntime_api.h"
#include "../ort_genai.h"
#include <iostream>

using namespace pybind11::literals;

// If a parameter to a C++ function is an array of float16, this type will let pybind11::array_t<Ort::Float16_t> map to numpy's float16 format
namespace pybind11 {
namespace detail {
template <>
struct npy_format_descriptor<Ort::Float16_t> {
  static constexpr auto name = _("float16");
  static pybind11::dtype dtype() {
    handle ptr = npy_api::get().PyArray_DescrFromType_(23 /*NPY_FLOAT16*/); /* import numpy as np; print(np.dtype(np.float16).num */
    return reinterpret_borrow<pybind11::dtype>(ptr);
  }
  static std::string format() {
    // following: https://docs.python.org/3/library/struct.html#format-characters
    return "e";
  }
};
}  // namespace detail
}  // namespace pybind11

template <typename T>
std::span<T> ToSpan(pybind11::array_t<T> v) {
  if constexpr (std::is_const_v<T>)
    return {v.data(), static_cast<size_t>(v.size())};
  else
    return {v.mutable_data(), static_cast<size_t>(v.size())};
}

template <typename T>
std::span<T> ToSpan(OgaTensor& v) {
  assert(static_cast<ONNXTensorElementDataType>(v.Type()) == Ort::TypeToTensorType<T>);
  auto shape = v.Shape();
  auto element_count = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
  return {reinterpret_cast<T*>(v.Data()), static_cast<size_t>(element_count)};
}

template <typename T>
pybind11::array_t<T> ToPython(std::span<T> v) {
  return pybind11::array_t<T>{{v.size()}, {sizeof(T)}, v.data()};
}

ONNXTensorElementDataType ToTensorType(const pybind11::dtype& type) {
  switch (type.num()) {
    case pybind11::detail::npy_api::NPY_BOOL_:
      return Ort::TypeToTensorType<bool>;
    case pybind11::detail::npy_api::NPY_UINT8_:
      return Ort::TypeToTensorType<uint8_t>;
    case pybind11::detail::npy_api::NPY_INT8_:
      return Ort::TypeToTensorType<int8_t>;
    case pybind11::detail::npy_api::NPY_UINT16_:
      return Ort::TypeToTensorType<uint16_t>;
    case pybind11::detail::npy_api::NPY_INT16_:
      return Ort::TypeToTensorType<int16_t>;
    case pybind11::detail::npy_api::NPY_UINT32_:
      return Ort::TypeToTensorType<uint32_t>;
    case pybind11::detail::npy_api::NPY_INT32_:
      return Ort::TypeToTensorType<int32_t>;
    case pybind11::detail::npy_api::NPY_UINT64_:
      return Ort::TypeToTensorType<uint64_t>;
    case pybind11::detail::npy_api::NPY_INT64_:
      return Ort::TypeToTensorType<int64_t>;
    case 23 /*NPY_FLOAT16*/:
      return Ort::TypeToTensorType<Ort::Float16_t>;
    case pybind11::detail::npy_api::NPY_FLOAT_:
      return Ort::TypeToTensorType<float>;
    case pybind11::detail::npy_api::NPY_DOUBLE_:
      return Ort::TypeToTensorType<double>;
    default:
      throw std::runtime_error("Unsupported numpy type");
  }
}

int ToNumpyType(ONNXTensorElementDataType type) {
  switch (type) {
    case Ort::TypeToTensorType<bool>:
      return pybind11::detail::npy_api::NPY_BOOL_;
    case Ort::TypeToTensorType<int8_t>:
      return pybind11::detail::npy_api::NPY_INT8_;
    case Ort::TypeToTensorType<uint8_t>:
      return pybind11::detail::npy_api::NPY_UINT8_;
    case Ort::TypeToTensorType<int16_t>:
      return pybind11::detail::npy_api::NPY_INT16_;
    case Ort::TypeToTensorType<uint16_t>:
      return pybind11::detail::npy_api::NPY_UINT16_;
    case Ort::TypeToTensorType<int32_t>:
      return pybind11::detail::npy_api::NPY_INT32_;
    case Ort::TypeToTensorType<uint32_t>:
      return pybind11::detail::npy_api::NPY_UINT32_;
    case Ort::TypeToTensorType<int64_t>:
      return pybind11::detail::npy_api::NPY_INT64_;
    case Ort::TypeToTensorType<uint64_t>:
      return pybind11::detail::npy_api::NPY_UINT64_;
    case Ort::TypeToTensorType<Ort::Float16_t>:
      return 23 /*NPY_FLOAT16*/;
    case Ort::TypeToTensorType<float>:
      return pybind11::detail::npy_api::NPY_FLOAT_;
    case Ort::TypeToTensorType<double>:
      return pybind11::detail::npy_api::NPY_DOUBLE_;
    default:
      throw std::runtime_error("Unsupported onnx type");
  }
}

template <typename... Types>
std::string ToFormatDescriptor(ONNXTensorElementDataType type, Ort::TypeList<Types...>) {
  std::string result;
  (void)((type == Ort::TypeToTensorType<Types> ? result = pybind11::format_descriptor<Types>::format(), true : false) || ...);
  if (result.empty())
    throw std::runtime_error("Unsupported onnx type");
  return result;
}

std::string ToFormatDescriptor(ONNXTensorElementDataType type) {
  return ToFormatDescriptor(type, Ort::TensorTypes{});
}

std::unique_ptr<OgaTensor> ToOgaTensor(pybind11::array& v, bool copy = true) {
  auto type = ToTensorType(v.dtype());

  std::vector<int64_t> shape(v.ndim());
  for (pybind11::ssize_t i = 0; i < v.ndim(); i++)
    shape[i] = v.shape()[i];

  // Check if v is contiguous
  if ((v.flags() & (pybind11::array::c_style | pybind11::array::f_style)) == 0)
    throw std::runtime_error("Array must be contiguous. Please use NumPy's 'ascontiguousarray' method on the value.");

  auto tensor = OgaTensor::Create(copy ? nullptr : v.mutable_data(), shape, static_cast<OgaElementType>(type));
  if (copy) {
    auto ort_data = reinterpret_cast<uint8_t*>(tensor->Data());
    auto python_data = reinterpret_cast<const uint8_t*>(v.data());
    std::copy(python_data, python_data + v.nbytes(), ort_data);
  }
  return tensor;
}

pybind11::array ToNumpy(OgaTensor& v) {
  auto shape = v.Shape();
  auto type = static_cast<ONNXTensorElementDataType>(v.Type());
  auto element_size = Ort::SizeOf(type);
  auto data = v.Data();

  std::vector<int64_t> strides(shape.size());
  {
    auto size = element_size;
    for (size_t i = strides.size(); i-- > 0;) {
      strides[i] = size;
      size *= shape[i];
    }
  }

  pybind11::buffer_info bufinfo{
      data,                                          // Pointer to memory buffer
      static_cast<pybind11::ssize_t>(element_size),  // Size of underlying scalar type
      ToFormatDescriptor(type),                      // Python struct-style format descriptor
      static_cast<pybind11::ssize_t>(shape.size()),  // Number of dimensions
      shape,                                         // Buffer dimensions
      strides                                        // Strides (in bytes) for each index
  };

  return pybind11::array{bufinfo};
}

struct PyGeneratorParams {
  PyGeneratorParams(const OgaModel& model) : params_{OgaGeneratorParams::Create(model)} {}

  operator const OgaGeneratorParams&() const { return *params_; }

  std::unique_ptr<OgaGeneratorParams> params_;

  void SetSearchOptions(const pybind11::kwargs& dict) {
    for (auto& entry : dict) {
      auto name = entry.first.cast<std::string>();
      if (pybind11::isinstance<pybind11::float_>(entry.second)) {
        params_->SetSearchOption(name.c_str(), entry.second.cast<double>());
      } else if (pybind11::isinstance<pybind11::bool_>(entry.second)) {
        params_->SetSearchOptionBool(name.c_str(), entry.second.cast<bool>());
      } else if (pybind11::isinstance<pybind11::int_>(entry.second)) {
        params_->SetSearchOption(name.c_str(), entry.second.cast<int>());
      } else
        throw std::runtime_error("Unknown search option type, can be float/bool/int:" + name);
    }
  }

  void SetSpeculativeOptions(const pybind11::kwargs& dict) {
    for (auto& entry : dict) {
      auto name = entry.first.cast<std::string>();
      if (pybind11::isinstance<pybind11::float_>(entry.second)) {
        params_->SetSpeculativeNumber(name.c_str(), entry.second.cast<double>());
      } else if (pybind11::isinstance<pybind11::int_>(entry.second)) {
        params_->SetSpeculativeNumber(name.c_str(), entry.second.cast<int>());
      } else
        throw std::runtime_error("Unknown speculative option type, must be int or float: " + name);
    }
  }

  pybind11::dict GetSpeculativeOptions() {
    pybind11::dict d;
    d["max_draft_tokens"] = params_->GetSpeculativeNumber("max_draft_tokens");
    return d;
  }

  void SetGuidance(const std::string& type, const std::string& data, bool enable_ff_tokens = false) {
    params_->SetGuidance(type.c_str(), data.c_str(), enable_ff_tokens);
  }

  pybind11::dict GetSearchOptions() {
    pybind11::dict d;
    d["batch_size"] = params_->GetSearchNumber("batch_size");
    d["chunk_size"] = params_->GetSearchNumber("chunk_size");
    d["diversity_penalty"] = params_->GetSearchNumber("diversity_penalty");
    d["do_sample"] = params_->GetSearchBool("do_sample");
    d["early_stopping"] = params_->GetSearchBool("early_stopping");
    d["length_penalty"] = params_->GetSearchNumber("length_penalty");
    d["max_length"] = params_->GetSearchNumber("max_length");
    d["min_length"] = params_->GetSearchNumber("min_length");
    d["no_repeat_ngram_size"] = params_->GetSearchNumber("no_repeat_ngram_size");
    d["num_beams"] = params_->GetSearchNumber("num_beams");
    d["num_return_sequences"] = params_->GetSearchNumber("num_return_sequences");
    d["past_present_share_buffer"] = params_->GetSearchBool("past_present_share_buffer");
    d["random_seed"] = params_->GetSearchNumber("random_seed");
    d["repetition_penalty"] = params_->GetSearchNumber("repetition_penalty");
    d["temperature"] = params_->GetSearchNumber("temperature");
    d["top_k"] = params_->GetSearchNumber("top_k");
    d["top_p"] = params_->GetSearchNumber("top_p");
    return d;
  }

  std::vector<pybind11::object> refs_;  // References to data we want to ensure doesn't get garbage collected
};

struct PyGenerator {
  PyGenerator(const OgaModel& model, PyGeneratorParams& params) {
    generator_ = OgaGenerator::Create(model, *params.params_);
  }

  pybind11::array_t<int32_t> GetNextTokens() {
    return ToPython(generator_->GetNextTokens());
  }

  pybind11::array_t<int32_t> GetSequence(int index) {
    return ToPython(generator_->GetSequence(index));
  }

  pybind11::array GetInput(const std::string& name) {
    return ToNumpy(*generator_->GetInput(name.c_str()));
  }

  pybind11::array GetOutput(const std::string& name) {
    return ToNumpy(*generator_->GetOutput(name.c_str()));
  }

  void SetModelInput(const std::string& name, pybind11::array& value) {
    generator_->SetModelInput(name.c_str(), *ToOgaTensor(value, false));
  }

  void SetInputs(OgaNamedTensors& named_tensors) {
    generator_->SetInputs(named_tensors);
  }

  void AppendTokens(OgaTensor& tokens) {
    generator_->AppendTokens(ToSpan<int32_t>(tokens));
  }

  void AppendTokens(pybind11::array_t<int32_t>& tokens) {
    generator_->AppendTokens(ToSpan(tokens));
  }

  size_t TokenCount() const {
    return generator_->TokenCount();
  }

  pybind11::array_t<float> GetLogits() {
    return ToNumpy(*generator_->GetLogits());
  }

  void SetLogits(pybind11::array_t<float> new_logits) {
    generator_->SetLogits(*ToOgaTensor(new_logits, false));
  }

  void GenerateNextToken() {
    generator_->GenerateNextToken();
  }

  void RewindTo(size_t new_length) {
    generator_->RewindTo(new_length);
  }

  bool IsDone() {
    return generator_->IsDone();
  }

  void SetActiveAdapter(OgaAdapters& adapters, const std::string& adapter_name) {
    generator_->SetActiveAdapter(adapters, adapter_name.c_str());
  }

  void SetRuntimeOption(const std::string& key, const std::string& value) {
    generator_->SetRuntimeOption(key.c_str(), value.c_str());
  }

  pybind11::dict GetSpeculativeStats() {
    auto stats = generator_->GetSpeculativeStats();
    pybind11::dict d;
    for (const char* key : {"rounds", "completed_rounds", "interrupted_rounds", "active_rounds",
                            "draft_tokens_proposed", "draft_tokens_evaluated", "draft_tokens_accepted",
                            "correction_tokens", "bonus_tokens", "tokens_queued", "tokens_emitted",
                            "tokens_discarded", "tokens_buffered", "draft_forward_passes",
                            "target_forward_passes"})
      d[key] = stats->GetCount(key);
    d["formula_supported"] = stats->GetBool("formula_supported");
    for (const char* key : {"total_draft_ms", "total_target_ms", "total_reconciliation_ms",
                            "avg_draft_ms_per_token", "acceptance_rate", "avg_draft_tokens_per_round",
                            "mean_emitted_tokens_per_round", "expected_tokens_per_round",
                            "avg_target_ms_per_round", "target_baseline_ms_per_token",
                            "target_overhead_ratio", "estimated_speedup", "observed_speedup"})
      d[key] = stats->GetNumber(key);
    return d;
  }

 private:
  std::unique_ptr<OgaGenerator> generator_;
};

void SetLogOptions(const pybind11::kwargs& dict) {
  for (auto& entry : dict) {
    auto name = entry.first.cast<std::string>();
    if (pybind11::isinstance<pybind11::bool_>(entry.second)) {
      Oga::SetLogBool(name.c_str(), entry.second.cast<bool>());
    } else if (pybind11::isinstance<pybind11::str>(entry.second)) {
      Oga::SetLogString(name.c_str(), entry.second.cast<std::string>().c_str());
    } else
      throw std::runtime_error("Unknown log option type, can be bool/string:" + name);
  }
}

void SetLogCallback(std::optional<const pybind11::function> callback) {
  static std::optional<pybind11::function> log_callback;
  log_callback = callback;

  if (log_callback.has_value()) {
    Oga::SetLogCallback([](const char* message, size_t length) {
      pybind11::gil_scoped_acquire gil;
      (*log_callback)(std::string_view(message, length));
    });
  } else {
    Oga::SetLogCallback(nullptr);
  }
}

PYBIND11_MODULE(onnxruntime_genai, m) {
  m.doc() = R"pbdoc(
        Ort Generators library
        ----------------------

        .. currentmodule:: cmake_example

        .. autosummary::
           :toctree: _generate

    )pbdoc";

  // Add a cleanup call to happen before global variables are destroyed
  static int unused{};  // The capsule needs something to reference
  pybind11::capsule cleanup(
      &unused, "cleanup", [](PyObject*) {
        OgaShutdown();
      });
  m.add_object("_cleanup", cleanup);

  pybind11::class_<PyGeneratorParams>(m, "GeneratorParams")
      .def(pybind11::init<const OgaModel&>())
      .def("set_search_options", &PyGeneratorParams::SetSearchOptions)  // See config.h 'struct Search' for the options
      .def("set_speculative_options", &PyGeneratorParams::SetSpeculativeOptions)
      .def("get_speculative_options", &PyGeneratorParams::GetSpeculativeOptions)
      .def("set_guidance", &PyGeneratorParams::SetGuidance,
           pybind11::arg("type"), pybind11::arg("data"),
           pybind11::arg("enable_ff_tokens") = false)
      .def("get_search_options", &PyGeneratorParams::GetSearchOptions);

  pybind11::class_<OgaTokenizerStream>(m, "TokenizerStream")
      .def("decode", [](OgaTokenizerStream& t, int32_t token) { return t.Decode(token); });

  pybind11::class_<OgaNamedTensors>(m, "NamedTensors")
      .def(pybind11::init([]() { return OgaNamedTensors::Create(); }))
      .def("__getitem__", [](OgaNamedTensors& named_tensors, const std::string& name) {
        auto tensor = named_tensors.Get(name.c_str());
        if (!tensor)
          throw std::runtime_error("Tensor with name: " + name + " not found.");
        return tensor;
      })
      .def("__setitem__", [](OgaNamedTensors& named_tensors, const std::string& name, pybind11::array& value) {
        named_tensors.Set(name.c_str(), *ToOgaTensor(value));
      })
      .def("__setitem__", [](OgaNamedTensors& named_tensors, const std::string& name, OgaTensor& value) {
        named_tensors.Set(name.c_str(), value);
      })
      .def("__contains__", [](const OgaNamedTensors& named_tensors, const std::string& name) {
        return const_cast<OgaNamedTensors&>(named_tensors).Get(name.c_str()) != nullptr;
      })
      .def("__delitem__", [](OgaNamedTensors& named_tensors, const std::string& name) {
        named_tensors.Delete(name.c_str());
      })
      .def("__len__", &OgaNamedTensors::Count)
      .def("keys", [](OgaNamedTensors& named_tensors) {
        std::vector<std::string> keys;
        auto names = named_tensors.GetNames();
        for (size_t i = 0; i < names->Count(); i++)
          keys.push_back(names->Get(i));
        return keys;
      });

  pybind11::class_<OgaTensor>(m, "Tensor")
      .def(pybind11::init([](pybind11::array& v) { return ToOgaTensor(v); }))
      .def("shape", &OgaTensor::Shape)
      .def("type", &OgaTensor::Type)
      .def("data", &OgaTensor::Data)
      .def("as_numpy", [](OgaTensor& t) { return ToNumpy(t); });

  pybind11::class_<OgaTokenizer>(m, "Tokenizer")
      .def(pybind11::init([](const OgaModel& model) { return OgaTokenizer::Create(model); }))
      .def_property_readonly("bos_token_id", &OgaTokenizer::GetBosTokenId)
      .def_property_readonly("eos_token_ids", [](const OgaTokenizer& t) {
        return ToPython(t.GetEosTokenIds());
      })
      .def_property_readonly("pad_token_id", &OgaTokenizer::GetPadTokenId)
      .def_property_readonly("bot_token_id", &OgaTokenizer::GetBotTokenId)
      .def_property_readonly("eot_token_id", &OgaTokenizer::GetEotTokenId)
      .def_property_readonly("bor_token_id", &OgaTokenizer::GetBorTokenId)
      .def_property_readonly("eor_token_id", &OgaTokenizer::GetEorTokenId)
      .def("update_options", [](OgaTokenizer& t, pybind11::kwargs kwargs) {
        std::vector<std::string> key_storage;
        std::vector<std::string> value_storage;
        key_storage.reserve(kwargs.size());
        value_storage.reserve(kwargs.size());

        std::vector<const char*> keys;
        std::vector<const char*> values;
        keys.reserve(kwargs.size());
        values.reserve(kwargs.size());

        for (auto& item : kwargs) {
            key_storage.emplace_back(pybind11::str(item.first));
            value_storage.emplace_back(pybind11::str(item.second));
            keys.push_back(key_storage.back().c_str());
            values.push_back(value_storage.back().c_str());
        }

        t.UpdateOptions(keys.data(), values.data(), kwargs.size()); })
      .def("encode", [](const OgaTokenizer& t, std::string s) -> pybind11::array_t<int32_t> {
        auto sequences = OgaSequences::Create();
        t.Encode(s.c_str(), *sequences);
        return ToPython(sequences->Get(0)); })
      .def("to_token_id", &OgaTokenizer::ToTokenId)
      .def("decode", [](const OgaTokenizer& t, pybind11::array_t<int32_t> tokens) -> std::string { return t.Decode(ToSpan(tokens)).p_; })
      .def("apply_chat_template", [](const OgaTokenizer& t, const char* messages, const char* template_str, const char* tools, bool add_generation_prompt) -> std::string { return t.ApplyChatTemplate(template_str, messages, tools, add_generation_prompt).p_; }, pybind11::arg("messages"), pybind11::kw_only(), pybind11::arg("template_str") = nullptr, pybind11::arg("tools") = nullptr, pybind11::arg("add_generation_prompt") = true)
      .def("encode_batch", [](const OgaTokenizer& t, std::vector<std::string> strings) {
        std::vector<const char*> c_strings;
        for (const auto& s : strings)
          c_strings.push_back(s.c_str());
        return t.EncodeBatch(c_strings.data(), c_strings.size()); })
      .def("decode_batch", [](const OgaTokenizer& t, const OgaTensor& tokens) {
        std::vector<std::string> strings;
        auto decoded = t.DecodeBatch(tokens);
        for (size_t i = 0; i < decoded->Count(); i++)
          strings.push_back(decoded->Get(i));
        return strings; })
      .def("create_stream", [](const OgaTokenizer& t) { return OgaTokenizerStream::Create(t); });

  pybind11::class_<OgaConfig>(m, "Config")
      .def(pybind11::init([](const std::string& config_path) { return OgaConfig::Create(config_path.c_str()); }))
      .def_static(
          "from_package_ep",
          [](const std::string& config_path, const std::string& ep) {
            return OgaConfig::CreateFromPackageEp(config_path.c_str(), ep.empty() ? nullptr : ep.c_str());
          },
          pybind11::arg("config_path"), pybind11::arg("ep"),
          "Load an OgaConfig from a model package, selecting the variant whose execution "
          "provider matches `ep`. Pass an empty string to auto-detect when the package "
          "declares a single ep across all variants.")
      .def("append_provider", &OgaConfig::AppendProvider)
      .def("set_provider_option", &OgaConfig::SetProviderOption)
      .def("clear_providers", &OgaConfig::ClearProviders)
      .def("add_model_data", [](OgaConfig& config, const std::string& model_filename, pybind11::object obj) {
        if (pybind11::isinstance<pybind11::bytes>(obj)) {
          const auto model_bytes = obj.cast<pybind11::bytes>();
          char* model_data;
          Py_ssize_t model_length;
          if (PyBytes_AsStringAndSize(model_bytes.ptr(), &model_data, &model_length) != 0) {
            throw std::runtime_error("Failed to extract bytes from the object");
          }

          config.AddModelData(model_filename, model_data, static_cast<size_t>(model_length));
        } else if (pybind11::isinstance<pybind11::buffer>(obj)) {
          pybind11::buffer_info info = obj.cast<pybind11::buffer>().request();
          if (info.format != pybind11::format_descriptor<uint8_t>::format() || info.ndim != 1) {
            throw std::runtime_error("Expected a 1D buffer of uint8_t");
          }
          const std::byte* model_data = static_cast<std::byte*>(info.ptr);
          config.AddModelData(model_filename, model_data, info.size);
        } else {
          throw std::runtime_error("Unsupported input type. Expected bytes or buffer.");
        }
      })
      .def("remove_model_data", [](OgaConfig& config, const std::string& model_filename) {
        config.RemoveModelData(model_filename.c_str());
      })
      .def("overlay", &OgaConfig::Overlay)
      .def("set_decoder_provider_options_hardware_device_type", &OgaConfig::SetDecoderProviderOptionsHardwareDeviceType)
      .def("set_decoder_provider_options_hardware_device_id", &OgaConfig::SetDecoderProviderOptionsHardwareDeviceId)
      .def("set_decoder_provider_options_hardware_vendor_id", &OgaConfig::SetDecoderProviderOptionsHardwareVendorId)
      .def("clear_decoder_provider_options_hardware_device_type", &OgaConfig::ClearDecoderProviderOptionsHardwareDeviceType)
      .def("clear_decoder_provider_options_hardware_device_id", &OgaConfig::ClearDecoderProviderOptionsHardwareDeviceId)
      .def("clear_decoder_provider_options_hardware_vendor_id", &OgaConfig::ClearDecoderProviderOptionsHardwareVendorId);

  pybind11::class_<OgaModel>(m, "Model")
      .def(pybind11::init([](const OgaConfig& config) { return OgaModel::Create(config); }))
      .def(pybind11::init([](const std::string& config_path) { return OgaModel::Create(config_path.c_str()); }))
      .def_property_readonly("type", [](const OgaModel& model) -> std::string { return model.GetType().p_; })
      .def_property_readonly(
          "device_type", [](const OgaModel& model) -> std::string { return model.GetDeviceType().p_; }, "The device type the model is running on")
      .def("create_multimodal_processor", [](const OgaModel& model) { return OgaMultiModalProcessor::Create(model); })
      .def("create_streaming_processor", [](OgaModel& model) { return OgaStreamingProcessor::Create(model); }, "Create a StreamingProcessor for mel spectrogram extraction from raw audio.");

  pybind11::class_<PyGenerator>(m, "Generator")
      .def(pybind11::init<const OgaModel&, PyGeneratorParams&>())
      .def("is_done", &PyGenerator::IsDone)
      .def("get_input", &PyGenerator::GetInput)
      .def("get_output", &PyGenerator::GetOutput)
      .def("set_inputs", &PyGenerator::SetInputs)
      .def("set_model_input", &PyGenerator::SetModelInput)
      .def("append_tokens", pybind11::overload_cast<pybind11::array_t<int32_t>&>(&PyGenerator::AppendTokens))
      .def("append_tokens", pybind11::overload_cast<OgaTensor&>(&PyGenerator::AppendTokens))
      .def("token_count", &PyGenerator::TokenCount)
      .def("get_logits", &PyGenerator::GetLogits)
      .def("set_logits", &PyGenerator::SetLogits)
      .def("generate_next_token", &PyGenerator::GenerateNextToken)
      .def("rewind_to", &PyGenerator::RewindTo)
      .def("get_next_tokens", &PyGenerator::GetNextTokens)
      .def("get_sequence", &PyGenerator::GetSequence)
      .def("set_active_adapter", &PyGenerator::SetActiveAdapter)
      .def("set_runtime_option", &PyGenerator::SetRuntimeOption)
      .def("get_speculative_stats", &PyGenerator::GetSpeculativeStats);

  pybind11::class_<OgaImages>(m, "Images")
      .def_static("open", [](pybind11::args image_paths) {
        std::vector<std::string> image_paths_string;
        std::vector<const char*> image_paths_vector;
        for (auto image_path : image_paths) {
          if (!pybind11::isinstance<pybind11::str>(image_path))
            throw std::runtime_error("Image paths must be strings.");
          image_paths_string.push_back(image_path.cast<std::string>());
          image_paths_vector.push_back(image_paths_string.back().c_str());
        }

        return OgaImages::Load(image_paths_vector);
      })
      .def_static("open_bytes", [](pybind11::args image_datas) {
        std::vector<const void*> image_raw_data(image_datas.size());
        std::vector<size_t> image_sizes(image_datas.size());
        for (size_t i = 0; i < image_datas.size(); ++i) {
          if (!pybind11::isinstance<pybind11::bytes>(image_datas[i]))
            throw std::runtime_error("Image data must be bytes.");
          auto bytes = image_datas[i].cast<pybind11::bytes>();
          pybind11::buffer_info info(pybind11::buffer(bytes).request());
          image_raw_data[i] = reinterpret_cast<void*>(info.ptr);
          image_sizes[i] = info.size;
        }

        return OgaImages::Load(image_raw_data.data(), image_sizes.data(), image_raw_data.size());
      });

  pybind11::class_<OgaAudios>(m, "Audios")
      .def_static("open", [](pybind11::args audio_paths) {
        std::vector<std::string> audio_paths_string;
        std::vector<const char*> audio_paths_vector;

        for (const auto& audio_path : audio_paths) {
          if (!pybind11::isinstance<pybind11::str>(audio_path))
            throw std::runtime_error("Audio paths must be strings.");
          audio_paths_string.push_back(audio_path.cast<std::string>());
          audio_paths_vector.push_back(audio_paths_string.back().c_str());
        }

        return OgaAudios::Load(audio_paths_vector);
      })
      .def_static("open_bytes", [](pybind11::args audio_datas) {
        std::vector<const void*> audio_raw_data(audio_datas.size());
        std::vector<size_t> audio_sizes(audio_datas.size());
        for (size_t i = 0; i < audio_datas.size(); ++i) {
          if (!pybind11::isinstance<pybind11::bytes>(audio_datas[i]))
            throw std::runtime_error("Audio data must be bytes.");
          auto bytes = audio_datas[i].cast<pybind11::bytes>();
          pybind11::buffer_info info(pybind11::buffer(bytes).request());
          audio_raw_data[i] = reinterpret_cast<void*>(info.ptr);
          audio_sizes[i] = info.size;
        }

        return OgaAudios::Load(audio_raw_data.data(), audio_sizes.data(), audio_raw_data.size());
      });

  pybind11::class_<OgaMultiModalProcessor>(m, "MultiModalProcessor")
      .def(
          "__call__", [](OgaMultiModalProcessor& processor, pybind11::object prompts, const pybind11::kwargs& kwargs) {
            OgaImages* images{};
            OgaAudios* audios{};
            if (kwargs.contains("images")) {
              images = kwargs["images"].cast<OgaImages*>();
            }
            if (kwargs.contains("audios")) {
              audios = kwargs["audios"].cast<OgaAudios*>();
            }

            std::vector<std::string> prompts_str;
            std::vector<const char*> c_prompts;
            if (pybind11::isinstance<pybind11::str>(prompts)) {
              // One prompt
              return processor.ProcessImagesAndAudios(prompts.cast<std::string>().c_str(), images, audios);
            } else if (pybind11::isinstance<pybind11::list>(prompts)) {
              // Multiple prompts
              for (const auto& prompt : prompts) {
                if (!pybind11::isinstance<pybind11::str>(prompt)) {
                  throw std::runtime_error("One or more items in the list of provided prompts is not a string.");
                }
                prompts_str.push_back(prompt.cast<std::string>());
                c_prompts.push_back(prompts_str.back().c_str());
              }
            } else if (!prompts.is_none()) {
              // Unsupported type for prompts
              throw std::runtime_error("Unsupported type for prompts. Prompts must be a string or a list of strings.");
            }

            return processor.ProcessImagesAndAudios(c_prompts, images, audios);
          },
          pybind11::arg("prompt") = pybind11::none())
      .def("create_stream", [](OgaMultiModalProcessor& processor) { return OgaTokenizerStream::Create(processor); })
      .def("decode", [](OgaMultiModalProcessor& processor, pybind11::array_t<int32_t> tokens) -> std::string {
        return processor.Decode(ToSpan(tokens)).p_;
      });

  pybind11::class_<OgaAdapters>(m, "Adapters")
      .def(pybind11::init([](OgaModel& model) {
        return OgaAdapters::Create(model);
      }))
      .def("unload", &OgaAdapters::UnloadAdapter)
      .def("load", &OgaAdapters::LoadAdapter);

  pybind11::class_<OgaRequest>(m, "Request")
      .def(pybind11::init(
          [](PyGeneratorParams& params) {
            return OgaRequest::Create(*params.params_);
          }))
      .def("add_tokens", [](OgaRequest& request, pybind11::array_t<int32_t> tokens) {
        auto sequences = OgaSequences::Create();
        auto tokens_span = ToSpan(tokens);
        sequences->Append(tokens_span.data(), tokens_span.size());
        request.AddTokens(*sequences);
      })
      .def("has_unseen_tokens", &OgaRequest::HasUnseenTokens)
      .def("is_done", &OgaRequest::IsDone)
      .def("get_unseen_token", &OgaRequest::GetUnseenToken)
      .def("set_opaque_data", [](OgaRequest& request, pybind11::object opaque_data) {
        request.SetOpaqueData(opaque_data.ptr());
      })
      .def("get_opaque_data", [](OgaRequest& request) -> pybind11::object {
        auto opaque_data = request.GetOpaqueData();
        if (!opaque_data)
          return pybind11::none();
        return pybind11::reinterpret_borrow<pybind11::object>(static_cast<PyObject*>(opaque_data));
      });

  pybind11::class_<OgaEngine>(m, "Engine")
      .def(pybind11::init([](OgaModel& model) { return OgaEngine::Create(model); }))
      .def("add_request", &OgaEngine::Add)
      .def("step", &OgaEngine::Step)
      .def("remove_request", &OgaEngine::Remove)
      .def("has_pending_requests", &OgaEngine::HasPendingRequests);

  pybind11::class_<OgaStreamingProcessor>(m, "StreamingProcessor")
      .def(pybind11::init([](OgaModel& model) { return OgaStreamingProcessor::Create(model); }),
           "Create a StreamingProcessor for mel spectrogram extraction.\n"
           "The model must be of type 'nemotron_speech'.")
      .def(
          "process",
          [](OgaStreamingProcessor& proc, pybind11::array_t<float, pybind11::array::c_style | pybind11::array::forcecast> audio_chunk) -> pybind11::object {
            auto buf = audio_chunk.request();
            if (buf.ndim != 1) {
              throw std::runtime_error("audio_chunk must be a 1-D array, got " + std::to_string(buf.ndim) + "-D");
            }
            auto result = proc.Process(static_cast<const float*>(buf.ptr), static_cast<size_t>(buf.size));
            if (result) {
              return pybind11::cast(std::move(result));
            }
            return pybind11::none();
          },
          pybind11::arg("audio_chunk"),
          "Feed raw PCM audio. Returns a NamedTensors if a full chunk is ready, or None if more audio is needed.")
      .def(
          "flush",
          [](OgaStreamingProcessor& proc) -> pybind11::object {
            auto result = proc.Flush();
            if (result) {
              return pybind11::cast(std::move(result));
            }
            return pybind11::none();
          },
          "Flush remaining buffered audio (pads with silence). Returns NamedTensors or None.")
      .def(
          "set_option",
          [](OgaStreamingProcessor& proc, const std::string& key, const std::string& value) {
            proc.SetOption(key.c_str(), value.c_str());
          },
          pybind11::arg("key"),
          pybind11::arg("value"),
          "Set a processor option. Keys: 'use_vad', 'vad_threshold', 'silence_duration_ms', 'prefix_padding_ms'.")
      .def(
          "get_option",
          [](OgaStreamingProcessor& proc, const std::string& key) {
            return std::string(proc.GetOption(key.c_str()));
          },
          pybind11::arg("key"),
          "Get a processor option value by key.");

  m.def("set_log_options", &SetLogOptions);
  m.def("set_log_callback", &SetLogCallback);
  m.def("enable_telemetry_events", []() { OgaSetTelemetryEnabled(true); }, "Enable non-essential ONNX Runtime GenAI telemetry events.");
  m.def("disable_telemetry_events", []() { OgaSetTelemetryEnabled(false); }, "Disable non-essential ONNX Runtime GenAI telemetry events.");

  m.def("is_cuda_available", []() { return USE_CUDA != 0; });
  m.def("is_dml_available", []() { return USE_DML != 0; });
  m.def("is_webgpu_available", []() { return true; });
  m.def("is_qnn_available", []() { return true; });
  m.def("is_openvino_available", []() { return true; });

  m.def("set_current_gpu_device_id", [](int device_id) { Ort::SetCurrentGpuDeviceId(device_id); });
  m.def("get_current_gpu_device_id", []() { return Ort::GetCurrentGpuDeviceId(); });

  m.def("register_execution_provider_library", [](const std::string& provider_name, const std::string& path_str) {
    OgaRegisterExecutionProviderLibrary(provider_name.c_str(), path_str.c_str());
  });

  m.def("unregister_execution_provider_library", [](const std::string& provider_name) {
    OgaUnregisterExecutionProviderLibrary(provider_name.c_str());
  });
}