// SDQ expert pager：SSD → RAM → CPU 的 MoE 專家分層快取
//
// 設計要點
//   * 分頁單位 = 一個 expert（layer, expert id），大小 ~1.8 MB（三個矩陣：gate/up/down）
//   * RAM 端是一個固定槽位的 arena（slab），每槽 2 MiB，O(1) 配置、不碎片化
//   * 淘汰策略：LRU + 本 token 使用中（pinned）的保護
//   * 預取：依「上一個 token 的專家選擇」用一階 Markov 轉移預測下一個 token 的專家，
//           交給背景執行緒 pread()，量測「預取命中率 / 預取浪費量」
//   * 指標：hit/miss、SSD 讀取位元組、demand 停滯時間、prefetch 領先時間、arena 使用量
//
// 不使用 swap：所有權重資料只存在 arena（實體 RAM）與檔案（SSD）兩處。

#pragma once

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

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

namespace sdq {

// ---------------------------------------------------------------- GGUF 讀取
struct gguf_info {
    int64_t n_layer = 0;
    int64_t n_expert = 0;
    int64_t n_expert_used = 0;
    int64_t n_embd = 0;
    int64_t n_ff_exp = 0;
    int64_t n_ff_shexp = 0;
    int64_t n_head = 0;
    int64_t n_head_kv = 0;
    int64_t n_embd_head_k = 0;
    int64_t n_embd_head_v = 0;
    int64_t context_length = 0;
    int64_t full_attention_interval = 4;
    int64_t is_recr[512] = {0};
    std::string arch;
    std::string name;
};

// 一個 expert 槽位：gate/up/down 三段（有些模型把 gate+up 合成 gate_up）
struct seg_desc {
    uint64_t off = 0;  // 檔案內位元組偏移
    uint64_t size = 0; // 位元組數
};

enum part : int { GATE = 0, UP = 1, DOWN = 2, COUNT = 3 };

struct expert_desc {
    seg_desc seg[3];
    int64_t slot_off[3] = {0, 0, 0}; // 三段在 arena 槽位內的相對偏移
    int64_t ne0 = 0;         // gate/up row 長度（= n_embd）
    int64_t ne1 = 0;         // gate/up row 數（= n_ff_exp）
    int64_t down_ne0 = 0;    // down row 長度（= n_ff_exp）
    int64_t down_ne1 = 0;    // down row 數（= n_embd）
    ggml_type type[3] = {GGML_TYPE_COUNT, GGML_TYPE_COUNT, GGML_TYPE_COUNT};
    int64_t bytes = 0;       // 三段總大小
    bool fused_gate_up = false;
};

// 一層的所有 expert 描述（40 層 × 256 個）
struct layer_desc {
    std::vector<expert_desc> experts;
    int64_t bytes_per_expert = 0;
    // 這個層的 expert 比「大多數層」大（Q4_K/Q5_K vs Q6_K）→ 需要大槽位。
    // 不然的話所有槽位都得按最大的那種配，白白浪費 6.8% 的 arena。
    bool needs_big_slot = false;
};

// ---------------------------------------------------------------- 指標
struct stats {
    std::atomic<uint64_t> expert_requests{0};   // 每個 (layer, token) 需要幾個 expert
    std::atomic<uint64_t> expert_hits{0};
    std::atomic<uint64_t> expert_misses{0};
    std::atomic<uint64_t> ssb_read_bytes_demand{0};
    std::atomic<uint64_t> ssb_read_bytes_prefetch{0};
    std::atomic<uint64_t> prefetch_issued{0};
    std::atomic<uint64_t> prefetch_completed{0};
    std::atomic<uint64_t> prefetch_used{0};     // 預取後真的被用到
    std::atomic<uint64_t> prefetch_duplicate{0};// 預取時已經在 RAM 裡
    std::atomic<uint64_t> evictions{0};
    std::atomic<uint64_t> demand_wait_us{0};    // 因為缺頁而停下來等 SSD 的時間
    std::atomic<uint64_t> prefetch_wait_us{0};
    std::atomic<uint64_t> tokens{0};
    std::atomic<uint64_t> moe_calls{0};
    std::atomic<uint64_t> flops{0};
};

// ---------------------------------------------------------------- pager
struct config {
    size_t   arena_bytes   = 0;   // 0 = 自動（依 RAM 預算）
    size_t   slot_bytes    = 0;   // 0 = 自動（依模型推導）
    int      n_threads     = 0;
    bool     prefetch      = true;
    int      prefetch_ahead_layers = 0; // >0 時限制同時在飛的預取數量
    bool     fadvise_dontneed = true;  // 讀完丟掉 kernel page cache，讓 RAM 帳目誠實
    bool     verbose       = false;
    uint64_t ram_budget_bytes = 0;    // 0 = 自動（實體 RAM 的一小部分）
    std::string model_path;
    std::string stats_path;
    int      slot_bias = 0;

