// SDQ expert buffer type
//
// 這是一個「假」buffer type：它宣稱接受所有張量，但實際上
//   * 不配置任何 RAM（alloc_buffer 回傳 0 大小的 buffer）
//   * set_tensor 忽略傳入的資料（資料來自 llama.cpp 的 mmap，我們不碰）
//   * 真正的資料由 sdq::pager 從檔案 pread 到自己的 arena
//
// 用途：讓 llama.cpp 把 18.2 GiB 的 MoE 專家權重「掛在模型上」卻不放進 RAM，
// 改由分頁器管理 SSD → RAM → CPU 的三層資料流。

#include "sdq_pager.h"

#include "ggml-backend-impl.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <mutex>
#include <unordered_map>
#include <vector>

namespace sdq {

// 給 tensor->data 的假位址。每個 expert 權重張量分到一小塊（預設 4 KB），
// 指標必須有效（ggml 會做指標運算），但內容永遠不會被讀。
struct expert_buffer {
    std::vector<uint8_t> scratch;
};

static bool g_enabled = false;

void    sdq_set_enabled(bool on) { g_enabled = on; }
bool    sdq_enabled()           { return g_enabled; }

static config g_cfg;
void        sdq_set_config(const config & c) { g_cfg = c; }
const config & sdq_get_config() { return g_cfg; }

// ---------------------------------------------------------------- buffer iface

// ggml_backend_buffer_i 的欄位依賴下面的實作，先給宣告
struct ggml_backend_buffer_i sdq_buffer_iface();

static const char * sdq_buft_get_name(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return "SDQ_ExpertSSD";
}

static size_t sdq_buft_get_alignment(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return 4096;
}

static size_t sdq_buft_get_max_size(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return SIZE_MAX;
}

// 每個 expert 權重張量只回報 4 KB：18.2 GiB 的權重「不會」被載入 RAM 的關鍵。
static size_t sdq_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
    (void) buft;
    (void) tensor;
    return 4096;
}

static bool sdq_buft_is_host(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return false;
}

static ggml_backend_buffer_t sdq_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    // llama.cpp 會用 get_alloc_size() 計算總大小後呼叫本函式；
    // 我們只保留真正需要的假位址空間（每個張量 4 KB）。
    const size_t bytes = size < 4096 ? 4096 : size;
    auto * buf = new expert_buffer();
    buf->scratch.assign(bytes, 0);
    auto * b = new ggml_backend_buffer{
        /* .iface   = */ sdq_buffer_iface(),
        /* .buft    = */ buft,
        /* .context = */ buf,
        /* .size    = */ bytes,
        /* .usage   = */ GGML_BACKEND_BUFFER_USAGE_WEIGHTS,
    };
    return b;
}

// 一次配置整個 context 的張量：每個張量 4 KB 假位址，實際權重 0 RAM。
static ggml_backend_buffer_t sdq_buft_alloc_buffer_n(ggml_backend_buffer_type_t buft,
                                                     ggml_tensor ** tensors, int n_tensors) {
    const size_t slot = 4096;
    const size_t bytes = slot * (size_t) std::max(n_tensors, 1);
    ggml_backend_buffer_t buffer = sdq_buft_alloc_buffer(buft, bytes);
    if (!buffer) {
        return nullptr;
    }
    uint8_t * base = (uint8_t *) buffer->iface.get_base(buffer);
    size_t off = 0;
    for (int i = 0; i < n_tensors; i++) {
        ggml_tensor * t = tensors[i];
        t->data = base + off;
        t->buffer = buffer;
        off += slot;
        buffer->iface.init_tensor(buffer, t);
    }
    return buffer;
}

static void sdq_buft_free_buffer(ggml_backend_buffer_t buffer) {
    // 只釋放自己的資料；ggml_backend_buffer_free() 會負責 delete buffer 本體，
    // 這裡再 delete 一次會造成 double free。
    delete (expert_buffer *) buffer->context;
    buffer->context = nullptr;
}

