// [7] 自由后台读线程（de-risk）：双队列 + 低优并发上限的后台 pread 线程池。
#include "bg_reader.h"
#include "blob_io.h"

#include <condition_variable>
#include <functional>
#include <queue>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <fcntl.h>
#include <unistd.h>

namespace {
struct ReadOp { uint8_t* dst; size_t nbytes; off_t file_off; };
struct BgJob {
  std::vector<ReadOp> ops;
  std::string path;
  long ticket;
  std::vector<mx::array> keep;     // 持 buffer 引用，保证 dst 指针在读期间存活
  std::function<void()> task;      // 若设置，worker 直接执行它（侧区异步读用），不走 ops/ticket
  int prio = 0;                    // >0=高优(route 读)，=0=低优(投机兜底)
  bool nocache = true;             // demand（route）读设 false 走 page cache：段偏移非页对齐，
                                   // F_NOCACHE 下多 MB 非对齐 pread 会短读 → 池槽脏字节（实测）。
};

// 双队列 + 低优并发上限：高优(route)读永远优先取、且独占大多数 worker；
// 低优(投机)读最多占 low_cap_ 个 worker → 给 route 读留出 worker 和 SSD 带宽，
// 实现"高分抢带宽先就绪、低分延后不阻塞关键路径"的真 IO 优先级。
class BgReader {
 public:
  ~BgReader() { stop(); }            // 进程退出兜底：避免 joinable 线程析构触发 std::terminate
  void start(int workers, int low_cap) {
    { std::lock_guard<std::mutex> lk(m_); low_cap_ = low_cap; }
    ensure_started(workers);
  }
  // 懒启动：侧区预取的 GPU 回调可能在 Python 没调用 bg_reader_start 时就要派任务，
  // 故 submit/submit_task 自动确保线程已起（默认 4 worker），调用方无需显式 start。
  void ensure_started(int workers) {
    std::lock_guard<std::mutex> lk(m_);
    if (running_) return;
    running_ = true;
    for (int i = 0; i < workers; ++i) threads_.emplace_back([this] { loop(); });
  }
  void submit(BgJob job) {
    ensure_started(4);
    { std::lock_guard<std::mutex> lk(m_);
      if (job.prio > 0) high_q_.push(std::move(job)); else low_q_.push(std::move(job)); }
    cv_.notify_all();
  }
  void submit_task(std::function<void()> fn) {
    ensure_started(4);
    BgJob job;
    job.task = std::move(fn);
    { std::lock_guard<std::mutex> lk(m_); low_q_.push(std::move(job)); }   // 通用任务走低优
    cv_.notify_all();
  }
  bool ready(long ticket) {
    std::lock_guard<std::mutex> lk(dm_);
    return done_.count(ticket) > 0;
  }
  void wait(long ticket) {
    std::unique_lock<std::mutex> lk(dm_);
    dcv_.wait(lk, [&] { return done_.count(ticket) > 0; });
  }
  void stop() {
    { std::lock_guard<std::mutex> lk(m_); running_ = false; }
    cv_.notify_all();
    for (auto& t : threads_) if (t.joinable()) t.join();
    threads_.clear();
    { std::lock_guard<std::mutex> lk(dm_); done_.clear(); }
    { std::lock_guard<std::mutex> lk(m_); active_low_ = 0; }
  }
 private:
  int low_budget() const { return low_cap_ <= 0 ? 1 << 30 : low_cap_; }  // <=0 视为不限流
  bool can_take_low() const { return !low_q_.empty() && active_low_ < low_budget(); }
  void loop() {
    std::unordered_map<std::string, int> fds;       // 每线程各自缓存 fd
    while (true) {
      std::unique_lock<std::mutex> lk(m_);
      cv_.wait(lk, [this] { return !high_q_.empty() || can_take_low() || !running_; });
      if (!running_ && high_q_.empty() && low_q_.empty()) break;
      if (!running_ && high_q_.empty() && !can_take_low()) break;  // 退出时低优可超额排空
      bool is_low = false;
      BgJob job;
      if (!high_q_.empty()) {                          // 高优(route)读永远先取
        job = std::move(high_q_.front()); high_q_.pop();
      } else if (can_take_low()) {                     // 低优限流：仅在额度内取
        job = std::move(low_q_.front()); low_q_.pop(); ++active_low_; is_low = true;
      } else {
        continue;                                      // 假醒（低优已满额）→ 回去等
      }
      lk.unlock();
      if (job.task) { job.task(); }                    // 通用任务（侧区异步读）：直接执行，无 ticket
      else {
        int fd;
        std::string key = (job.nocache ? "N:" : "C:") + job.path;   // 按 nocache 分别缓存 fd
        auto it = fds.find(key);
        if (it == fds.end()) {
          fd = job.nocache ? open_blob_nocache(job.path.c_str()) : ::open(job.path.c_str(), O_RDONLY);
          fds[key] = fd;
        } else fd = it->second;
        if (fd >= 0)
          for (auto& op : job.ops) {
            ssize_t got = ::pread(fd, op.dst, op.nbytes, op.file_off);
            if (got != static_cast<ssize_t>(op.nbytes))
              fprintf(stderr, "[bg pread SHORT] got=%zd want=%zu off=%lld\n",
                      got, op.nbytes, static_cast<long long>(op.file_off));
          }
        { std::lock_guard<std::mutex> lk2(dm_); done_.insert(job.ticket); }
        dcv_.notify_all();
      }
      if (is_low) {                                    // 释放低优额度 → 唤醒别的 worker 再取
        std::lock_guard<std::mutex> lk2(m_); --active_low_; cv_.notify_all();
      }
    }
    for (auto& kv : fds) if (kv.second >= 0) ::close(kv.second);
  }
  std::mutex m_, dm_;
  std::condition_variable cv_, dcv_;
  std::queue<BgJob> high_q_, low_q_;
  int active_low_ = 0;             // 当前正在执行的低优读数
  int low_cap_ = 0;               // 低优并发上限（<=0 不限流，保持旧行为）
  std::unordered_set<long> done_;
  std::vector<std::thread> threads_;
  bool running_ = false;
};
BgReader g_bg;
}  // namespace