    // ---- I/O 執行緒池 ----
    // 為什麼需要：一個 expert 在檔案上是 gate / up / down 三段「不相鄰」的資料，
    // 循序讀要 3 次 pread，而單一 pread 在本機量到約 7 ms（1.9 MB / 275 MB/s）。
    // 也就是說「讀一個 expert」的延遲完全由這三段循序決定，執行緒再多也沒用。
    // 把三段交給三條執行緒同時讀，單一 expert 的延遵掉到約 1/2~1/3。
    // （實測：8 個 expert 全部 miss 的一層，序列讀 12.6 ms → 段平行 7.6 ms。）
    bool     io_parallel = true;      // SDQ_IO_PARALLEL=0 可關（回到序列讀）
    bool     io_uring    = false;     // SDQ_IO_URING=1：以 liburing 提交同批非同步讀取
    int      io_threads   = 0;         // 0 = 自動（3 × n_threads，上限 48）
    int      io_split     = 1;         // SDQ_IO_SPLIT：每一段再切成幾塊並行讀（實測 1 最好）

    // ---- Strata 式「熱門 expert 剖面」（OPT-1）----
    // 離線用代表性文字跑一次，把每個 (layer, expert) 的請求次數記成剖面；
    // 執行時把排行榜前 hot_pin 個釘成「熱頁」：熱頁可以被讀取，但**永遠不被淘汰**，
    // 淘汰時一律先從冷頁挑。這解決的是「prefill 的一次性專家把熱區洗掉」。
    std::string profile_in;    // 讀取剖面（SDQ_PROFILE_IN）
    std::string profile_out;   // 結束時輸出剖面（SDQ_PROFILE_OUT）
    int         hot_pin = 0;   // 釘選前幾名（SDQ_HOT_PIN）
    int         admit_min_hits = 1;  // 請求次數低於此值就不收進快取（SDQ_ADMIT_MIN_HITS）
};

class pager {
public:
    static pager & instance();

    // 從 GGUF 檔頭建立 expert 頁面表，並配置 RAM arena
    void init(const config & cfg, ggml_context * ctx);

    bool ready() const { return ready_; }
    size_t arena_bytes() const { return arena_bytes_; }
    size_t slot_bytes() const { return slot_bytes_; }
    size_t n_slots() const { return n_slots_; }
    size_t small_bytes() const { return small_bytes_; }
    size_t big_bytes() const { return big_bytes_; }
    size_t n_small_slots() const { return n_small_; }
    size_t n_big_slots() const { return n_big_; }
    const gguf_info & info() const { return info_; }

    // 取得 layer 描述（idx 越界或沒資料回 nullptr）
    const layer_desc * layer(int il) const;

    // 一次把「連續編號的一串 expert」讀進 RAM：同一層的三個矩陣（gate/up/down）本來
    // 就各自連續存放，所以 k 個連續 expert 只需要 3 次大讀取（而不是 3k 次）。
    // 對高延遲的儲存裝置（每個 request 都貴）這是最大的單項加速。
    //   experts 必須已依編號遞增排序；slots[i] 是 experts[i] 對應的槽位（由呼叫端配置）
    // 回傳實際讀取的位元組數。
    uint64_t read_expert_span(int il, const int32_t * experts, int n, const int32_t * slots);

    // 只配置槽位（不讀資料）：回傳槽位索引，-1 表示沒有空槽
    int reserve_slot(int il, int ie);

    // 查這個 expert 是否已經在 RAM：命中則 pin 並回傳槽位，未命中回傳 -1（不配置）
    int probe(int il, int ie);

    // 記錄一次合併讀取的位元組數（給指標用）
    void note_span_read(uint64_t bytes);

