#include "engine/models/fun_asr_nano/decoder.h"

#include "engine/framework/assets/tensor_source.h"
#include "engine/framework/core/backend.h"
#include "engine/framework/core/backend_weight_store.h"
#include "engine/framework/debug/profiler.h"
#include "engine/framework/modules/activation_modules.h"
#include "engine/framework/modules/linear_module.h"
#include "engine/framework/modules/lookup_modules.h"
#include "engine/framework/modules/norm_modules.h"
#include "engine/framework/modules/positional_modules.h"
#include "engine/framework/modules/primitive_modules.h"
#include "engine/framework/modules/structural_modules.h"
#include "engine/framework/modules/transformers/qwen_causal_decoder.h"
#include "engine/framework/runtime/kv_cache.h"

#include <ggml-backend.h>
#include <ggml.h>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>
#include <vector>

namespace engine::models::fun_asr_nano {
namespace {

namespace modules = engine::modules;
using Clock = std::chrono::steady_clock;

struct GgmlContextDeleter {
  void operator()(ggml_context *ctx) const noexcept {
    if (ctx != nullptr) {
      ggml_free(ctx);
    }
  }
};

struct TextLayerWeights {
  core::TensorValue input_norm;
  core::TensorValue q_proj;
  core::TensorValue k_proj;
  core::TensorValue v_proj;
  core::TensorValue o_proj;
  core::TensorValue q_norm;
  core::TensorValue k_norm;
  core::TensorValue post_norm;
  core::TensorValue gate_proj;
  core::TensorValue up_proj;
  core::TensorValue down_proj;
};

struct TextDecoderWeights {
  std::shared_ptr<core::BackendWeightStore> store;
  core::TensorValue token_embedding;
  std::vector<TextLayerWeights> layers;
  core::TensorValue norm;
  core::TensorValue lm_head;
};

struct PrefillOutput {
  std::vector<float> logits;
  runtime::TransformerKVState kv_state;
};

modules::QwenDecoderLayerWeights
to_qwen_layer_weights(const TextLayerWeights &weights) {
  modules::QwenDecoderLayerWeights out;
  out.input_norm = {weights.input_norm, std::nullopt};
  out.self_attention.q_weight = weights.q_proj;
  out.self_attention.k_weight = weights.k_proj;
  out.self_attention.v_weight = weights.v_proj;
  out.self_attention.out_weight = weights.o_proj;
  out.q_norm = {weights.q_norm, std::nullopt};
  out.k_norm = {weights.k_norm, std::nullopt};
  out.post_norm = {weights.post_norm, std::nullopt};
  out.mlp.gate_proj = {weights.gate_proj, std::nullopt};
  out.mlp.up_proj = {weights.up_proj, std::nullopt};
  out.mlp.down_proj = {weights.down_proj, std::nullopt};
  return out;
}

modules::QwenCausalDecoderConfig
make_qwen_decoder_config(const FunAsrNanoTextConfig &config) {
  modules::QwenCausalDecoderConfig out;
  out.stack.hidden_size = config.hidden_size;
  out.stack.num_attention_heads = config.attention_heads;
  out.stack.num_key_value_heads = config.key_value_heads;
  out.stack.head_dim = config.head_dim;
  out.stack.intermediate_size = config.intermediate_size;
  out.stack.layers = config.layers;
  out.stack.rms_norm_eps = config.rms_norm_eps;
  out.stack.rope_theta = config.rope_theta;
  out.stack.use_qk_norm = true;
  out.stack.runtime.static_cache.update_mode =
      modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows;
  out.logits_size = config.vocab_size;
  out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep;
  return out;
}

modules::QwenCausalDecoderWeights
make_qwen_decoder_weights(const TextDecoderWeights &weights) {
  modules::QwenCausalDecoderWeights out;
  out.stack.layers.reserve(weights.layers.size());
  for (const auto &layer : weights.layers) {
    out.stack.layers.push_back(to_qwen_layer_weights(layer));
  }
  out.final_norm = {weights.norm, std::nullopt};
  out.lm_head = {weights.lm_head, std::nullopt};
  return out;
}

int64_t head_dim(const FunAsrNanoTextConfig &config) {
  if (config.attention_heads <= 0 || config.key_value_heads <= 0 ||
      config.head_dim <= 0) {
    throw std::runtime_error(
        "Fun-ASR-Nano decoder attention config is invalid");
  }
  return config.head_dim;
}

core::TensorValue
prompt_embeddings(core::ModuleBuildContext &ctx,
                  const TextDecoderWeights &weights,
                  const FunAsrNanoTextConfig &config, ggml_tensor *token_ids,
                  ggml_tensor *audio_embeddings, ggml_tensor *audio_positions,
                  int64_t prompt_steps, int64_t audio_tokens) {
  auto ids = core::wrap_tensor(
      token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32);
  auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size})
               .build(ctx, ids, weights.token_embedding);
  if (audio_tokens > 0) {
    auto audio = core::wrap_tensor(
        audio_embeddings,
        core::TensorShape::from_dims({audio_tokens, config.hidden_size}),
        GGML_TYPE_F32);
    auto positions = core::wrap_tensor(
        audio_positions, core::TensorShape::from_dims({audio_tokens}),
        GGML_TYPE_I64);
    x = core::wrap_tensor(
        ggml_set_rows(ctx.ggml, x.tensor, audio.tensor, positions.tensor),
        x.shape, GGML_TYPE_F32);
  }
  return core::reshape_tensor(
      ctx, x,
      core::TensorShape::from_dims({1, prompt_steps, config.hidden_size}));
}

