#pragma once

#include "strata/spec/suffix_drafter.hpp"

#include <algorithm>
#include <cstdint>

namespace strata::spec {

inline bool suffix_stop(int32_t token) {
    return token == 248046 || token == 248044;
}

inline int suffix_limit(int requested, int remaining, int64_t pos, int64_t max_cells, bool unrestricted_qsa = true) {
    const int block_drafts = unrestricted_qsa ? 3 : (3 - int(pos & 3));
    const int64_t room = std::max<int64_t>(0, max_cells - pos - 3);
    const int64_t drafts = std::min<int64_t>({requested, remaining - 1, block_drafts, room});
    return int(std::max<int64_t>(0, drafts));
}

inline int suffix_keep(const int32_t* input, const int32_t* output, int nd, bool stop_on_eos = true) {
    int keep = 1;
    for (int i = 0; i < nd; ++i) {
        if ((stop_on_eos && suffix_stop(output[i])) || output[i] != input[i + 1]) break;
        ++keep;
    }
    return keep;
}

}  // namespace strata::spec