    // 確保 (layer, expert) 在 RAM，回傳槽位索引；-1 表示失敗
    // pinned: 本 token 使用中，不可被淘汰；blocking: 缺頁時是否同步讀 SSD
    int acquire(int il, int ie, bool blocking);

    // 背景預取：把 page 讀進 arena（非同步），依歷史預測下一個 token 會用到的 expert
    void prefetch_expert(int il, int ie);
    // 「下一層」預取：在 layer L 的 op 開頭呼叫，替 layer L+1 排預取。
    // 與 end_token() 的差別在時機 —— 後者在一個 token 的 40 層全算完才跑，
    // 那時排的預取沒有任何計算可以重疊，實測使用率只有 0~8 %。
    void prefetch_layer_ahead(int il);

    // 每層計算開始時呼叫：解除上一層的 pin（本層用到的槽位才受保護）
    void unpin_all();
    // 用完一個槽位後立刻放掉 pin（成對呼叫 acquire()）。
    //
    // 為什麼需要：原本 pin 是「整層釘住、下一層開頭才 unpin_all」。mode C 的
    // expert_forward 回傳之後就再也不讀那個槽位（最後的歸約只讀各執行緒自己的
    // partial 累加緩衝），所以繼續釘著只是白白把工作集撐大。
    // 撐大的後果很具體：prefill 一個 ubatch 有 128 個 token，一層最多需要
    // 256 個不同的 expert → 整層釘住 = 256 個槽位不能回收。大小兩種槽位之下，
    // 大槽位只有約 n_big_layers × n_expert_used 個，於是**必定**在 prefill 時
    // 「找不到能放 il=34 ie=… 的槽位」而中止。
    void release(int s);
    void note_used(int il, int ie);
    void end_token();

    // 依歷史預測某層下一個 token 的 top-k expert（回傳信心 0..1）
    void predict(int il, int32_t * out, int k);

    void shutdown();
    stats & get_stats() { return stats_; }

    // 把當前指標寫成一行 JSON（append 到 stats_path）
    void dump_stats(const char * phase);
    void print_summary() const;

    // arena 內的槽位指標
    // 槽位索引空間是統一的：[0, n_small_) 是小槽位（連續），
    // [n_small_, n_slots_) 是大槽位（接在小槽位區之後）。
    uint8_t * slot_ptr(int s) {
        return (size_t) s < n_small_
             ? arena_ + (size_t) s * small_bytes_
             : arena_ + small_bytes_ * n_small_ + (size_t)(s - (int) n_small_) * big_bytes_;
    }
    size_t slot_size_of(int s) const {
        return (size_t) s < n_small_ ? small_bytes_ : big_bytes_;
    }
    bool need_big_slot(int il, int ie) const;

    // 這個槽位算不算「大槽位」。
    //
    // 判斷依據是**槽位的實際大小**，不是索引 >= n_small_：
    // 統一尺寸（SDQ_FORCE_UNIFORM_SLOT=1，或模型只有一種 expert 大小）時，
    // 所有槽位的索引都 < n_small_ 但大小都是大的 —— 用索引判斷會讓大 expert
    // 永遠配不到槽位，整個推論卡在 admit 的重試迴圈裡。
    bool is_big_slot(int s) const {
        return big_bytes_ && slot_size_of(s) >= big_bytes_;
    }

    size_t used_bytes() const { return used_bytes_snapshot(); }
    double  hit_rate() const;

    // arena 裡實際被占用的位元組數。大小兩種槽位之後不能用「槽位數 × slot_bytes_」
    // 計算（slot_bytes_ 是大的那一種），所以真的把槽位走一遍。
    // 呼叫端：dump_stats（每次解碼階段結束）與 print_summary，不是熱路徑。
    size_t used_bytes_snapshot() const {
        size_t b = 0;
        for (size_t s = 0; s < slots_.size(); s++) {
            if (slots_[s].layer >= 0) {
                b += slot_size_of((int) s);
            }
        }
        return b;
    }

    // 讀 SSD（同步，給 acquire 用）
    bool read_expert(int il, int ie, uint8_t * dst);