TextDecoderWeights load_weights(const FunAsrNanoAssets &assets,
                                ggml_backend_t backend,
                                core::BackendType backend_type,
                                size_t weight_context_bytes,
                                assets::TensorStorageType storage_type) {
  const auto &config = assets.config.text;
  const auto &source = *assets.model_weights;
  TextDecoderWeights weights;
  weights.store = std::make_shared<core::BackendWeightStore>(
      backend, backend_type, "fun_asr_nano.decoder.weights",
      weight_context_bytes);
  weights.token_embedding = weights.store->load_tensor(
      source, "model.language_model.embed_tokens.weight", storage_type,
      {config.vocab_size, config.hidden_size});
  weights.layers.reserve(static_cast<size_t>(config.layers));
  const int64_t dim = head_dim(config);
  for (int64_t layer = 0; layer < config.layers; ++layer) {
    const std::string prefix =
        "model.language_model.layers." + std::to_string(layer);
    TextLayerWeights w;
    w.input_norm = weights.store->load_f32_tensor(
        source, prefix + ".input_layernorm.weight", {config.hidden_size});
    w.q_proj = weights.store->load_tensor(
        source, prefix + ".self_attn.q_proj.weight", storage_type,
        {config.attention_heads * dim, config.hidden_size});
    w.k_proj = weights.store->load_tensor(
        source, prefix + ".self_attn.k_proj.weight", storage_type,
        {config.key_value_heads * dim, config.hidden_size});
    w.v_proj = weights.store->load_tensor(
        source, prefix + ".self_attn.v_proj.weight", storage_type,
        {config.key_value_heads * dim, config.hidden_size});
    w.o_proj = weights.store->load_tensor(
        source, prefix + ".self_attn.o_proj.weight", storage_type,
        {config.hidden_size, config.attention_heads * dim});
    w.q_norm = weights.store->load_f32_tensor(
        source, prefix + ".self_attn.q_norm.weight", {dim});
    w.k_norm = weights.store->load_f32_tensor(
        source, prefix + ".self_attn.k_norm.weight", {dim});
    w.post_norm = weights.store->load_f32_tensor(
        source, prefix + ".post_attention_layernorm.weight",
        {config.hidden_size});
    w.gate_proj = weights.store->load_tensor(
        source, prefix + ".mlp.gate_proj.weight", storage_type,
        {config.intermediate_size, config.hidden_size});
    w.up_proj = weights.store->load_tensor(
        source, prefix + ".mlp.up_proj.weight", storage_type,
        {config.intermediate_size, config.hidden_size});
    w.down_proj = weights.store->load_tensor(
        source, prefix + ".mlp.down_proj.weight", storage_type,
        {config.hidden_size, config.intermediate_size});
    weights.layers.push_back(std::move(w));
  }
  weights.norm = weights.store->load_f32_tensor(
      source, "model.language_model.norm.weight", {config.hidden_size});
  weights.lm_head = weights.token_embedding;
  weights.store->upload();
  return weights;
}

