/* qxmx_cpu: see qxmx_cpu.h.  All real work is in qxmx_ref (ref_decode). */
#include "qxmx_cpu.h"

#include <cmath>

namespace qx {

bool cpu_engine::init(qx_model* m, int max_ctx) {
    this->m = m;
    s.init(m, max_ctx);
    logits_buf.assign(QX_N_VOCAB, 0.0f);
    x_buf.assign(QX_D_MODEL, 0.0f);
    logits = logits_buf.data();
    return true;
}

void cpu_engine::reset(int /*slot*/) {
    s.init(m, s.max_ctx);   /* re-init zeroes SSM/KV state and n_tokens */
}

void cpu_engine::forward(int /*slot*/, int token_id) {
    ref_dequant_row(&m->token_embd, token_id, x_buf.data());
    ref_decode(m, &s, x_buf.data(), logits);
}

void cpu_engine::prefill(int /*slot*/, const int* token_ids, int n) {
    for (int i = 0; i < n; i++) {
        ref_dequant_row(&m->token_embd, token_ids[i], x_buf.data());
        ref_decode(m, &s, x_buf.data(), logits);
    }
}

void cpu_engine::prefill_step(int /*slot*/, const int* token_ids, int n,
                              int& i, int& /*last_chunk_m*/, bool& done) {
    /* CPU path: one forward() per step.  last_chunk_m is unused (no chunking
     * here); the GPU engine's _b path is the only place chunking matters. */
    if (i >= n) { done = true; return; }
    ref_dequant_row(&m->token_embd, token_ids[i], x_buf.data());
    ref_decode(m, &s, x_buf.data(), logits);
    i++;
    done = (i >= n);
}

int64_t cpu_engine::slot_n_tokens(int /*slot*/) const {
    return s.n_tokens;
}

}  // namespace qx