// SDQ 聊天 / benchmark 主程式
//
//   sdq-chat -m <model.gguf> [-c ctx] [-t threads] [-n n_predict] [-p prompt] ...
//
// 這個工具存在的理由：官方 llama-cli 不會安裝我們的 tensor buffer override，
// 而且需要我們在載入模型前後插入 pager 的初始化/收尾。用最少的相依度
// （只鏈 llama + ggml）自己寫一支，行為完全可控。

#include "llama.h"
#include "sdq_pager.h"

#include <algorithm>
#include <chrono>
#include <cinttypes>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <thread>
#include <iostream>
#include <string>
#include <vector>

using namespace sdq;

namespace {

struct options {
    std::string model;
    std::string prompt;
    std::string system = "You are a helpful assistant. Answer in the user's language.";
    std::string stats_file;
    int32_t n_ctx = 4096;          // 驗證範圍：4K（理論支援 1M）
    int32_t n_predict = 256;
    int32_t n_threads = 0;
    int32_t n_ubatch = 0;
    int32_t n_keep = 0;
    float temp = 0.7f;
    float top_p = 0.95f;
    int32_t top_k = 20;
    bool chat = true;
    bool interactive = true;
    bool bench = false;
    bool no_prefetch = false;
    bool dump_logits = false;
    std::string logits_out;
    std::string selftest;
    std::string span_test;
    std::string dump_tensors;   // ggml eval callback：把指定張量的資料存檔（A/B 比對用）
    bool no_pager = false;      // A/B 測試用：不啟用分頁，走上游 mul_mat_id（整個模型在 RAM）
    int32_t arena_mb = 0;
    int32_t ram_budget_mb = 0;
    bool fadvise = true;
};

double now_s() {
    static auto t0 = std::chrono::steady_clock::now();
    return std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
}

size_t rss_mb() { return process_rss_bytes() / (1024 * 1024); }

// 從 /proc/self/status 讀 VmHWM（峰值 RSS）/ VmSwap（這個行程真的用了多少 swap）
size_t proc_status_kb(const char * key) {
    FILE * f = fopen("/proc/self/status", "r");
    if (!f) return 0;
    char line[256];
    const size_t klen = strlen(key);
    size_t val = 0;
    while (fgets(line, sizeof(line), f)) {
        if (strncmp(line, key, klen) == 0) {
            val = (size_t) strtoull(line + klen, nullptr, 10);
            break;
        }
    }
    fclose(f);
    return val;
}

size_t peak_rss_mb() { return proc_status_kb("VmHWM:") / 1024; }
size_t swap_mb()    { return proc_status_kb("VmSwap:") / 1024; }

// /proc/self/smaps_rollup：把 RSS 拆成「檔案對映」與「匿名」兩塊。
// 這很重要：模型檔是 mmap 進來的，核心預讀會讓「檔案對映頁」也算進 RSS，
// 報告時必須分開看，否則會誤以為 RAM 用量比實際少（或多）。
size_t smaps_kb(const char * key) {
    FILE * f = fopen("/proc/self/smaps_rollup", "r");
    if (!f) return 0;
    char line[256];
    const size_t klen = strlen(key);
    size_t val = 0;
    while (fgets(line, sizeof(line), f)) {
        if (strncmp(line, key, klen) == 0) {
            val = (size_t) strtoull(line + klen, nullptr, 10);
            break;
        }
    }
    fclose(f);
    return val;
}

void print_mem_report(const char * when) {
    fprintf(stderr,
            "[sdq] 記憶體（%s）：匿名 RSS %zu MB / 檔案對映 RSS %zu MB / 總 RSS %zu MB / 峰值 %zu MB / swap %zu MB\n",
            when, smaps_kb("Anonymous:") / 1024, smaps_kb("RssFile:") / 1024,
            smaps_kb("Rss:") / 1024, peak_rss_mb(), swap_mb());
}

// 讀 /proc/self/io：從 block device 實際讀了多少（SSD 讀取）
uint64_t read_bytes_io() {
    FILE * f = fopen("/proc/self/io", "r");
    if (!f) return 0;
    char line[256];
    uint64_t v = 0;
    while (fgets(line, sizeof(line), f)) {
        if (strncmp(line, "read_bytes:", 11) == 0) {
            v = strtoull(line + 11, nullptr, 10);
        }
    }
    fclose(f);
    return v;
}

struct chat_state {
    std::vector<llama_chat_message> msgs;
};

std::string build_prompt(const options & opt, const chat_state & st) {
    if (!opt.chat) {
        return opt.prompt;
    }
    const char * tmpl = llama_model_chat_template(nullptr, nullptr);
    // 沒有 model 指標時用內建模板；下面 main() 會用實際模型的模板重做
    (void) tmpl;
    size_t n_chars = 0;
    for (const auto & m : st.msgs) {
        n_chars += strlen(m.role) + strlen(m.content) + 32;
    }
    std::vector<char> buf(2 * n_chars + 1024);
    const char * t = st.msgs.empty() ? nullptr : nullptr;
    (void) t;
    int32_t len = llama_chat_apply_template(tmpl, st.msgs.data(), st.msgs.size(), true,
                                            buf.data(), (int32_t) buf.size());
    if (len > (int32_t) buf.size()) {
        buf.resize(len + 1);
        llama_chat_apply_template(tmpl, st.msgs.data(), st.msgs.size(), true,
                                  buf.data(), (int32_t) buf.size());
    }
    return std::string(buf.data());
}

} // namespace