int32_t argmax_index(const std::vector<float> &values) {
  if (values.empty()) {
    throw std::runtime_error(
        "Fun-ASR-Nano decoder cannot select from empty logits");
  }
  size_t best = 0;
  for (size_t i = 1; i < values.size(); ++i) {
    if (values[i] > values[best]) {
      best = i;
    }
  }
  return static_cast<int32_t>(best);
}

bool is_eos(const FunAsrNanoTextConfig &config, int32_t token) {
  return token == config.eos_token_id;
}

std::shared_ptr<const FunAsrNanoAssets>
require_assets(std::shared_ptr<const FunAsrNanoAssets> assets) {
  if (assets == nullptr) {
    throw std::runtime_error("Fun-ASR-Nano decoder requires assets");
  }
  return assets;
}

class TextDecoderWeightsRuntime {
public:
  TextDecoderWeightsRuntime(std::shared_ptr<const FunAsrNanoAssets> assets,
                            core::ExecutionContext &execution,
                            size_t weight_context_bytes,
                            assets::TensorStorageType storage_type)
      : assets_(require_assets(std::move(assets))),
        backend_(execution.backend()), backend_type_(execution.backend_type()),
        threads_(std::max(1, execution.config().threads)),
        weights_(std::make_shared<TextDecoderWeights>(
            load_weights(*assets_, backend_, backend_type_,
                         weight_context_bytes, storage_type))) {
    if (backend_ == nullptr) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder backend is not initialized");
    }
  }

  const FunAsrNanoAssets &assets() const noexcept { return *assets_; }

  const TextDecoderWeights &weights() const noexcept { return *weights_; }

  ggml_backend_t backend() const noexcept { return backend_; }

  core::BackendType backend_type() const noexcept { return backend_type_; }

  int threads() const noexcept { return threads_; }

private:
  std::shared_ptr<const FunAsrNanoAssets> assets_;
  ggml_backend_t backend_ = nullptr;
  core::BackendType backend_type_ = core::BackendType::Cpu;
  int threads_ = 1;
  std::shared_ptr<const TextDecoderWeights> weights_;
};

