#include "compute/fused_moe.h"
#include "io/blob_load.h"
#include "io/bg_reader.h"
#include "prefetch/prefetch.h"
#include "pool/owned_pool.h"
#include "pool/side_region.h"
#include "pool/demand.h"
using namespace nb::literals;

NB_MODULE(native_moe_ext, m) {
  nb::module_::import_("mlx.core");
  m.doc() = "Native MLX MoE extension.";

  // ---- 融合 MoE 计算核（native_fused.cpp）----
  m.def(
      "fused_moe",
      &fused_moe,
      "x"_a,
      "expert_ids"_a,
      "scores"_a,
      "compute_dir"_a,
      "layer"_a,
      "hidden"_a,
      "inter"_a,
      "group"_a,
      "bits"_a,
      "num_experts"_a,
      "synthetic"_a,
      nb::kw_only(),
      "stream"_a = nb::none());
  m.def(
      "fused_moe_staged",
      &fused_moe_staged,
      "x"_a,
      "scores"_a,
      "gate_w"_a,
      "gate_s"_a,
      "gate_b"_a,
      "up_w"_a,
      "up_s"_a,
      "up_b"_a,
      "down_w"_a,
      "down_s"_a,
      "down_b"_a,
      "hidden"_a,
      "inter"_a,
      "group"_a,
      "bits"_a,
      nb::kw_only(),
      "stream"_a = nb::none());
  // ---- [1] blob 直读 ----
  m.def(
      "blob_load",
      &blob_load,
      "path"_a,
      "expert_ids"_a,
      "stride"_a,
      nb::kw_only(),
      "stream"_a = nb::none());
  // ---- [2] 轻量预取(无 staging，仅预热 page cache) ----
  m.def(
      "prefetch_on_complete",
      &prefetch_on_complete,
      "expert_ids"_a,
      "path"_a,
      "stride"_a,
      "do_read"_a = true,
      nb::kw_only(),
      "stream"_a = nb::none());
  // ---- [3] staging miss→hit + 完成回调时刻探针（诊断）----
  m.def(
      "prefetch_into_staging",
      &prefetch_into_staging,
      "staging"_a,
      "expert_ids"_a,
      "layer"_a,
      "gen"_a,
      "path"_a,
      "stride"_a,
      "resident"_a,
      "cap"_a,
      "parallel"_a,
      nb::kw_only(),
      "stream"_a = nb::none());
  m.def("prefetch_staging_take", &prefetch_staging_take, "layer"_a);
  m.def("staging_hprof_enable", &staging_hprof_enable, "on"_a);
  m.def("staging_hprof_now", &staging_hprof_now);
  m.def("staging_hprof_get", &staging_hprof_get);
  // ---- [4] 段散写侧区缓存（zero-copy dual-source 默认路径）----
  m.def("prefetch_pool_sideregion", &prefetch_pool_sideregion,
        "pool_list"_a, "seg_nbytes"_a, "expert_ids"_a, "layer"_a, "path"_a, "stride"_a,
        "resident"_a, "spec_slots"_a, "base_row"_a, "gen"_a = 0, nb::kw_only(),
        "stream"_a = nb::none());
  m.def("sideregion_contents", &sideregion_contents, "layer"_a, "gen"_a = 0);
  m.def("sideregion_kv", &sideregion_kv, "layer"_a, "gen"_a = 0);
  m.def("sideregion_reset", &sideregion_reset);
  m.def("sideregion_drain", &sideregion_drain,
        nb::call_guard<nb::gil_scoped_release>());   // 等待时释放 GIL
  // ---- [5] Route 3 owned 池底座（C++ 拥有 buffer + 直写，删 MLX scatter）----
  m.def("pool_owned_zeros", &pool_owned_zeros, "shape"_a, "dtype"_a);
  m.def("pool_write_rows", &pool_write_rows, "pool_list"_a, "srcs_flat"_a, "slots"_a);
  m.def("pool_write_stacked", &pool_write_stacked, "pool_list"_a, "stacked_list"_a, "slots"_a);
  m.def("array_data_ptr", &array_data_ptr, "a"_a);
  // ---- [6] 方案B 真实区槽状态 C++ 接管 + demand_dual ----
  m.def("real_init", &real_init, "layer"_a, "cap"_a);
  m.def("real_region_contents", &real_region_contents, "layer"_a);
  m.def("real_region_count", &real_region_count, "layer"_a);
  m.def("real_reset", &real_reset);
  m.def("real_pin", &real_pin, "layer"_a, "experts"_a);
  m.def("real_pinned_count", &real_pinned_count, "layer"_a);
  m.def("real_freq_dump", &real_freq_dump);
  m.def("demand_dual", &demand_dual, "inds"_a, "pool_list"_a, "seg_nbytes"_a, "layer"_a,
        "side_gen"_a, "path"_a, "stride"_a, "cap"_a, "lfu"_a, "decay_interval"_a, nb::kw_only(),
        "stream"_a = nb::none());
  m.def("demand_last_stats", &demand_last_stats);
  m.def("demand_timings", &demand_timings);
  m.def("demand_timing_enable", &demand_timing_enable, "on"_a);
  m.def("real_debug_place", &real_debug_place, "layer"_a, "experts_flat"_a, "cap"_a, "lfu"_a,
        "decay_interval"_a);
  // ---- [7] 自由后台读线程 ----
  m.def("bg_reader_start", &bg_reader_start, "workers"_a = 1, "low_cap"_a = 0);
  m.def("bg_reader_submit", &bg_reader_submit,
        "dst"_a, "experts"_a, "rows"_a, "path"_a, "stride"_a, "ticket"_a, "prio"_a = 0);
  m.def("bg_reader_ready", &bg_reader_ready, "ticket"_a);
  m.def("bg_reader_wait", &bg_reader_wait, "ticket"_a,
        nb::call_guard<nb::gil_scoped_release>());   // 阻塞等时释放 GIL
  m.def("bg_reader_stop", &bg_reader_stop);
  m.def("bg_pread_into_pool", &bg_pread_into_pool,
        "dst"_a, "seg_off"_a, "seg_nb"_a, "slot"_a, "expert"_a,
        "path"_a, "stride"_a, "ticket"_a, "prio"_a = 0, "nocache"_a = true);
  // ---- 融合 MoE 计算核（slots 版，native_fused.cpp）----
  m.def(
      "fused_moe_slots",
      &fused_moe_slots,
      "x"_a,
      "local_slots"_a,
      "scores"_a,
      "gate_w"_a,
      "gate_s"_a,
      "gate_b"_a,
      "up_w"_a,
      "up_s"_a,
      "up_b"_a,
      "down_w"_a,
      "down_s"_a,
      "down_b"_a,
      "hidden"_a,
      "inter"_a,
      "group"_a,
      "bits"_a,
      nb::kw_only(),
      "stream"_a = nb::none());
}