#include <execinfo.h>
#include <signal.h>
#include <unistd.h>
static void sdq_segv(int sig) {
    void * bt[64];
    const int n = backtrace(bt, 64);
    fprintf(stderr, "\n*** sdq: signal %d, backtrace:\n", sig);
    backtrace_symbols_fd(bt, n, 2);
    _exit(9);
}

// eval callback：把感興趣的張量（用 SDQ_DUMP_NAMES 指定）寫進檔案
struct dump_state {
    FILE * fp = nullptr;
    uint64_t step = 0;
    std::vector<std::string> names;
};

static dump_state g_dump;

static bool sdq_eval_callback(ggml_tensor * t, bool ask, void * user_data) {
    (void) user_data;
    auto & ds = g_dump;
    if (ask) {
        // 除錯模式：每個節點都要資料 → 會強制逐節點同步（慢，但能逐節點比對）
        return ds.fp != nullptr;
    }
    if (!ds.fp) {
        return false;
    }
    const char * nm = ggml_get_name(t);
    const size_t n_len = nm ? strlen(nm) : 0;
    bool want = ds.names.empty();
    for (const auto & n : ds.names) {
        if (nm && strstr(nm, n.c_str())) {
            want = true;
            break;
        }
    }
    const int32_t ne[4] = {(int32_t) t->ne[0], (int32_t) t->ne[1], (int32_t) t->ne[2], (int32_t) t->ne[3]};
    const int32_t type = (int32_t) t->type;
    fwrite(&n_len, sizeof(size_t), 1, ds.fp);
    if (n_len) {
        fwrite(nm, 1, n_len, ds.fp);
    }
    fwrite(&ds.step, sizeof(uint64_t), 1, ds.fp);
    fwrite(ne, sizeof(int32_t), 4, ds.fp);
    fwrite(&type, sizeof(int32_t), 1, ds.fp);

    // 永遠寫「摘要」（前 8 個值 + 平方和 + 最大值），這樣即使沒有名字也能 A/B 找出第一個分歧點
    float summary[10] = {0};
    if (t->type == GGML_TYPE_F32 && t->data) {
        const size_t n = (size_t) ne[0] * ne[1] * ne[2] * ne[3];
        const float * p = (const float *) t->data;
        double ss = 0;
        float mx = 0;
        for (size_t i = 0; i < n; i++) {
            ss += (double) p[i] * p[i];
            if (std::fabs(p[i]) > std::fabs(mx)) {
                mx = p[i];
            }
        }
        for (int i = 0; i < 8 && (size_t) i < n; i++) {
            summary[i] = p[i];
        }
        summary[8] = (float) std::sqrt(ss);
        summary[9] = mx;
        fwrite(summary, sizeof(float), 10, ds.fp);
        if (want) {
            fwrite(p, sizeof(float), n, ds.fp);
        }
    } else {
        fwrite(summary, sizeof(float), 10, ds.fp);
    }
    ds.step++;
    return true;   // 這個節點的資料我們已經讀過了，標記成已完成
}