class PrefillGraph {
public:
  PrefillGraph(std::shared_ptr<TextDecoderWeightsRuntime> runtime,
               int64_t prompt_steps, int64_t audio_tokens,
               size_t graph_arena_bytes)
      : runtime_(std::move(runtime)), prompt_steps_(prompt_steps),
        audio_tokens_(audio_tokens) {
    if (prompt_steps_ <= 0) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill requires positive prompt length");
    }
    if (audio_tokens_ < 0 || audio_tokens_ > prompt_steps_) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill audio token count is invalid");
    }
    const auto build_start = Clock::now();
    ggml_init_params params{graph_arena_bytes, nullptr, true};
    ctx_.reset(ggml_init(params));
    if (ctx_ == nullptr) {
      throw std::runtime_error(
          "failed to initialize Fun-ASR-Nano decoder prefill graph context");
    }
    const auto &config = runtime_->assets().config.text;
    const auto &weights = runtime_->weights();
    core::ModuleBuildContext ctx{ctx_.get(), "fun_asr_nano.decoder.prefill",
                                 runtime_->backend_type()};
    token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_);
    audio_embeddings_ =
        ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, config.hidden_size,
                           std::max<int64_t>(audio_tokens_, 1));
    audio_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64,
                                          std::max<int64_t>(audio_tokens_, 1));
    auto x =
        prompt_embeddings(ctx, weights, config, token_ids_, audio_embeddings_,
                          audio_positions_, prompt_steps_, audio_tokens_);
    positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_);
    auto positions = core::wrap_tensor(
        positions_, core::TensorShape::from_dims({prompt_steps_}),
        GGML_TYPE_I32);

    auto decoder_out =
        modules::QwenCausalDecoderModule(make_qwen_decoder_config(config))
            .build(ctx, x, positions, make_qwen_decoder_weights(weights));
    for (const auto &layer : decoder_out.state.layers) {
      if (!layer.key.has_value() || !layer.value.has_value()) {
        throw std::runtime_error(
            "Fun-ASR-Nano decoder prefill decoder did not return K/V state");
      }
      auto *key_readback =
          ggml_cpy(ctx_.get(), layer.key->tensor,
                   ggml_dup_tensor(ctx_.get(), layer.key->tensor));
      auto *value_readback =
          ggml_cpy(ctx_.get(), layer.value->tensor,
                   ggml_dup_tensor(ctx_.get(), layer.value->tensor));
      ggml_set_output(key_readback);
      ggml_set_output(value_readback);
      keys_.push_back(key_readback);
      values_.push_back(value_readback);
    }
    logits_ = decoder_out.logits.tensor;
    ggml_set_output(logits_);
    graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
    for (auto *key : keys_) {
      ggml_build_forward_expand(graph_, key);
    }
    for (auto *value : values_) {
      ggml_build_forward_expand(graph_, value);
    }
    ggml_build_forward_expand(graph_, logits_);
    gallocr_ = ggml_gallocr_new(
        ggml_backend_get_default_buffer_type(runtime_->backend()));
    if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) ||
        !ggml_gallocr_alloc_graph(gallocr_, graph_)) {
      throw std::runtime_error(
          "failed to allocate Fun-ASR-Nano decoder prefill graph");
    }
    const auto pos = modules::qwen_position_ids(prompt_steps_);
    ggml_backend_tensor_set(positions_, pos.data(), 0,
                            pos.size() * sizeof(int32_t));
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prefill.graph.build_ms",
        engine::debug::elapsed_ms(build_start, Clock::now()));
    debug::trace_log_scalar("fun_asr_nano.decoder.prefill_prompt_steps",
                            prompt_steps_);
  }

  ~PrefillGraph() {
    engine::core::release_backend_graph_resources(runtime_->backend(), graph_);
    if (gallocr_ != nullptr) {
      ggml_gallocr_free(gallocr_);
    }
  }

  bool matches(const TextDecoderWeightsRuntime &runtime, int64_t prompt_steps,
               int64_t audio_tokens) const {
    return runtime_.get() == &runtime && prompt_steps_ == prompt_steps &&
           audio_tokens_ == audio_tokens;
  }

  PrefillOutput run(const std::vector<int32_t> &token_ids,
                    const std::vector<float> &audio_embeddings,
                    const std::vector<int32_t> &audio_positions) {
    const auto &config = runtime_->assets().config.text;
    if (static_cast<int64_t>(token_ids.size()) != prompt_steps_) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill token id count mismatch");
    }
    if (static_cast<int64_t>(audio_embeddings.size()) !=
        audio_tokens_ * config.hidden_size) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill audio embedding size mismatch");
    }
    if (static_cast<int64_t>(audio_positions.size()) != audio_tokens_) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill audio position count mismatch");
    }
    auto timing_start = Clock::now();
    ggml_backend_tensor_set(token_ids_, token_ids.data(), 0,
                            token_ids.size() * sizeof(int32_t));
    if (audio_tokens_ > 0) {
      std::vector<int64_t> positions(audio_positions.begin(),
                                     audio_positions.end());
      ggml_backend_tensor_set(audio_embeddings_, audio_embeddings.data(), 0,
                              audio_embeddings.size() * sizeof(float));
      ggml_backend_tensor_set(audio_positions_, positions.data(), 0,
                              positions.size() * sizeof(int64_t));
    }
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prefill_input_upload_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    core::set_backend_threads(runtime_->backend(), runtime_->threads());
    timing_start = Clock::now();
    const ggml_status status =
        engine::core::compute_backend_graph(runtime_->backend(), graph_);
    ggml_backend_synchronize(runtime_->backend());
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prefill.graph.compute_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    if (status != GGML_STATUS_SUCCESS) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder prefill graph compute failed");
    }
    PrefillOutput out;
    out.logits.resize(static_cast<size_t>(config.vocab_size));
    timing_start = Clock::now();
    ggml_backend_tensor_get(logits_, out.logits.data(), 0,
                            out.logits.size() * sizeof(float));
    out.kv_state.current_end = prompt_steps_;
    out.kv_state.layers.resize(keys_.size());
    const size_t layer_values = static_cast<size_t>(
        prompt_steps_ * config.key_value_heads * head_dim(config));
    for (size_t layer = 0; layer < keys_.size(); ++layer) {
      auto &state = out.kv_state.layers[layer];
      state.valid_steps = prompt_steps_;
      state.key = core::read_tensor_f32(keys_[layer]);
      state.value = core::read_tensor_f32(values_[layer]);
      if (state.key.size() != layer_values ||
          state.value.size() != layer_values) {
        throw std::runtime_error(
            "Fun-ASR-Nano decoder prefill K/V export size mismatch");
      }
    }
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prefill_output_read_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    return out;
  }