void bg_submit_task(std::function<void()> fn) { g_bg.submit_task(std::move(fn)); }

void bg_reader_start(int workers, int low_cap) { g_bg.start(workers, low_cap); }

long bg_reader_submit(const mx::array& dst, const std::vector<int>& experts,
                      const std::vector<int>& rows, const std::string& path,
                      int stride, long ticket, int prio) {
  mx::array d = dst;
  d.eval();
  uint8_t* base = d.data<uint8_t>();
  size_t st = static_cast<size_t>(stride);
  BgJob job;
  job.path = path;
  job.ticket = ticket;
  job.prio = prio;
  job.keep.push_back(d);
  for (size_t i = 0; i < experts.size(); ++i)
    job.ops.push_back(ReadOp{base + static_cast<size_t>(rows[i]) * st, st,
                             static_cast<off_t>(static_cast<size_t>(experts[i]) * st)});
  g_bg.submit(std::move(job));
  return ticket;
}

// 专家各段直写进多个池段张量的 slot 行（消费侧零 MLX 算子）。
long bg_pread_into_pool(
    const std::vector<mx::array>& dst,
    const std::vector<long>& seg_off,
    const std::vector<long>& seg_nb,
    long slot, long expert,
    const std::string& path, long stride, long ticket, int prio, bool nocache) {
  BgJob job;
  job.path = path;
  job.ticket = ticket;
  job.prio = prio;
  job.nocache = nocache;
  for (size_t i = 0; i < dst.size(); ++i) {
    mx::array d = dst[i];
    d.eval();
    uint8_t* base = d.data<uint8_t>();
    job.keep.push_back(d);
    job.ops.push_back(ReadOp{
        base + static_cast<size_t>(slot) * static_cast<size_t>(seg_nb[i]),
        static_cast<size_t>(seg_nb[i]),
        static_cast<off_t>(static_cast<size_t>(expert) * static_cast<size_t>(stride)
                           + static_cast<size_t>(seg_off[i]))});
  }
  g_bg.submit(std::move(job));
  return ticket;
}

bool bg_reader_ready(long ticket) { return g_bg.ready(ticket); }
void bg_reader_wait(long ticket) { g_bg.wait(ticket); }
void bg_reader_stop() { g_bg.stop(); }