/*
 * Kazeia ExecuTorch+QNN LLM JNI bridge.
 *
 * Wraps example::Runner<T> (the qnn_llama_runner CLI logic) into a JNI-callable
 * shared library `libkazeia_pte.so` so the Kazeia Android app can load and run
 * dense .pte LLMs (Qwen3-4B/8B, Qwen2.5-7B, Qwen3Guard-4B) on the Hexagon V79
 * NPU from Kotlin.
 *
 * Java class:  com.kazeia.llm.ExecuTorchJni   (library name: kazeia_pte)
 *
 *   long   load(String ptePath, String tokenizerPath, int seqLen, int evalMode)
 *   long[] generateFromIds(long handle, long[] promptIds, int maxNewTokens)
 *   long[] lastStats(long handle)   // [prefillMs, decodeMs, nGenerated]
 *   void   free(long handle)
 *
 * Also exposes plain extern "C" entrypoints (kz_pte_*) for an on-device test
 * harness that drives the same code path without a JVM.
 */

#include <jni.h>

#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include <vector>

#include <executorch/backends/qualcomm/runtime/QnnExecuTorch.h>
#include <executorch/examples/qualcomm/oss_scripts/llama/runner/runner.h>
#include <executorch/extension/module/module.h>
#include <executorch/runtime/platform/log.h>

using executorch::extension::Module;
using executorch::runtime::Error;

namespace {

// Type-erased wrapper so the opaque handle does not depend on the kv-bitwidth
// template parameter of example::Runner<T>.
struct IRunnerHolder {
  virtual ~IRunnerHolder() = default;
  virtual Error load() = 0;
  virtual Error generate_from_ids(
      const std::vector<uint64_t>& prompt_ids,
      int32_t max_new_tokens,
      std::vector<uint64_t>& out_ids) = 0;
  virtual int64_t last_prefill_ms() const = 0;
  virtual int64_t last_decode_ms() const = 0;
  virtual int64_t last_num_generated() const = 0;
};

template <typename T>
struct RunnerHolder : public IRunnerHolder {
  example::Runner<T> runner;
  RunnerHolder(
      std::unique_ptr<Module> module,
      const std::string& decoder_model_version,
      const std::string& model_path,
      const std::string& tokenizer_path,
      int eval_mode,
      bool shared_buffer)
      : runner(
            std::move(module),
            decoder_model_version,
            model_path,
            tokenizer_path,
            /*dump_logits_path=*/"",
            /*performance_output_path=*/"/data/local/tmp/pte_jni_test/perf.txt",
            /*temperature=*/0.0f,
            eval_mode,
            shared_buffer,
            /*ngram=*/0,
            /*window=*/0,
            /*gcap=*/0,
            /*tokenizer=*/nullptr,
            /*attention_sink_rope_module=*/nullptr) {}

  Error load() override {
    return runner.load();
  }
  Error generate_from_ids(
      const std::vector<uint64_t>& prompt_ids,
      int32_t max_new_tokens,
      std::vector<uint64_t>& out_ids) override {
    return runner.generate_from_ids(prompt_ids, max_new_tokens, out_ids);
  }
  int64_t last_prefill_ms() const override {
    return runner.last_prefill_ms();
  }
  int64_t last_decode_ms() const override {
    return runner.last_decode_ms();
  }
  int64_t last_num_generated() const override {
    return runner.last_num_generated();
  }
};

// Infer the decoder model version string the Runner constructor expects.
// All current Kazeia targets are Qwen; default to "qwen3" (EOS <|im_end|>),
// switch to "qwen2_5" if the path clearly names a 2.5 checkpoint.
std::string infer_decoder_model_version(const std::string& pte_path) {
  auto contains = [&](const char* needle) {
    return pte_path.find(needle) != std::string::npos;
  };
  if (contains("qwen2_5") || contains("qwen2.5") || contains("q25") ||
      contains("Qwen2.5") || contains("qwen25")) {
    return "qwen2_5";
  }
  return "qwen3";
}

// Build a runner holder (load module, pick kv bitwidth, run load()).
// Returns nullptr on failure.
IRunnerHolder* create_runner(
    const std::string& pte_path,
    const std::string& tokenizer_path,
    int /*seq_len*/,
    int eval_mode,
    bool shared_buffer) {
  std::unique_ptr<Module> module = std::make_unique<Module>(
      pte_path.c_str(), Module::LoadMode::MmapUseMlockIgnoreErrors);

  // Older models default to 8-bit KV IO; newer ones publish a width method.
  example::KvBitWidth kv_bitwidth = example::KvBitWidth::kWidth8;
  if (module->method_names().ok() &&
      module->method_names()->count("get_kv_io_bit_width") > 0) {
    auto res = module->get("get_kv_io_bit_width");
    if (res.ok()) {
      kv_bitwidth = static_cast<example::KvBitWidth>(
          res.get().toScalar().to<int64_t>());
    }
  }

  std::string dmv = infer_decoder_model_version(pte_path);

  IRunnerHolder* holder = nullptr;
  if (kv_bitwidth == example::KvBitWidth::kWidth16) {
    holder = new RunnerHolder<uint16_t>(
        std::move(module), dmv, pte_path, tokenizer_path, eval_mode,
        shared_buffer);
  } else {
    holder = new RunnerHolder<uint8_t>(
        std::move(module), dmv, pte_path, tokenizer_path, eval_mode,
        shared_buffer);
  }

  Error err = holder->load();
  if (err != Error::Ok) {
    ET_LOG(Error, "kazeia_pte: runner.load() failed (%d)", (int)err);
    delete holder;
    return nullptr;
  }
  return holder;
}

} // namespace