int main(int argc, char ** argv) {
    signal(SIGSEGV, sdq_segv);
    signal(SIGBUS, sdq_segv);
    signal(SIGABRT, sdq_segv);
    options opt;
    const char * env_model = getenv("SDQ_MODEL_PATH");
    if (env_model) {
        opt.model = env_model;
    }

    for (int i = 1; i < argc; i++) {
        std::string a = argv[i];
        auto next = [&](const char * what) -> std::string {
            if (i + 1 >= argc) {
                fprintf(stderr, "%s 需要參數\n", what);
                exit(1);
            }
            return argv[++i];
        };
        if (a == "-m" || a == "--model")        opt.model = next("-m");
        else if (a == "-p" || a == "--prompt")  opt.prompt = next("-p");
        else if (a == "--sys")                  opt.system = next("--sys");
        else if (a == "-c" || a == "--ctx")     opt.n_ctx = atoi(next("-c").c_str());
        else if (a == "-n")                     opt.n_predict = atoi(next("-n").c_str());
        else if (a == "-t")                     opt.n_threads = atoi(next("-t").c_str());
        else if (a == "-b" || a == "--ubatch")  opt.n_ubatch = atoi(next("-b").c_str());
        else if (a == "--temp")                 opt.temp = atof(next("--temp").c_str());
        else if (a == "--top-p")                opt.top_p = atof(next("--top-p").c_str());
        else if (a == "--top-k")                opt.top_k = atoi(next("--top-k").c_str());
        else if (a == "--no-chat")              opt.chat = false;
        else if (a == "--no-interactive")       opt.interactive = false;
        else if (a == "--bench")                opt.bench = true;
        else if (a == "--no-prefetch")          opt.no_prefetch = true;
        else if (a == "--arena-mb")             opt.arena_mb = atoi(next("--arena-mb").c_str());
        else if (a == "--ram-budget-mb")        opt.ram_budget_mb = atoi(next("--ram-budget-mb").c_str());
        else if (a == "--no-fadvise")           opt.fadvise = false;
        else if (a == "--stats-file")           opt.stats_file = next("--stats-file");
        else if (a == "--dump-logits")          opt.dump_logits = true;
        else if (a == "--logits-out")           opt.logits_out = next("--logits-out");
        else if (a == "--selftest")            opt.selftest = next("--selftest");
        else if (a == "--dump-tensors")        opt.dump_tensors = next("--dump-tensors");
        else if (a == "--no-pager")            opt.no_pager = true;
        else if (a == "--span-test")           opt.span_test = next("--span-test");
        else if (a == "-h" || a == "--help") {
            printf(
                "用法：sdq-chat -m <model.gguf> [選項]\n"
                "  -c, --ctx N          上下文長度（預設 4096；模型原生 262144）\n"
                "  -n N                 要生成的 token 數\n"
                "  -t N                 CPU 執行緒數\n"
                "  -b N                 ubatch 大小\n"
                "  -p, --prompt TEXT    提示詞\n"
                "  --sys TEXT           system 訊息\n"
                "  --no-chat            不要套 chat template（純 completion）\n"
                "  --no-interactive     跑完提示詞就結束\n"
                "  --bench              benchmark 模式（輸出 JSON 到 stats 檔）\n"
                "  --arena-mb N         expert arena 大小（MB）\n"
                "  --ram-budget-mb N    RAM 總預算（MB），arena 會配合縮減\n"
                "  --no-prefetch        關閉背景預取\n"
                "  --no-fadvise         讀完不丟 kernel page cache（會讓 RAM 帳目失真）\n"
                "  --stats-file PATH    指標 JSONL 輸出\n"
                "  --dump-logits        把每個 token 的 logits 存檔（驗證用）\n"
                "  --logits-out PATH    logits 輸出路徑\n");
            return 0;
        } else {
            fprintf(stderr, "未知參數：%s（用 --help 看用法）\n", a.c_str());
            return 1;
        }
    }

    if (opt.model.empty()) {
        fprintf(stderr, "缺少模型檔（-m 或 SDQ_MODEL_PATH）\n");
        return 1;
    }

    // ---- 環境變數（pager 讀）
    setenv("SDQ_MODEL_PATH", opt.model.c_str(), 1);
    // I/O 讀取執行緒數：預設 3 × 計算執行緒（上限 48）。
    // 必須在這裡設定 —— pager 的 I/O 池是在「載入模型」時啟動的，那已經比
    // context 建立早了；放到後面設定會讓池子只拿到下限 4 條執行緒。
    // 量測（repeat=3，中位數）：序列讀 3.79 → 12 條 5.84 → 24 條 6.68 tok/s。
    if (!getenv("SDQ_IO_THREADS")) {
        const int hw_io = (int) std::max(1u, std::thread::hardware_concurrency());
        const int dflt_io = hw_io < 8 ? hw_io : 8;
        const int n_io   = opt.n_threads > 0 ? opt.n_threads : dflt_io;
        setenv("SDQ_IO_THREADS", std::to_string(n_io * 3).c_str(), 1);
    }
    // 只有明確指定時才設定，否則保留呼叫者從環境變數傳入的值
    // （舊版無條件 setenv 會讓 SDQ_PREFETCH=0 被悄悄蓋回 1，除錯時極易誤判）
    if (opt.no_prefetch) {
        setenv("SDQ_PREFETCH", "0", 1);
    }
    setenv("SDQ_FADV_DONTNEED", opt.fadvise ? "1" : "0", 1);
    if (opt.arena_mb > 0) {
        setenv("SDQ_ARENA_MB", std::to_string(opt.arena_mb).c_str(), 1);
    }
    if (opt.ram_budget_mb > 0) {
        setenv("SDQ_RAM_BUDGET_MB", std::to_string(opt.ram_budget_mb).c_str(), 1);
    }
    if (!opt.stats_file.empty()) {
        setenv("SDQ_STATS_FILE", opt.stats_file.c_str(), 1);
    }

    llama_backend_init();
    llama_numa_init(GGML_NUMA_STRATEGY_DISABLED);

    // ---- 模型載入：expert 權重交給 pager 的 buffer type（0 RAM）
    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = 0;
    mparams.load_mode = LLAMA_LOAD_MODE_MMAP;
    if (!opt.no_pager) {
        sdq_set_enabled(true);
        sdq_install_buft_overrides(&mparams);
    } else {
        printf("[sdq] --no-pager：對照模式，專家權重全部留在 RAM（需要 %s RAM）\n", "22GB");
    }

    printf("[sdq] 載入模型：%s\n", opt.model.c_str());
    const double t_load0 = now_s();
    llama_model * model = llama_model_load_from_file(opt.model.c_str(), mparams);
    if (!model) {
        fprintf(stderr, "[sdq] 模型載入失敗\n");
        return 1;
    }
    const double t_load = now_s() - t_load0;

    // ---- pager 初始化（讀 GGUF 檔頭、建頁面表、配置 arena）
    auto & P = pager::instance();
    if (!opt.no_pager && !sdq_init(opt.model)) {
        fprintf(stderr, "[sdq] pager 初始化失敗\n");
        return 1;
    }

    // ---- 張量 dump（A/B 比對用）：在 op 真的算完的那一刻把資料寫出來
    FILE * tensor_fp = nullptr;
    dump_state dump_state_;
    if (!opt.dump_tensors.empty()) {
        tensor_fp = fopen(opt.dump_tensors.c_str(), "wb");
        dump_state_.fp = tensor_fp;
        if (const char * names = getenv("SDQ_DUMP_NAMES")) {
            std::string s = names;
            size_t pos = 0;
            while (pos <= s.size()) {
                size_t c = s.find(',', pos);
                if (c == std::string::npos) c = s.size();
                dump_state_.names.push_back(s.substr(pos, c - pos));
                pos = c + 1;
            }
        }
        g_dump = dump_state_;
    }

    // ---- context
    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx = (uint32_t) opt.n_ctx;
    cparams.n_batch = (uint32_t) opt.n_ctx;
    // ubatch 預設 128（不是 512）。理由是「8 GB 預算要真的裝得下」：
    // llama.cpp 的 CPU compute buffer 是按 n_ubatch 線性成長的，實測
    //   ubatch=512 → 502 MiB、256 → 251 MiB、128 → 125.5 MiB
    // 而 compute buffer 是在 pager 算完 arena 大小之後才配置的，所以它不算在
    // arena 的扣減裡。實測在 RLIMIT_DATA 8 GB 下：
    //   ubatch=512 → 提示詞超過 ~256 token 直接 std::bad_alloc（崩潰）
    //   ubatch=256 → 同一個提示詞一樣崩潰
    //   ubatch=128 → 正常，681 token 的 prefill 跑到 28.8 tok/s
    // 而且 ubatch=128 還更快（docs/04 實測 +7%），因為 prefill 分更多批、
    // 每批的 expert 聯集比較小、命中率比較高。
    cparams.n_ubatch = opt.n_ubatch > 0 ? (uint32_t) opt.n_ubatch : (uint32_t) std::min(opt.n_ctx, 128);
    // 執行緒預設值：不要直接用 hardware_concurrency()。
    // 實測（docs/04，repeat=3）：mode C 下 t=8 最快，給更多執行緒反而變慢 ——
    //   t=8 4.02 / t=10 3.79 / t=12 3.55 / t=16 3.60 tok/s
    // 原因很直接：每層每 token 只有 n_expert_used = 8 個 expert 可以平行，
    // 執行緒超過 8 之後就只是多開執行緒搶同一批 CPU。
    // 下限取實際核心數，避免在 4 核心的目標機器上開太多。
    const int hw = (int) std::max(1u, std::thread::hardware_concurrency());
    const int dflt_threads = hw < 8 ? hw : 8;
    cparams.n_threads = opt.n_threads > 0 ? opt.n_threads : (int32_t) dflt_threads;
    cparams.n_threads_batch = cparams.n_threads;
    if (!opt.stats_file.empty()) {
        cparams.offload_kqv = false;
    }

    if (tensor_fp) {
        cparams.cb_eval = sdq_eval_callback;
        cparams.cb_eval_user_data = &g_dump;
    }

    llama_context * ctx = llama_init_from_model(model, cparams);
    if (!ctx) {
        fprintf(stderr, "[sdq] context 建立失敗（ctx=%d 太大？）\n", opt.n_ctx);
        return 1;
    }

    printf("[sdq] ctx=%d  threads=%d  ubatch=%u  RAM(RSS)=%zu MB  模型載入 %.1fs\n",
           opt.n_ctx, cparams.n_threads, cparams.n_ubatch, rss_mb(), t_load);
    print_mem_report("載入 + pager 後");

    // ---- sampler
    auto sparams = llama_sampler_chain_default_params();
    sparams.no_perf = true;
    llama_sampler * smpl = llama_sampler_chain_init(sparams);
    llama_sampler_chain_add(smpl, llama_sampler_init_top_k(opt.top_k));
    llama_sampler_chain_add(smpl, llama_sampler_init_top_p(opt.top_p, 1));
    llama_sampler_chain_add(smpl, llama_sampler_init_temp(opt.temp));
    llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));

    chat_state st;
    if (!opt.system.empty()) {
        st.msgs.push_back({"system", opt.system.c_str()});
    }

    std::vector<llama_token> decoded;

    const llama_vocab * vocab = llama_model_get_vocab(model);

    auto tokenize = [&](const std::string & text) {
        std::vector<llama_token> toks(text.size() + 16);
        int32_t n = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(),
                                   toks.data(), (int32_t) toks.size(), true, true);
        if (n < 0) {
            toks.resize((size_t) -n);
            n = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(),
                               toks.data(), (int32_t) toks.size(), true, true);
        }
        toks.resize(n);
        return toks;
    };

    // ---- 送進去一個 ubatch，回傳生成的 token（-1 = 需要繼續生成）
    std::function<int(llama_token, bool)> submit;
    uint64_t n_prefill_tok = 0;
    double t_prefill = 0;
    uint64_t n_decode_tok = 0;
    double t_decode = 0;
    uint64_t io_read0 = read_bytes_io();

    std::vector<float> logits_store;
    FILE * logits_fp = nullptr;
    if (opt.dump_logits) {
        logits_fp = fopen(opt.logits_out.empty() ? "sdq_logits.bin" : opt.logits_out.c_str(), "wb");
    }

    submit = [&](llama_token tok, bool last) -> int {
        const uint32_t ctx_size = llama_n_ctx(ctx);
        llama_memory_t mem = llama_get_memory(ctx);
        if (tok == 1) { // BOS：重開記憶體
            llama_memory_clear(mem, true);
            decoded.clear();
        }
        if (decoded.size() >= opt.n_ctx - 4) {
            // 環形丟掉最舊的 token（保留 opt.n_keep）
            const size_t drop = decoded.size() - (size_t) opt.n_keep;
            llama_memory_seq_rm(mem, 0, 0, (int32_t) drop);
            decoded.erase(decoded.begin(), decoded.begin() + (long) drop);
        }

        llama_batch b = llama_batch_get_one((llama_token *) &tok, 1);

        const bool is_prefill = decoded.empty();
        const double t0 = now_s();
        const int rc = llama_decode(ctx, b);
        const double dt = now_s() - t0;
        if (rc != 0) {
            fprintf(stderr, "\n[sdq] llama_decode 失敗 rc=%d（記憶體不足？）\n", rc);
            exit(3);
        }
        decoded.push_back(tok);

        // 每個 ubatch 結束 → 更新路由歷史並觸發下一輪預取
        P.end_token();

        if (is_prefill) {
            n_prefill_tok++;
            t_prefill += dt;
        } else {
            n_decode_tok++;
            t_decode += dt;
            P.get_stats().tokens.fetch_add(1, std::memory_order_relaxed);
        }

        if (logits_fp) {
            const int32_t n_logits = llama_vocab_n_tokens(vocab);
            const float * lg = llama_get_logits_ith(ctx, -1);
            if (lg) {
                logits_store.assign(lg, lg + n_logits);
                fwrite(&n_logits, sizeof(int32_t), 1, logits_fp);
                fwrite(logits_store.data(), sizeof(float), n_logits, logits_fp);
                fflush(logits_fp);
            }
        }

        if (last) {
            return -1;
        }
        const llama_token id = llama_sampler_sample(smpl, ctx, -1);
        return (int) id;
    };

    // ---- 產生 prompt
    auto run_prompt = [&](const std::string & user_text, bool add_to_history) {
        if (opt.chat) {
            st.msgs.push_back({"user", user_text.c_str()});
        }
        std::string text;
        if (opt.chat) {
            const char * tmpl = llama_model_chat_template(model, nullptr);
            size_t n_chars = 0;
            for (const auto & m : st.msgs) {
                n_chars += strlen(m.role) + strlen(m.content) + 32;
            }
            std::vector<char> buf(2 * n_chars + 1024);
            int32_t len = llama_chat_apply_template(tmpl, st.msgs.data(), st.msgs.size(),
                                                    true, buf.data(), (int32_t) buf.size());
            if (len > (int32_t) buf.size()) {
                buf.resize((size_t) len + 1);
                llama_chat_apply_template(tmpl, st.msgs.data(), st.msgs.size(), true,
                                          buf.data(), (int32_t) buf.size());
            }
            text.assign(buf.data());
        } else {
            text = user_text;
        }

        auto toks = tokenize(text);
        const int n_prompt = (int) toks.size();
        const uint64_t io0 = read_bytes_io();

        // 分批送進去（每批 = 一個 ubatch，控制記憶體）
        const int32_t ub = (int32_t) cparams.n_ubatch;
        llama_sampler_reset(smpl);
        int32_t pos = 0;
        for (int32_t i = 0; i < n_prompt; i += ub) {
            const int32_t n = std::min(ub, n_prompt - i);
            llama_batch b = llama_batch_init((uint32_t) n, 0, 1);
            for (int32_t j = 0; j < n; j++) {
                b.token[j] = toks[i + j];
                b.pos[j] = pos++;
                b.n_seq_id[j] = 1;
                b.seq_id[j][0] = 0;
            }
            b.n_tokens = (uint32_t) n;
            const double t0 = now_s();
            const int rc = llama_decode(ctx, b);
            t_prefill += now_s() - t0;
            n_prefill_tok += (uint64_t) n;
            if (rc != 0) {
                fprintf(stderr, "\n[sdq] prefill 失敗 rc=%d（%d tokens）\n", rc, n);
                exit(3);
            }
            P.end_token();
            llama_batch_free(b);
        }
        const double t_p = now_s();
        (void) t_p;
        (void) io0;

        if (add_to_history && opt.chat) {
            st.msgs.push_back({"assistant", ""});
        }
        return n_prompt;
    };

    // ---- 產生
    auto generate = [&](int32_t n_max) {
        const uint64_t io0 = read_bytes_io();
        const uint64_t io_read_before = io0;
        (void) io_read_before;
        const double t0 = now_s();
        int32_t n_gen = 0;
        llama_token id = llama_sampler_sample(smpl, ctx, -1);
        fputs("assistant: ", stdout);
        fflush(stdout);
        while (n_gen < n_max) {
            const int rc = submit(id, /*last=*/n_gen + 1 >= n_max);
            if (rc < 0) {
                break;
            }
            id = (llama_token) rc;
            char piece_buf[256];
            const int n_piece = llama_token_to_piece(vocab, id, piece_buf, sizeof(piece_buf), 0, true);
            const std::string piece(piece_buf, (size_t) std::max(0, n_piece));
            fputs(piece.c_str(), stdout);
            fflush(stdout);
            n_gen++;
        }
        fputs("\n", stdout);
        const double dt = now_s() - t0;
        const uint64_t io1 = read_bytes_io();
        fprintf(stderr, "[sdq] decode: %d tok / %.2fs = %.2f tok/s | SSD 讀 %.1f MB (%.1f MB/tok) | RSS %zu MB\n",
                n_gen, dt, dt > 0 ? n_gen / dt : 0.0,
                (io1 - io0) / 1048576.0, n_gen ? (io1 - io0) / 1048576.0 / n_gen : 0.0, rss_mb());
        if (opt.chat) {
            for (auto & m : st.msgs) {
                if (strcmp(m.role, "assistant") == 0) {
                    m.content = "";
                }
            }
        }
        P.dump_stats("decode");
        return n_gen;
    };

    // ---- 合併讀取自我測試：--span-test IL:E0:E1:...
    if (!opt.span_test.empty()) {
        std::vector<int> parts;
        std::string cur_s;
        for (char ch : opt.span_test + ":") {
            if (ch == ':') { parts.push_back(atoi(cur_s.c_str())); cur_s.clear(); }
            else cur_s += ch;
        }
        if (parts.size() < 3) {
            fprintf(stderr, "用法：--span-test IL:E0:E1[:E2...]（E 要連續）\n");
            return 1;
        }
        const int il = parts[0];
        std::vector<int32_t> ex;
        for (size_t i = 1; i < parts.size(); i++) {
            ex.push_back(parts[i]);
        }
        std::vector<int32_t> slots(ex.size(), -1);
        for (size_t i = 0; i < ex.size(); i++) {
            slots[i] = P.reserve_slot(il, ex[i]);
        }
        const uint64_t got = P.read_expert_span(il, ex.data(), (int) ex.size(), slots.data());
        printf("[sdq-span-test] L%d experts=%zu 讀取 %.2f MB\n", il, ex.size(), got / 1048576.0);
        int bad = 0;
        for (size_t i = 0; i < ex.size(); i++) {
            std::vector<uint8_t> one((size_t) P.layer(il)->experts[ex[i]].bytes, 0);
            if (!P.read_expert(il, ex[i], one.data())) {
                printf("  單獨讀取失敗\n");
                return 1;
            }
            const uint8_t * span_ptr = P.slot_ptr(slots[i]);
            size_t diff_at = (size_t) -1;
            for (size_t b = 0; b < one.size(); b++) {
                if (one[b] != span_ptr[b]) { diff_at = b; break; }
            }
            printf("  expert %4d: %s", ex[i], diff_at == (size_t) -1 ? "一致 ✅" : "不一致 ❌");
            if (diff_at != (size_t) -1) {
                printf("（第一個不同的 byte 在 %zu / %zu，slot=%d）", diff_at, one.size(), slots[i]);
                bad++;
            }
            printf("\n");
        }
        sdq_shutdown();
        llama_free(ctx);
        llama_model_free(model);
        llama_backend_free();
        return bad ? 1 : 0;
    }

    // ---- 自我測試模式：--selftest IL:IE:SEED:OUT.bin
    if (!opt.selftest.empty()) {
        int il = 0, ie = 0, seed = 1234, n_tok = 1;
        std::string outpath;
        {
            std::vector<std::string> parts;
            std::string cur_s;
            for (char c : opt.selftest + ":") {
                if (c == ':') { parts.push_back(cur_s); cur_s.clear(); }
                else cur_s += c;
            }
            if (parts.size() >= 3) {
                il = atoi(parts[0].c_str());
                ie = atoi(parts[1].c_str());
                seed = atoi(parts[2].c_str());
            }
            if (parts.size() >= 4) {
                outpath = parts[3];
            }
            if (parts.size() >= 5) {
                n_tok = atoi(parts[4].c_str());
            }
        }
        const int64_t n_embd = P.info().n_embd;
        const int64_t n_ff = P.info().n_ff_exp;
        std::vector<float> x((size_t) n_embd * n_tok);
        uint32_t rng = (uint32_t) seed * 2654435761u + 1u;
        for (auto & v : x) {
            rng = rng * 1664525u + 1013904223u;
            v = ((float) (rng >> 8) / (float) (1 << 24)) * 2.0f - 1.0f;
        }
        std::vector<float> io((size_t) n_embd * n_tok * 2);
        const int rc = sdq_moe_selftest(il, ie, x.data(), io.data(), n_embd, n_ff, n_tok);
        if (rc == 0 && !outpath.empty()) {
            FILE * f = fopen(outpath.c_str(), "wb");
            fwrite(io.data(), sizeof(float), io.size(), f);
            fclose(f);
            printf("[sdq] selftest 結果已寫到 %s（layer=%d expert=%d seed=%d n_embd=%lld n_ff=%lld）\n",
                   outpath.c_str(), il, ie, seed, (long long) n_embd, (long long) n_ff);
        }
        sdq_shutdown();
        llama_free(ctx);
        llama_model_free(model);
        llama_backend_free();
        return rc;
    }

    // ---- 主流程
    if (!opt.prompt.empty()) {
        run_prompt(opt.prompt, true);
        generate(opt.n_predict);
    }

    if (opt.bench) {
        // benchmark：固定提示詞 ×3，量 prefill / decode / 分頁指標
        const std::vector<std::string> prompts = {
            "請用中文簡短介紹 SSD、RAM、CPU 之間的資料分層是什麼，以及為什麼快取很重要。",
            "Write a short paragraph explaining how a mixture-of-experts model routes tokens to experts.",
            "請列出十種日常會用到電腦的工作，並各用一句話說明。",
        };
        for (int round = 0; round < 3; round++) {
            fprintf(stderr, "[bench] round %d\n", round);
            llama_memory_clear(llama_get_memory(ctx), true);
            decoded.clear();
            P.get_stats().tokens.store(0);
            run_prompt(prompts[round % prompts.size()], false);
            generate(48);
            P.dump_stats(round == 0 ? "bench_warm" : "bench_hot");
        }
    }

    if (opt.interactive) {
        printf("\n（輸入文字開始對話；/exit 離開、/reset 重設、/stats 顯示指標）\n");
        std::string line;
        while (true) {
            printf("\nuser> ");
            fflush(stdout);
            if (!std::getline(std::cin, line)) {
                break;
            }
            if (line == "/exit" || line == "/quit") {
                break;
            }
            if (line == "/reset") {
                llama_memory_clear(llama_get_memory(ctx), true);
                decoded.clear();
                st.msgs.clear();
                if (!opt.system.empty()) {
                    st.msgs.push_back({"system", opt.system.c_str()});
                }
                printf("（已重設）\n");
                continue;
            }
            if (line == "/stats") {
                P.dump_stats("manual");
                fprintf(stderr, "hit_rate=%.1f%%  arena=%zu MB  RSS=%zu MB\n",
                        P.hit_rate() * 100, P.used_bytes() / 1048576, rss_mb());
                continue;
            }
            if (line.empty()) {
                continue;
            }
            run_prompt(line, true);
            generate(opt.n_predict);
        }
    }

    P.get_stats().moe_calls.fetch_add(0);
    sdq_shutdown();

    fprintf(stderr, "\n[sdq] ===== 總結 =====\n");
    fprintf(stderr, "[sdq] prefill: %" PRIu64 " tok / %.2fs = %.1f tok/s\n",
            n_prefill_tok, t_prefill, t_prefill > 0 ? n_prefill_tok / t_prefill : 0.0);
    const uint64_t io_end = read_bytes_io();
    fprintf(stderr, "[sdq] 全程 SSD 實際讀取: %.2f GB\n", (io_end - io_read0) / 1073741824.0);
    print_mem_report("結束時");

    if (logits_fp) {
        fclose(logits_fp);
    }
    llama_sampler_free(smpl);
    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
    return 0;
}