private:
  std::shared_ptr<TextDecoderWeightsRuntime> runtime_;
  int64_t prompt_steps_ = 0;
  int64_t audio_tokens_ = 0;
  std::unique_ptr<ggml_context, GgmlContextDeleter> ctx_;
  ggml_tensor *token_ids_ = nullptr;
  ggml_tensor *audio_embeddings_ = nullptr;
  ggml_tensor *audio_positions_ = nullptr;
  ggml_tensor *positions_ = nullptr;
  ggml_tensor *logits_ = nullptr;
  std::vector<ggml_tensor *> keys_;
  std::vector<ggml_tensor *> values_;
  ggml_cgraph *graph_ = nullptr;
  ggml_gallocr_t gallocr_ = nullptr;
};

class DecodeGraph {
public:
  DecodeGraph(std::shared_ptr<TextDecoderWeightsRuntime> runtime,
              int64_t cache_steps, size_t graph_arena_bytes)
      : runtime_(std::move(runtime)), real_cache_steps_(cache_steps),
        cache_steps_(cache_steps) {
    if (real_cache_steps_ <= 0) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder decode requires positive cache length");
    }
    const auto build_start = Clock::now();
    ggml_init_params params{graph_arena_bytes, nullptr, true};
    ctx_.reset(ggml_init(params));
    if (ctx_ == nullptr) {
      throw std::runtime_error(
          "failed to initialize Fun-ASR-Nano decoder decode graph context");
    }
    const auto &config = runtime_->assets().config.text;
    const auto &weights = runtime_->weights();
    core::ModuleBuildContext ctx{ctx_.get(), "fun_asr_nano.decoder.decode",
                                 runtime_->backend_type()};
    token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1);
    ggml_set_input(token_id_);
    auto token_id = core::wrap_tensor(
        token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32);
    auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size})
                 .build(ctx, token_id, weights.token_embedding);
    x = core::reshape_tensor(
        ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size}));
    positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1);
    ggml_set_input(positions_);
    auto positions = core::wrap_tensor(
        positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32);
    cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1);
    ggml_set_input(cache_slot_);
    auto cache_slot = core::wrap_tensor(
        cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32);
    attention_mask_ =
        ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1);
    ggml_set_input(attention_mask_);
    auto attention_mask = core::wrap_tensor(
        attention_mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}),
        GGML_TYPE_F16);
    graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
    auto decoder_out =
        modules::QwenCausalDecoderModule(make_qwen_decoder_config(config))
            .build_static_cache_tail(ctx, graph_, x, positions,
                                     make_qwen_decoder_weights(weights),
                                     cache_steps_, attention_mask, cache_slot);
    step_cache_ = std::move(decoder_out.cache);
    logits_ = decoder_out.logits.tensor;
    ggml_set_output(logits_);
    ggml_build_forward_expand(graph_, logits_);
    buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend());
    if (buffer_ == nullptr) {
      throw std::runtime_error(
          "failed to allocate Fun-ASR-Nano decoder decode graph");
    }
    attention_mask_values_.assign(static_cast<size_t>(cache_steps_),
                                  ggml_fp32_to_fp16(-INFINITY));
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.decode.graph.build_ms",
        engine::debug::elapsed_ms(build_start, Clock::now()));
    debug::trace_log_scalar("fun_asr_nano.decoder.decode_cache_steps",
                            real_cache_steps_);
  }

  ~DecodeGraph() {
    engine::core::release_backend_graph_resources(runtime_->backend(), graph_);
    if (buffer_ != nullptr) {
      ggml_backend_buffer_free(buffer_);
    }
  }

  bool can_run(const TextDecoderWeightsRuntime &runtime,
               int64_t required_steps) const {
    return runtime_.get() == &runtime && real_cache_steps_ >= required_steps;
  }

  void import_state(const runtime::TransformerKVState &state) {
    step_cache_.import_state(state);
  }

  std::vector<float> run_step(int32_t token) {
    const auto &config = runtime_->assets().config.text;
    if (step_cache_.valid_steps() >= real_cache_steps_) {
      throw std::runtime_error("Fun-ASR-Nano decoder decode cache exhausted");
    }
    ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t));
    const int32_t position = static_cast<int32_t>(step_cache_.current_end());
    ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t));
    const int32_t cache_slot = static_cast<int32_t>(step_cache_.valid_steps());
    ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t));
    modules::write_qwen_cached_step_mask(
        attention_mask_, attention_mask_values_, cache_steps_,
        step_cache_.valid_steps(), step_cache_.valid_steps());
    core::set_backend_threads(runtime_->backend(), runtime_->threads());
    const ggml_status status =
        engine::core::compute_backend_graph(runtime_->backend(), graph_);
    ggml_backend_synchronize(runtime_->backend());
    if (status != GGML_STATUS_SUCCESS) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder decode graph compute failed");
    }
    std::vector<float> logits(static_cast<size_t>(config.vocab_size));
    ggml_backend_tensor_get(logits_, logits.data(), 0,
                            logits.size() * sizeof(float));
    step_cache_.advance_after_direct_append(1);
    return logits;
  }

