// [2] 轻量预取(无 staging，仅预热 page cache) + [3] staging 版 miss→hit + STAGING_HPROF 探针。
#include "prefetch.h"
#include "../io/bg_reader.h"

#include <chrono>
#include <functional>
#include <map>
#include <tuple>
#include <unordered_set>
#include <fcntl.h>
#include <unistd.h>

// ====== [2] 轻量预取(无 staging 模式)：GPU 完成回调里读 inds、pread 预热 page cache ======
// 命门验证:回调在 command buffer 完成后触发,此时 inds 已算完,读到的是正确值。
class PrefetchOnCompletePrimitive : public mx::Primitive {
 public:
  PrefetchOnCompletePrimitive(mx::Stream s, std::string path, size_t stride, bool do_read)
      : Primitive(s), path_(std::move(path)), stride_(stride), do_read_(do_read) {}
  const char* name() const override { return "PrefetchOnCompletePrimitive"; }

  void eval_cpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
    outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes()));
    mx::array ids = inputs[0];
    ids.eval();
    record_and_read(ids.data<uint32_t>(), ids.size(), path_, stride_, do_read_);
  }

  void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
    outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes()));
    mx::array ids = inputs[0];                 // by-value：保活到回调结束
    const uint32_t* ptr = ids.data<uint32_t>();  // 指针此刻有效（buffer 已分配），值待 GPU 算完
    size_t n = ids.size();
    std::string path = path_;
    size_t stride = stride_;
    bool do_read = do_read_;
    auto& enc = mx::metal::get_command_encoder(stream());
    MTL::CommandBuffer* cb = enc.get_command_buffer();
    cb->addCompletedHandler([ids, ptr, n, path, stride, do_read](MTL::CommandBuffer*) {
      // 此时 buffer 已完成 → ptr 指向已算好的 inds 值。ids 按值捕获保证 buffer 不被释放。
      record_and_read(ptr, n, path, stride, do_read);
    });
  }

 private:
  static void record_and_read(const uint32_t* p, size_t n, const std::string& path,
                              size_t stride, bool do_read) {
    int fd = (do_read && !path.empty()) ? ::open(path.c_str(), O_RDONLY) : -1;
    static thread_local std::vector<uint8_t> buf;
    if (do_read && buf.size() < stride) buf.resize(stride);
    for (size_t i = 0; i < n; ++i) {
      int e = static_cast<int>(p[i]);
      if (fd >= 0) ::pread(fd, buf.data(), stride, static_cast<off_t>(static_cast<size_t>(e) * stride));
    }
    if (fd >= 0) ::close(fd);
  }
  std::string path_;
  size_t stride_;
  bool do_read_;
};

// 返回一个 dummy 输出（挂进图里，eval 时触发上面的回调）。
mx::array prefetch_on_complete(
    const mx::array& expert_ids,
    const std::string& path,
    int stride,
    bool do_read = true,
    mx::StreamOrDevice s = {}) {
  return mx::array(
      mx::Shape{1},
      mx::uint8,
      std::make_shared<PrefetchOnCompletePrimitive>(
          mx::to_stream(s), path, static_cast<size_t>(stride), do_read),
      std::vector<mx::array>{expert_ids});
}

// ---- handler 触发时刻探针(STAGING_HPROF)：记录每次 staging 完成回调被 Metal 触发的时刻 ----
// 用于实测"完成回调到底扎堆在 eval 尾、还是逐层铺开"。仅诊断用，默认关。
static std::mutex g_hprof_mutex;
static std::vector<std::tuple<long, int, double>> g_hprof_log;  // (gen, layer, t_fire_seconds)
static bool g_hprof_on = false;

static inline double hprof_steady_now() {
  return std::chrono::duration<double>(
             std::chrono::steady_clock::now().time_since_epoch())
      .count();
}

// ====== [3] staging 版 miss→hit：handler pread 进 per-layer staging + 记录 (expert→row) ======
static std::mutex g_stg_mutex;
// layer -> (gen, [(expert, row)])；handler 原子写 gen+映射，主线程按 gen 匹配 buffer 后 take。
static std::map<int, std::pair<long, std::vector<std::pair<int, int>>>> g_stg_ready;

class PrefetchStagingPrimitive : public mx::Primitive {
 public:
  // 方案B：expert_ids 是"预测宽集合"(top-N，按门控分降序)；resident 是目标层当前常驻专家。
  // handler 在回调里过滤掉常驻、按降序取前 cap 个缺口 pread 进 staging（cap=buffer 行数）。
  // 这样预测可以很宽(高 recall)，而 staging 内存只按 cap 预留(小，覆盖缺口分布即可)。
  PrefetchStagingPrimitive(mx::Stream s, int layer, long gen, std::string path, size_t stride,
                           std::vector<int> resident, int cap, bool parallel)
      : Primitive(s), layer_(layer), gen_(gen), path_(std::move(path)), stride_(stride),
        resident_(std::move(resident)), cap_(cap), parallel_(parallel) {}
  const char* name() const override { return "PrefetchStagingPrimitive"; }

