#include <float.h>
#include <math.h>
#include <string.h>

#include "bio_forward.h"
#include "bio_quant.h"

static const bio_gguf_tensor_span *select_projection_tensor(const bio_gguf_model *model) {
    const bio_gguf_tensor_span *t = bio_gguf_find_tensor(model, "output.weight");
    if (t != NULL) {
        return t;
    }
    t = bio_gguf_find_tensor(model, "token_embd.weight");
    if (t != NULL) {
        return t;
    }
    if (model != NULL) {
        for (uint32_t i = 0; i < model->registered_tensor_count; ++i) {
            if (model->tensors[i].int4_compatible != 0u) {
                return &model->tensors[i];
            }
        }
    }
    return NULL;
}

static uint32_t sample_argmax(const float *logits, uint32_t n, uint64_t seed,
                              float *score_out) {
    uint32_t best = 0;
    float best_score = -FLT_MAX;
    for (uint32_t i = 0; i < n; ++i) {
        const float jitter = (float)((seed >> (i & 15u)) & 0xffu) / 65536.0f;
        const float s = logits[i] + jitter;
        if (s > best_score) {
            best_score = s;
            best = i;
        }
    }
    if (score_out != NULL) {
        *score_out = best_score;
    }
    return best;
}

bio_v1_status bio_forward_next_token(const bio_gguf_model *model,
                                     const uint32_t *tokens,
                                     uint32_t token_count,
                                     uint64_t seed,
                                     bio_forward_result *out) {
    float logits[BIO_FORWARD_LOGITS];
    const bio_gguf_tensor_span *tensor = select_projection_tensor(model);
    if (model == NULL || tokens == NULL || token_count == 0u || out == NULL ||
        tensor == NULL) {
        return BIO_ERR_ARGS;
    }
    memset(logits, 0, sizeof(logits));
    if (bio_quant_project_avx2(tensor, tokens, token_count, logits,
                               BIO_FORWARD_LOGITS) != BIO_OK) {
        return BIO_ERR_INTERNAL;
    }
    out->token = sample_argmax(logits, BIO_FORWARD_LOGITS, seed, &out->score);
    out->executed_ops = (uint64_t)token_count * BIO_FORWARD_LOGITS;
    out->baseline_ops = out->executed_ops;
    return BIO_OK;
}