private:
  std::shared_ptr<TextDecoderWeightsRuntime> runtime_;
  int64_t real_cache_steps_ = 0;
  int64_t cache_steps_ = 0;
  std::unique_ptr<ggml_context, GgmlContextDeleter> ctx_;
  ggml_tensor *token_id_ = nullptr;
  ggml_tensor *positions_ = nullptr;
  ggml_tensor *cache_slot_ = nullptr;
  ggml_tensor *attention_mask_ = nullptr;
  ggml_tensor *logits_ = nullptr;
  std::vector<ggml_fp16_t> attention_mask_values_;
  runtime::TransformerKVCache step_cache_;
  ggml_cgraph *graph_ = nullptr;
  ggml_backend_buffer_t buffer_ = nullptr;
};

} // namespace

struct FunAsrNanoDecoderRuntime::Impl {
  Impl(std::shared_ptr<const FunAsrNanoAssets> assets,
       core::ExecutionContext &execution, size_t prefill_graph_arena_bytes,
       size_t decode_graph_arena_bytes, size_t weight_context_bytes,
       assets::TensorStorageType storage_type)
      : weights(std::make_shared<TextDecoderWeightsRuntime>(
            std::move(assets), execution, weight_context_bytes, storage_type)),
        prefill_graph_arena_bytes(prefill_graph_arena_bytes),
        decode_graph_arena_bytes(decode_graph_arena_bytes) {}