    // ---- I/O 執行緒池（見 config::io_parallel 的說明）----
    // 一個「讀取批次」的完成狀態；放在呼叫端的 stack 上，worker 只在 remaining==0
    // 之後才會再碰它，所以離開 read_expert_par() 之後不會有 dangling pointer。
    struct io_ctx {
        std::atomic<int> remaining{0};
        bool ok = true;
    };
    struct io_job {
        uint64_t off;
        uint8_t * dst;
        size_t   size;
        io_ctx * ctx;
    };

    void io_pool_start();
    void io_pool_stop();
    void io_worker(size_t wid);
    bool read_expert_par(int il, int ie, uint8_t * dst);
    bool read_expert_uring(int il, int ie, uint8_t * dst);
    bool read_range_staged(uint64_t off, uint8_t * dst, size_t size,
                           uint8_t * stage, size_t stage_sz);

    // I/O 執行緒池每條 worker 自帶的 O_DIRECT 對齊緩衝（啟動時配置，已計入 RSS）
    struct io_stage { uint8_t * buf = nullptr; size_t sz = 0; };
    std::vector<io_stage> io_stage_;

    bool     io_pool_on_ = false;
    int      io_threads_ = 0;
    int      io_split_   = 1;
    std::vector<std::thread> io_pool_;
    std::mutex io_mu_;
    std::condition_variable io_cv_;       // 有工作 / 要停
    std::condition_variable io_done_cv_;  // 某個批次讀完了
    std::deque<io_job> io_q_;
    bool io_stop_ = false;
    std::atomic<uint64_t> io_jobs_{0};
    std::atomic<uint64_t> io_wait_ns_{0};

private:
    pager() = default;
    ~pager();

    struct slot_state {
        int32_t layer  = -1;
        int32_t expert = -1;
        uint64_t last_use = 0;   // LRU 記時
        uint32_t freq = 0;       // 這個 expert 被用到幾次（衰減式；用來做 LFU 混合策略）
        uint32_t pin = 0;        // >0 時不可淘汰
        bool busy = false;       // 背景執行緒正在寫入
        bool is_hot = false;     // 屬於剖面排行榜：淘汰時最後才會被選中
    };

    struct pf_item {
        int slot;
        int il;
        int ie;
        uint64_t issue_time_us;
    };

    void touch_freq_locked(slot_state & st);
    // ---- OPT-1：剖面 / 熱頁 ----
    void  count_request(int il, int ie);            // 記一次請求（給剖面用）
    bool  load_profile(const std::string & path);   // 讀剖面並標記前 hot_pin 名
    void  dump_profile(const std::string & path);   // 輸出排行榜
    bool  hot_lookup(int il, int ie) const;
    void load_gguf_header(const std::string & path);
    void build_page_map(const std::string & path);
    void alloc_arena();
    int  find_slot_locked(int il, int ie);
    int  admit_slot_locked(int il, int ie);
    int  victim_locked() { return admit_slot_locked(-1, -1); }
    int  victim_stale_locked(uint64_t stale_before);
    int  victim_stale_for(int il, int ie, uint64_t stale_before);
    void pf_worker();
    bool read_range(uint64_t off, uint8_t * dst, size_t size);
    bool read_range_direct(uint64_t off, uint8_t * dst, size_t size);
    void open_model_fd();

    bool     ready_ = false;
    config   cfg_;
    gguf_info info_;
    std::vector<layer_desc> layers_;

    uint8_t * arena_ = nullptr;
    size_t    arena_bytes_ = 0;
    size_t    slot_bytes_ = 0;   // = max(small_bytes_, big_bytes_)，保留給舊的訊息輸出
    size_t    small_bytes_ = 0;  // 小槽位大小（大多數層用）
    size_t    big_bytes_   = 0;  // 大槽位大小；0 = 這個模型沒有大層
    size_t    n_small_     = 0;
    size_t    n_big_       = 0;
    size_t    n_slots_ = 0;
    int       fd_ = -1;
    int       fd_direct_ = -1;   // O_DIRECT：完全繞過 kernel page cache
    bool      use_direct_ = false;

    mutable std::mutex mu_;
    std::vector<slot_state> slots_;
    std::unordered_map<int64_t, int32_t> page_; // (layer<<16|expert) -> slot
    std::vector<int32_t> free_slots_;   // 空的小槽位
    std::vector<int32_t> free_big_slots_;