  void eval_cpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
    outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes()));
    mx::array ids = inputs[0]; ids.eval();
    mx::array stg = inputs[1]; stg.eval();
    fill(ids.data<uint32_t>(), ids.size(), stg.data<uint8_t>(), layer_, gen_, path_, stride_,
         resident_, cap_);
  }
  void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
    outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes()));
    mx::array ids = inputs[0];
    mx::array stg = inputs[1];
    const uint32_t* idp = ids.data<uint32_t>();
    uint8_t* sp = stg.data<uint8_t>();
    size_t n = ids.size();
    int layer = layer_; long gen = gen_; std::string path = path_; size_t stride = stride_;
    std::vector<int> resident = resident_; int cap = cap_; bool parallel = parallel_;
    auto& enc = mx::metal::get_command_encoder(stream());
    MTL::CommandBuffer* cb = enc.get_command_buffer();
    cb->addCompletedHandler(
        [ids, stg, idp, sp, n, layer, gen, path, stride, resident, cap, parallel](MTL::CommandBuffer*) {
          // 探针：记录本回调被 Metal 触发的时刻(= 该层 pread 能开始跑的时刻)。
          if (g_hprof_on) {
            double t = hprof_steady_now();
            std::lock_guard<std::mutex> lk(g_hprof_mutex);
            g_hprof_log.emplace_back(gen, layer, t);
          }
          // ids/stg 按值捕获保活 buffer；派后台线程时再拷一份保活到 fill 跑完。
          if (parallel) {
            bg_submit_task([ids, stg, idp, sp, n, layer, gen, path, stride, resident, cap]() {
              fill(idp, n, sp, layer, gen, path, stride, resident, cap);
            });
          } else {
            fill(idp, n, sp, layer, gen, path, stride, resident, cap);
          }
        });
  }

 private:
  static void fill(const uint32_t* idp, size_t n, uint8_t* stg,
                   int layer, long gen, const std::string& path, size_t stride,
                   const std::vector<int>& resident, int cap) {
    int fd = ::open(path.c_str(), O_RDONLY);
    if (fd < 0) return;
    std::unordered_set<int> res(resident.begin(), resident.end());
    std::unordered_set<int> done;                 // 去重：同一缺口只 pread 一次
    std::vector<std::pair<int, int>> ready;
    ready.reserve(static_cast<size_t>(cap));
    int row = 0;                                  // 写入的 staging 行（≤ cap）
    for (size_t i = 0; i < n && row < cap; ++i) {
      int e = static_cast<int>(idp[i]);
      if (res.count(e) || !done.insert(e).second) continue;  // 已常驻/已取过 → 跳过
      if (::pread(fd, stg + static_cast<size_t>(row) * stride, stride,
                  static_cast<off_t>(static_cast<size_t>(e) * stride))
          == static_cast<ssize_t>(stride)) {
        ready.emplace_back(e, row);
        ++row;
      }
    }
    ::close(fd);
    std::lock_guard<std::mutex> lk(g_stg_mutex);
    g_stg_ready[layer] = {gen, std::move(ready)};   // 原子：gen 与映射一起写
  }
  int layer_;
  long gen_;
  std::string path_;
  size_t stride_;
  std::vector<int> resident_;   // 目标层提交时刻的常驻专家快照（过滤用）
  int cap_;                     // staging buffer 行数上限
  bool parallel_;               // true: fill 派后台线程池并行；false: 回调线程同步(旧行为)
};

mx::array prefetch_into_staging(
    const mx::array& staging, const mx::array& expert_ids, int layer, long gen,
    const std::string& path, int stride, const std::vector<int>& resident, int cap,
    bool parallel, mx::StreamOrDevice s = {}) {
  return mx::array(
      mx::Shape{1}, mx::uint8,
      std::make_shared<PrefetchStagingPrimitive>(
          mx::to_stream(s), layer, gen, path, static_cast<size_t>(stride), resident, cap, parallel),
      std::vector<mx::array>{expert_ids, staging});
}

// 取走某层就绪记录：[gen, e0,r0,e1,r1,...]（首元素是 generation）；空表示无就绪。
std::vector<long> prefetch_staging_take(int layer) {
  std::lock_guard<std::mutex> lk(g_stg_mutex);
  std::vector<long> out;
  auto it = g_stg_ready.find(layer);
  if (it != g_stg_ready.end()) {
    out.push_back(it->second.first);                 // gen
    for (auto& p : it->second.second) { out.push_back(p.first); out.push_back(p.second); }
    g_stg_ready.erase(it);
  }
  return out;
}

// ---- handler 触发时刻探针接口 ----
void staging_hprof_enable(bool on) {
  std::lock_guard<std::mutex> lk(g_hprof_mutex);
  g_hprof_on = on;
  g_hprof_log.clear();             // 开启即清零，便于每次采集干净
}
double staging_hprof_now() { return hprof_steady_now(); }  // 与日志同一时钟，供 Python 标 eval 边界
// 扁平返回 [gen0,layer0,t0, gen1,layer1,t1, ...]（避开 nanobind tuple caster）。
std::vector<double> staging_hprof_get() {
  std::lock_guard<std::mutex> lk(g_hprof_mutex);
  std::vector<double> out;
  out.reserve(g_hprof_log.size() * 3);
  for (auto& r : g_hprof_log) {
    out.push_back(static_cast<double>(std::get<0>(r)));
    out.push_back(static_cast<double>(std::get<1>(r)));
    out.push_back(std::get<2>(r));
  }
  return out;
}