static void * sdq_buft_get_base(ggml_backend_buffer_t buffer) {
    return ((expert_buffer *) buffer->context)->scratch.data();
}

static ggml_status sdq_buft_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
    (void) buffer;
    (void) tensor;
    return GGML_STATUS_SUCCESS;
}

static void sdq_buft_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
                                   uint8_t value, size_t offset, size_t size) {
    (void) buffer; (void) tensor; (void) value; (void) offset; (void) size;
}

static void sdq_buft_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
                                const void * data, size_t offset, size_t size) {
    // 關鍵：忽略。llama.cpp 會把 mmap 的資料交給我們，這裡一個位元組都不落地，
    // 專家權利永遠不會佔用 RAM。
    (void) buffer; (void) tensor; (void) data; (void) offset; (void) size;
}

static void sdq_buft_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor,
                                void * data, size_t offset, size_t size) {
    (void) buffer; (void) tensor; (void) offset;
    memset(data, 0, size);
}

static bool sdq_buft_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
    (void) buffer; (void) src; (void) dst;
    return true;
}

static void sdq_buft_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    (void) buffer; (void) value;
}

struct ggml_backend_buffer_i sdq_buffer_iface() {
    static ggml_backend_buffer_i i = [] {
        ggml_backend_buffer_i x {};
        x.free_buffer   = sdq_buft_free_buffer;
        x.get_base      = sdq_buft_get_base;
        x.init_tensor   = sdq_buft_init_tensor;
        x.memset_tensor = sdq_buft_memset_tensor;
        x.set_tensor    = sdq_buft_set_tensor;
        x.get_tensor    = sdq_buft_get_tensor;
        x.cpy_tensor    = sdq_buft_cpy_tensor;
        x.clear         = sdq_buft_clear;
        return x;
    }();
    return i;
}

ggml_backend_buffer_type_t sdq_expert_buft() {
    static ggml_backend_buffer_type_t buft = [] {
        return new ggml_backend_buffer_type{
            /* .iface   = */ {
                /* .get_name         = */ sdq_buft_get_name,
                /* .alloc_buffer     = */ sdq_buft_alloc_buffer,
                /* .alloc_buffer_n   = */ sdq_buft_alloc_buffer_n,
                /* .get_alignment    = */ sdq_buft_get_alignment,
                /* .get_max_size     = */ sdq_buft_get_max_size,
                /* .get_alloc_size   = */ sdq_buft_get_alloc_size,
                /* .get_alloc_size_n = */ nullptr,
                /* .is_host          = */ sdq_buft_is_host,
            },
            /* .device  = */ nullptr,
            /* .context = */ nullptr,
        };
    }();
    return buft;
}

// llama.cpp 端需要的介面宣告
void sdq_install_buft_overrides(llama_model_params * params);

static const llama_model_tensor_buft_override * sdq_overrides() {
    static llama_model_tensor_buft_override ov[] = {
        { "blk\\.[0-9]+\\.ffn_(gate|up|down|gate_up)_exps\\.weight", sdq_expert_buft() },
        { nullptr, nullptr },
    };
    return ov;
}

void sdq_install_buft_overrides(llama_model_params * params) {
    if (!params) {
        return;
    }
    if (params->tensor_buft_overrides == nullptr) {
        params->tensor_buft_overrides = sdq_overrides();
    } else {
        // 使用者已有 override：把手動指定的順序原樣保留，附加在尾端
        size_t n = 0;
        while (params->tensor_buft_overrides[n].pattern != nullptr) {
            n++;
        }
        static std::vector<llama_model_tensor_buft_override> merged;
        merged.clear();
        for (size_t i = 0; i < n; i++) {
            merged.push_back(params->tensor_buft_overrides[i]);
        }
        for (size_t i = 0; sdq_overrides()[i].pattern != nullptr; i++) {
            merged.push_back(sdq_overrides()[i]);
        }
        merged.push_back({ nullptr, nullptr });
        params->tensor_buft_overrides = merged.data();
    }
}

} // namespace sdq