  void validate_prompt_audio(
      const FunAsrNanoPrompt &prompt,
      const FunAsrNanoAdaptorEmbeddings &audio_embeddings) const {
    const auto &config = weights->assets().config.text;
    if (audio_embeddings.hidden_size != config.hidden_size) {
      throw std::runtime_error(
          "Fun-ASR-Nano audio embedding hidden size mismatch");
    }
    if (audio_embeddings.tokens !=
        static_cast<int64_t>(prompt.audio_token_positions.size())) {
      throw std::runtime_error("Fun-ASR-Nano audio embedding token count does "
                               "not match prompt placeholders");
    }
    if (static_cast<int64_t>(audio_embeddings.values.size()) !=
        audio_embeddings.tokens * config.hidden_size) {
      throw std::runtime_error(
          "Fun-ASR-Nano audio embedding value count mismatch");
    }
    std::vector<int32_t> expected_audio_positions;
    for (size_t index = 0; index < prompt.input_ids.size(); ++index) {
      if (prompt.input_ids[index] == weights->assets().config.audio_token_id) {
        expected_audio_positions.push_back(static_cast<int32_t>(index));
      }
    }
    if (prompt.audio_token_positions != expected_audio_positions) {
      throw std::runtime_error(
          "Fun-ASR-Nano audio positions must cover every audio placeholder in "
          "prompt order");
    }
    if (prompt.attention_mask.size() != prompt.input_ids.size() ||
        std::any_of(prompt.attention_mask.begin(), prompt.attention_mask.end(),
                    [](int32_t value) { return value != 1; })) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder requires an unpadded prompt mask");
    }
  }

  FunAsrNanoGeneratedTokens
  generate(const FunAsrNanoPrompt &prompt,
           const FunAsrNanoAdaptorEmbeddings &audio_embeddings,
           const FunAsrNanoGenerationOptions &options,
           const FunAsrNanoTokenCallback &token_callback) {
    const auto &config = weights->assets().config.text;
    if (prompt.input_ids.empty()) {
      throw std::runtime_error("Fun-ASR-Nano decoder prompt is empty");
    }
    if (options.max_new_tokens <= 0) {
      throw std::runtime_error("Fun-ASR-Nano max_new_tokens must be positive");
    }
    const int64_t prompt_steps = static_cast<int64_t>(prompt.input_ids.size());
    if (prompt_steps > config.max_position_embeddings ||
        options.max_new_tokens >
            config.max_position_embeddings - prompt_steps) {
      throw std::runtime_error(
          "Fun-ASR-Nano decoder request exceeds max_position_embeddings");
    }
    auto timing_start = Clock::now();
    validate_prompt_audio(prompt, audio_embeddings);
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prompt_prepare_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    if (prefill_graph == nullptr ||
        !prefill_graph->matches(*weights, prompt_steps,
                                audio_embeddings.tokens)) {
      prefill_graph = std::make_unique<PrefillGraph>(weights, prompt_steps,
                                                     audio_embeddings.tokens,
                                                     prefill_graph_arena_bytes);
    } else {
      debug::timing_log_scalar("fun_asr_nano.decoder.prefill.graph.build_ms",
                               0.0);
      debug::trace_log_scalar("fun_asr_nano.decoder.prefill_prompt_steps",
                              prompt_steps);
    }
    timing_start = Clock::now();
    auto prefill = prefill_graph->run(prompt.input_ids, audio_embeddings.values,
                                      prompt.audio_token_positions);
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.prefill_total_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    const int64_t required_cache_steps =
        prompt_steps + std::max<int64_t>(options.max_new_tokens - 1, 0);
    if (required_cache_steps > prompt_steps) {
      if (decode_graph == nullptr ||
          !decode_graph->can_run(*weights, required_cache_steps)) {
        decode_graph = std::make_unique<DecodeGraph>(
            weights, required_cache_steps, decode_graph_arena_bytes);
      } else {
        debug::timing_log_scalar("fun_asr_nano.decoder.decode.graph.build_ms",
                                 0.0);
        debug::trace_log_scalar("fun_asr_nano.decoder.decode_cache_steps",
                                required_cache_steps);
      }
      decode_graph->import_state(prefill.kv_state);
    }

    FunAsrNanoGeneratedTokens out;
    std::vector<float> logits = std::move(prefill.logits);
    timing_start = Clock::now();
    for (int64_t step = 0; step < options.max_new_tokens; ++step) {
      if (options.capture_logits) {
        out.step_logits.push_back(logits);
      }
      const int32_t token = argmax_index(logits);
      if (is_eos(config, token)) {
        break;
      }
      out.token_ids.push_back(token);
      if (token_callback) {
        token_callback(out);
      }
      if (step + 1 >= options.max_new_tokens) {
        break;
      }
      if (decode_graph == nullptr) {
        throw std::runtime_error(
            "Fun-ASR-Nano decoder decode graph is not initialized");
      }
      logits = decode_graph->run_step(token);
    }
    debug::timing_log_scalar(
        "fun_asr_nano.decoder.decode_total_ms",
        engine::debug::elapsed_ms(timing_start, Clock::now()));
    return out;
  }

  std::shared_ptr<TextDecoderWeightsRuntime> weights;
  size_t prefill_graph_arena_bytes = 0;
  size_t decode_graph_arena_bytes = 0;
  std::unique_ptr<PrefillGraph> prefill_graph;
  std::unique_ptr<DecodeGraph> decode_graph;
};

FunAsrNanoDecoderRuntime::FunAsrNanoDecoderRuntime(
    std::shared_ptr<const FunAsrNanoAssets> assets,
    core::ExecutionContext &execution, size_t prefill_graph_arena_bytes,
    size_t decode_graph_arena_bytes, size_t weight_context_bytes,
    assets::TensorStorageType weight_storage_type)
    : impl_(std::make_unique<Impl>(std::move(assets), execution,
                                   prefill_graph_arena_bytes,
                                   decode_graph_arena_bytes,
                                   weight_context_bytes, weight_storage_type)) {
}

FunAsrNanoDecoderRuntime::~FunAsrNanoDecoderRuntime() = default;

FunAsrNanoGeneratedTokens FunAsrNanoDecoderRuntime::generate(
    const FunAsrNanoPrompt &prompt,
    const FunAsrNanoAdaptorEmbeddings &audio_embeddings,
    const FunAsrNanoGenerationOptions &options,
    const FunAsrNanoTokenCallback &token_callback) {
  return impl_->generate(prompt, audio_embeddings, options, token_callback);
}

} // namespace engine::models::fun_asr_nano