    // ---- OPT-1：剖面 / 熱頁（Strata 式：靜態排行榜 + 熱頁不淘汰）----
    std::vector<uint32_t> prof_;   // 每個 (layer, expert) 的累計請求次數
    std::vector<uint8_t>  hot_;    // 是否在排行榜前 hot_pin 名
    int                   n_hot_ = 0;
    std::string           profile_out_;
    uint32_t freq_tick_ = 0;
    std::atomic<size_t> used_{0};
    uint64_t clock_ = 0;

    // 背景預取
    std::thread pf_;
    std::mutex pf_mu_;
    std::condition_variable pf_cv_;
    std::deque<pf_item> pf_q_;
    bool pf_stop_ = false;
    std::atomic<uint64_t> pf_inflight_{0};

    // 路由歷史（一階 Markov：上個 token 的 expert -> 這次被選到的次數）
    struct trans { int32_t to[64]; uint32_t cnt[64]; int32_t n = 0; };
    std::vector<std::vector<trans>> hist_; // [layer][prev_expert]
    std::vector<std::vector<int32_t>>  prev_used_; // [layer] 上一個 token 用到的 expert（預測用）
    std::vector<std::vector<int32_t>>  cur_used_;  // [layer] 當前 token 用到的 expert
    std::mutex hist_mu_;   // 保護 cur_used_ / prev_used_ / hist_（多執行緒會同時 note_used）
    std::atomic<uint64_t> hist_lock_{0};

    stats stats_;
};

// ---------------------------------------------------------------- 工具
uint64_t now_us();
size_t   process_rss_bytes();
size_t   process_anon_bytes();
size_t   enforced_memory_ceiling_bytes();
size_t   physical_ram_bytes();

// ---------------------------------------------------------------- ggml 介接
// 自訂 buffer type：expert 權重佔 0 個 RAM 位元組，資料留在檔案裡
ggml_backend_buffer_type_t sdq_expert_buft();

// 註冊 model params 的 tensor buffer override（expert 權重 → 我們的 buffer）
void sdq_install_buft_overrides(llama_model_params * params);

// 建立被分頁的 MoE FFN（替代 build_moe_ffn 裡的三次 mul_mat_id）
ggml_tensor * sdq_build_moe(ggml_context * ctx, ggml_tensor * cur,
                            ggml_tensor * ids, ggml_tensor * weights, int il,
                            int64_t n_expert, int64_t n_expert_used,
                            float swiglu_limit);

// 設定 arena 大小 / 預取開關等（CLI 在載入模型前呼叫）
bool sdq_init(const std::string & model_path);
void sdq_set_config(const config & c);
const config & sdq_get_config();
void sdq_set_enabled(bool on);
void sdq_shutdown();

bool sdq_enabled();

// 效能計時探針的結果（SDQ_PROFILE_TIME=1）
void sdq_print_time_profile();

// 自我測試：算一次指定 expert 的 MoE 輸出（給 NumPy 對照用）
int sdq_moe_selftest(int il, int ie, const float * x, float * out_io,
                     int64_t n_embd, int64_t n_ff, int n_tok);

// 建立我們的 MoE 自訂節點。
// **必須把 gf 傳進來並呼叫 ggml_build_forward_expand(gf, out)**：
// ggml 的排程器只會為「進了圖的節點」配置 compute buffer，沒註冊的節點其輸出
// 緩衝區會被 liveness 分析判成可重用，導致 layer>=1 的結果被其他張量覆蓋
// （實測症狀：layer 0 正確、layer 1 起整個錯掉）。
ggml_tensor * sdq_build_moe(ggml_cgraph * gf, ggml_context * ctx, ggml_tensor * cur,
                            ggml_tensor * ids, ggml_tensor * weights, int il,
                            int64_t n_expert, int64_t n_expert_used, float swiglu_limit);

// 這一層的分頁是否可用（graph build 時呼叫）
bool sdq_pager_active_for_layer(int il);

// 除錯／驗證用：把某個 F32 張量 dump 到檔案。
// 必須在「原生 mul_mat_id 路徑」與「SDQ 分頁路徑」兩邊都呼叫，才能做真正的 A/B 對照
// （env: SDQ_MOE_OUT_DUMP=prefix、SDQ_MOE_OUT_DUMP_LAYER=層編號，-1 = 全部）
ggml_tensor * sdq_maybe_dump_tensor(ggml_context * ctx, ggml_tensor * t, int il);


} // namespace sdq