// ----------------------------------------------------------------------------
// Plain C entrypoints (for the on-device test harness; no JVM required).
// ----------------------------------------------------------------------------
extern "C" {

void* kz_pte_load(
    const char* pte_path,
    const char* tokenizer_path,
    int seq_len,
    int eval_mode) {
  // shared_buffer=true mirrors the validated qnn_llama_runner invocation.
  return create_runner(pte_path, tokenizer_path, seq_len, eval_mode, true);
}

// Writes generated ids into out (caller-allocated, capacity out_cap).
// Returns the number of generated ids (>=0), or -1 on error. If the count
// exceeds out_cap, only out_cap ids are written but the true count is returned.
int kz_pte_generate_from_ids(
    void* handle,
    const uint64_t* prompt_ids,
    int n_prompt,
    int max_new_tokens,
    uint64_t* out,
    int out_cap) {
  if (handle == nullptr || prompt_ids == nullptr || n_prompt <= 0) {
    return -1;
  }
  auto* holder = reinterpret_cast<IRunnerHolder*>(handle);
  std::vector<uint64_t> in(prompt_ids, prompt_ids + n_prompt);
  std::vector<uint64_t> gen;
  Error err = holder->generate_from_ids(in, max_new_tokens, gen);
  if (err != Error::Ok) {
    return -1;
  }
  int n = (int)gen.size();
  int n_copy = (n < out_cap) ? n : out_cap;
  if (out != nullptr && n_copy > 0) {
    std::memcpy(out, gen.data(), n_copy * sizeof(uint64_t));
  }
  return n;
}

void kz_pte_last_stats(void* handle, int64_t out[3]) {
  if (handle == nullptr) {
    out[0] = out[1] = out[2] = 0;
    return;
  }
  auto* holder = reinterpret_cast<IRunnerHolder*>(handle);
  out[0] = holder->last_prefill_ms();
  out[1] = holder->last_decode_ms();
  out[2] = holder->last_num_generated();
}

void kz_pte_free(void* handle) {
  if (handle != nullptr) {
    delete reinterpret_cast<IRunnerHolder*>(handle);
  }
}

} // extern "C"

// ----------------------------------------------------------------------------
// JNI entrypoints: com.kazeia.llm.ExecuTorchJni
// ----------------------------------------------------------------------------
extern "C" {

JNIEXPORT jlong JNICALL
Java_com_kazeia_llm_ExecuTorchJni_load(
    JNIEnv* env,
    jclass /*clazz*/,
    jstring ptePath,
    jstring tokenizerPath,
    jint seqLen,
    jint evalMode) {
  const char* pte = env->GetStringUTFChars(ptePath, nullptr);
  const char* tok = env->GetStringUTFChars(tokenizerPath, nullptr);
  std::string pte_s = pte ? pte : "";
  std::string tok_s = tok ? tok : "";
  if (pte) env->ReleaseStringUTFChars(ptePath, pte);
  if (tok) env->ReleaseStringUTFChars(tokenizerPath, tok);

  IRunnerHolder* holder =
      create_runner(pte_s, tok_s, (int)seqLen, (int)evalMode, /*shared=*/true);
  return reinterpret_cast<jlong>(holder);
}

JNIEXPORT jlongArray JNICALL
Java_com_kazeia_llm_ExecuTorchJni_generateFromIds(
    JNIEnv* env,
    jclass /*clazz*/,
    jlong handle,
    jlongArray promptIds,
    jint maxNewTokens) {
  if (handle == 0 || promptIds == nullptr) {
    return env->NewLongArray(0);
  }
  auto* holder = reinterpret_cast<IRunnerHolder*>(handle);

  jsize n = env->GetArrayLength(promptIds);
  std::vector<uint64_t> in(n);
  {
    jlong* p = env->GetLongArrayElements(promptIds, nullptr);
    for (jsize i = 0; i < n; ++i) {
      in[i] = static_cast<uint64_t>(p[i]);
    }
    env->ReleaseLongArrayElements(promptIds, p, JNI_ABORT);
  }

  std::vector<uint64_t> gen;
  Error err = holder->generate_from_ids(in, (int)maxNewTokens, gen);
  if (err != Error::Ok) {
    return env->NewLongArray(0);
  }

  jlongArray result = env->NewLongArray((jsize)gen.size());
  if (result == nullptr) {
    return env->NewLongArray(0);
  }
  if (!gen.empty()) {
    std::vector<jlong> tmp(gen.size());
    for (size_t i = 0; i < gen.size(); ++i) {
      tmp[i] = static_cast<jlong>(gen[i]);
    }
    env->SetLongArrayRegion(result, 0, (jsize)tmp.size(), tmp.data());
  }
  return result;
}

JNIEXPORT jlongArray JNICALL
Java_com_kazeia_llm_ExecuTorchJni_lastStats(
    JNIEnv* env,
    jclass /*clazz*/,
    jlong handle) {
  jlongArray result = env->NewLongArray(3);
  if (result == nullptr) {
    return nullptr;
  }
  jlong vals[3] = {0, 0, 0};
  if (handle != 0) {
    auto* holder = reinterpret_cast<IRunnerHolder*>(handle);
    vals[0] = holder->last_prefill_ms();
    vals[1] = holder->last_decode_ms();
    vals[2] = holder->last_num_generated();
  }
  env->SetLongArrayRegion(result, 0, 3, vals);
  return result;
}

JNIEXPORT void JNICALL
Java_com_kazeia_llm_ExecuTorchJni_free(
    JNIEnv* /*env*/,
    jclass /*clazz*/,
    jlong handle) {
  if (handle != 0) {
    delete reinterpret_cast<IRunnerHolder*>(handle);
  }
}

} // extern "C"