#include "utils.h"
#include "common.h"
#include "llama.h"

#include <cmath>
#include <cstring>
#include <fstream>
#include <regex>
#include <locale>
#include <codecvt>
#include <sstream>
#include <ctime>

#define MINIAUDIO_IMPLEMENTATION
#ifndef MTMD_AUDIO_DEBUG
#   define MA_NO_ENCODING
#endif
#define MA_NO_DEVICE_IO
#define MA_NO_RESOURCE_MANAGER
#define MA_NO_NODE_GRAPH
#define MA_NO_ENGINE
#define MA_NO_GENERATION
// #define MA_API static
#include "miniaudio/miniaudio.h"

#include "acestep/mp3/mp3enc.h"

void utreplace(std::string & str, const std::string & needle, const std::string & replacement) {
    size_t pos = 0;
    while ((pos = str.find(needle, pos)) != std::string::npos) {
        str.replace(pos, needle.length(), replacement);
        pos += replacement.length();
    }
}

std::map<std::string, int32_t> json_parse(const std::string & fname) {
    std::map<std::string, int32_t> result;

    // read file into string
    std::string json;
    {
        std::ifstream ifs(fname);
        if (!ifs) {
            fprintf(stderr, "Failed to open %s\n", fname.c_str());
            exit(1);
        }

        json = std::string((std::istreambuf_iterator<char>(ifs)),
                (std::istreambuf_iterator<char>()));
    }

    if (json[0] != '{') {
        return result;
    }

    // parse json
    {
        bool has_key  = false;
        bool in_token = false;

        std::string str_key = "";
        std::string str_val = "";

        int n = json.size();
        for (int i = 1; i < n; ++i) {
            if (!in_token) {
                if (json[i] == ' ') continue;
                if (json[i] == '"') {
                    in_token = true;
                    continue;
                }
            } else {
                if (json[i] == '\\' && i+1 < n) {
                    if (has_key == false) {
                        str_key += json[i];
                    } else {
                        str_val += json[i];
                    }
                    ++i;
                } else if (json[i] == '"') {
                    if (has_key == false) {
                        has_key = true;
                        ++i;
                        while (json[i] == ' ') ++i;
                        ++i; // :
                        while (json[i] == ' ') ++i;
                        if (json[i] != '\"') {
                            while (json[i] != ',' && json[i] != '}') {
                                str_val += json[i++];
                            }
                            has_key = false;
                        } else {
                            in_token = true;
                            continue;
                        }
                    } else {
                        has_key = false;
                    }

                    ::utreplace(str_key, "\\u0120", " " ); // \u0120 -> space
                    ::utreplace(str_key, "\\u010a", "\n"); // \u010a -> new line
                    ::utreplace(str_key, "\\\"",    "\""); // \\\"   -> "

                    try {
                        result[str_key] = std::stoi(str_val);
                    } catch (...) {
                        //fprintf(stderr, "%s: ignoring key '%s' with value '%s'\n", fname.c_str(), str_key.c_str(), str_val.c_str());

                    }
                    str_key = "";
                    str_val = "";
                    in_token = false;
                    continue;
                }
                if (has_key == false) {
                    str_key += json[i];
                } else {
                    str_val += json[i];
                }
            }
        }
    }

    return result;
}


void gpt_vocab::add_special_token(const std::string & token) {
    special_tokens.push_back(token);
}


std::string convert_to_utf8(const std::wstring & input) {
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    return converter.to_bytes(input);
}


std::wstring convert_to_wstring(const std::string & input) {
    try {
        std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
        return converter.from_bytes(input);
    } catch (const std::range_error& e) {
        return L"";
    } catch (...) {
        return L"";
    }
}

void gpt_split_words(std::string str, std::vector<std::string>& words) {
    const std::string pattern = R"('s|'t|'re|'ve|'m|'ll|'d| ?[[:alpha:]]+| ?[[:digit:]]+| ?[^\s[:alpha:][:digit:]]+|\s+(?!\S)|\s+)";
    const std::regex re(pattern);
    std::smatch m;

    while (std::regex_search(str, m, re)) {
        for (auto x : m) {
            words.push_back(x);
        }
        str = m.suffix();
    }
}

std::vector<gpt_vocab::id> gpt_tokenize(const gpt_vocab & vocab, const std::string & text) {
    std::vector<std::string> words;

    // first split the text into words
    {
        std::string str = text;

        // Generate the subpattern from the special_tokens vector if it's not empty
        if (!vocab.special_tokens.empty()) {
            const std::regex escape(R"([\[\\\^\$\.\|\?\*\+\(\)\{\}])");
            std::string special_tokens_subpattern;
            for (const auto & token : vocab.special_tokens) {
                if (!special_tokens_subpattern.empty()) {
                    special_tokens_subpattern += "|";
                }
                special_tokens_subpattern += std::regex_replace(token, escape, R"(\$&)");
            }

            std::regex re(special_tokens_subpattern);
            std::smatch m;
            // Split the text by special tokens.
            while (std::regex_search(str, m, re)) {
                // Split the substrings in-between special tokens into words.
                gpt_split_words(m.prefix(), words);
                // Add matched special tokens as words.
                for (auto x : m) {
                    words.push_back(x);
                }
                str = m.suffix();
            }
            // Remaining text without special tokens will be handled below.
        }

        gpt_split_words(str, words);
    }

    // find the longest token that forms each word in words:
    std::vector<gpt_vocab::id> tokens;
    for (const auto & word : words) {
        for (int i = 0; i < word.size(); ){
            for (int j = word.size() - 1; j >= i; j--){
                auto cand = word.substr(i, j-i+1);
                auto it = vocab.token_to_id.find(cand);
                if (it != vocab.token_to_id.end()){ // word.substr(i, j-i+1) in vocab
                    tokens.push_back(it->second);
                    i = j + 1;
                    break;
                }
                else if (j == i){ // word.substr(i, 1) has no matching
                    fprintf(stderr, "%s: unknown token '%s'\n", __func__, word.substr(i, 1).data());
                    i++;
                }
            }
        }
    }


    return tokens;
}

bool should_transpose_layer(std::string name)
{

    if(name.find(".mlp.fc_in.weight")!=std::string::npos ||
    name.find(".attn.out_proj.weight")!=std::string::npos ||
    name.find(".attn.q_proj.weight")!=std::string::npos ||
    name.find(".attn.k_proj.weight")!=std::string::npos ||
    name.find(".attn.v_proj.weight")!=std::string::npos ||
    name.find("/attn/c_attn/w")!=std::string::npos ||
    name.find("/attn/c_proj/w")!=std::string::npos ||
    name.find("/mlp/c_fc/w")!=std::string::npos ||
    name.find("/mlp/c_proj/w")!=std::string::npos)
    {
        return true;
    }
    return false;
}

static const std::string kcpp_base64_chars =
             "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
             "abcdefghijklmnopqrstuvwxyz"
             "0123456789+/";
static inline bool kcpp_is_base64(uint8_t c)
{
    return (isalnum(c) || (c == '+') || (c == '/'));
}
std::vector<uint8_t> kcpp_base64_decode(const std::string & encoded_string)
{
    int i = 0;
    int j = 0;
    int in_ = 0;

    int in_len = encoded_string.size();

    uint8_t char_array_4[4];
    uint8_t char_array_3[3];

    std::vector<uint8_t> ret;

    while (in_len-- && (encoded_string[in_] != '=') && kcpp_is_base64(encoded_string[in_]))
    {
        char_array_4[i++] = encoded_string[in_]; in_++;
        if (i == 4)
        {
            for (i = 0; i <4; i++)
            {
                char_array_4[i] = kcpp_base64_chars.find(char_array_4[i]);
            }

            char_array_3[0] = ((char_array_4[0]      ) << 2) + ((char_array_4[1] & 0x30) >> 4);
            char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
            char_array_3[2] = ((char_array_4[2] & 0x3) << 6) +   char_array_4[3];

            for (i = 0; (i < 3); i++)
            {
                ret.push_back(char_array_3[i]);
            }
            i = 0;
        }
    }

    if (i)
    {
        for (j = i; j <4; j++)
        {
            char_array_4[j] = 0;
        }

        for (j = 0; j <4; j++)
        {
            char_array_4[j] = kcpp_base64_chars.find(char_array_4[j]);
        }

        char_array_3[0] = ((char_array_4[0]      ) << 2) + ((char_array_4[1] & 0x30) >> 4);
        char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
        char_array_3[2] = ((char_array_4[2] & 0x3) << 6) +   char_array_4[3];

        for (j = 0; (j < i - 1); j++)
        {
            ret.push_back(char_array_3[j]);
        }
    }

    return ret;
}
std::string kcpp_base64_encode(const unsigned char* data, unsigned int data_length) {
    const std::string base64_chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::string encoded;
    encoded.reserve(((data_length + 2) / 3) * 4);
    for (unsigned int i = 0; i < data_length; i += 3) {
        unsigned int triple = (data[i] << 16) + (i + 1 < data_length ? data[i + 1] << 8 : 0) + (i + 2 < data_length ? data[i + 2] : 0);
        encoded.push_back(base64_chars[(triple >> 18) & 0x3F]);
        encoded.push_back(base64_chars[(triple >> 12) & 0x3F]);
        if (i + 1 < data_length) {
            encoded.push_back(base64_chars[(triple >> 6) & 0x3F]);
        } else {
            encoded.push_back('=');
        }
        if (i + 2 < data_length) {
            encoded.push_back(base64_chars[triple & 0x3F]);
        } else {
            encoded.push_back('=');
        }
    }
    return encoded;
}
std::string kcpp_base64_encode(const std::string &data) {
    static const char lookup[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::string encoded;
    int val = 0, valb = -6;
    for (unsigned char c : data) {
        val = (val << 8) + c;
        valb += 8;
        while (valb >= 0) {
            encoded.push_back(lookup[(val >> valb) & 0x3F]);
            valb -= 6;
        }
    }
    if (valb > -6) {
        encoded.push_back(lookup[((val << 8) >> (valb + 8)) & 0x3F]);
    }
    while (encoded.size() % 4) {
        encoded.push_back('=');
    }
    return encoded;
}

std::string get_timestamp_str()
{
    std::time_t t = std::time(nullptr);
    std::tm* now = std::localtime(&t);
    char buffer[16]; // Buffer to hold "hh:mm:ss" and null terminator
    std::sprintf(buffer, "%02d:%02d:%02d", now->tm_hour, now->tm_min, now->tm_sec);
    // Convert the buffer to a std::string
    std::string timestamp(buffer);
    return timestamp;
}

//split a big vector into multiple small vectors of chunk size or less
std::vector<std::vector<int>> split_big_vector(const std::vector<int>& big_arr, size_t chunk_size) {
    std::vector<std::vector<int>> small_arrs;
    for (size_t i = 0; i < big_arr.size(); i += chunk_size) {
        size_t end = std::min(i + chunk_size, big_arr.size());
        small_arrs.emplace_back(big_arr.begin() + i, big_arr.begin() + end);
    }
    return small_arrs;
}

std::vector<std::vector<int>> split_big_vector_in_two(const std::vector<int>& big_arr, size_t chunk_size)
{
    std::vector<std::vector<int>> result;
    if (chunk_size == 0 || big_arr.empty())
        return result;

    if (big_arr.size() <= chunk_size) {
        // Only one chunk (all elements)
        result.emplace_back(big_arr);
        return result;
    }
    size_t split_point = big_arr.size() - chunk_size;
    result.emplace_back(big_arr.begin(), big_arr.begin() + split_point);  // First big chunk
    result.emplace_back(big_arr.begin() + split_point, big_arr.end()); // Last chunk (size <= chunk_size)
    return result;
}

static double audio_resample_bessel_i0(double x) {
    double sum  = 1.0;
    double term = 1.0;
    double y    = x * x * 0.25;
    for (int k = 1; k < 30; k++) {
        term *= y / ((double) k * (double) k);
        sum += term;
        if (term < sum * 1e-15) {
            break;
        }
    }
    return sum;
}

std::vector<float> resample_wav(int num_channels,const std::vector<float>& input,uint32_t input_rate,uint32_t output_rate)
{
    if (input.empty() || num_channels <= 0 || input_rate == 0 || output_rate == 0)
        return {};

    if (input.size() % num_channels != 0)
        return {};

    const int n_in = input.size() / num_channels;

    if (input_rate == output_rate)
        return input;

    const double ratio = (double)output_rate / (double)input_rate;
    const int n_out = (int)std::lround(n_in * ratio);

    std::vector<float> output((size_t)n_out * num_channels);

    const int half_len = 32;
    const int taps = half_len * 2;

    const double beta = 9.0;

    const double inv_i0b = 1.0 / audio_resample_bessel_i0(beta);
    const double fc = 0.5 * ((ratio < 1.0) ? ratio : 1.0);

    // PRECOMPUTE KAISER WINDOW
    std::vector<double> window(taps + 1);

    for (int k = -half_len; k <= half_len; k++)
    {
        double t = (double)k / (double)half_len;

        double win;

        if (t < -1.0 || t > 1.0)
            win = 0.0;
        else
            win = audio_resample_bessel_i0(beta * std::sqrt(1.0 - t * t)) * inv_i0b;

        window[k + half_len] = win;
    }

    for (int ch = 0; ch < num_channels; ch++)
    {
        const float* src = input.data() + ch * n_in;
        float* dst = output.data() + ch * n_out;

        for (int i = 0; i < n_out; i++)
        {
            double center = (double)i / ratio;

            int base = (int)std::floor(center);

            int start = base - half_len + 1;
            int end   = base + half_len;

            double sum = 0.0;
            double wgt = 0.0;

            for (int j = start; j <= end; j++)
            {
                double d = center - (double)j;

                double sinc_val;

                if (std::fabs(d) < 1e-9)
                    sinc_val = 2.0 * fc;
                else
                    sinc_val = std::sin(2.0 * M_PI * fc * d) / (M_PI * d);

                double win = window[j - start];

                double h = sinc_val * win;

                int idx = j;

                if (idx < 0) idx = 0;
                if (idx >= n_in) idx = n_in - 1;

                sum += src[idx] * h;
                wgt += h;
            }

            dst[i] = (wgt > 1e-12) ? (float)(sum / wgt) : 0.0f;
        }
    }

    return output;
}

std::vector<float> mix_planar_stereo_to_mono(const float* audio, int T_audio)
{
    std::vector<float> mono(T_audio);
    const float* left  = audio;
    const float* right = audio + T_audio;
    for (int t = 0; t < T_audio; ++t)
    {
        mono[t] = 0.5f * (left[t] + right[t]);
    }
    return mono;
}

static uint8_t linear_to_mulaw(int16_t sample)
{
    const int16_t BIAS = 0x84;        // 132
    const int16_t CLIP = 32635;

    int16_t sign = (sample >> 8) & 0x80;
    if (sign)
        sample = -sample;

    if (sample > CLIP)
        sample = CLIP;

    sample += BIAS;

    int16_t exponent = 7;
    for (int16_t expMask = 0x4000;
         (sample & expMask) == 0 && exponent > 0;
         exponent--, expMask >>= 1);

    int16_t mantissa = (sample >> (exponent + 3)) & 0x0F;

    uint8_t ulaw = ~(sign | (exponent << 4) | mantissa);
    return ulaw;
}

std::string save_ulaw_wav8_base64(const std::vector<float> &data, int sample_rate)
{
    std::ostringstream oss;
    wav_ulaw_header header;

    header.sample_rate = sample_rate;
    header.byte_rate   = sample_rate;      // 1 byte per sample (mono)
    header.block_align = 1;
    header.data_size   = static_cast<uint32_t>(data.size());
    header.chunk_size  = 4                       // "WAVE"
                       + 8 + header.fmt_chunk_size
                       + 8 + header.data_size;

    // Write header
    oss.write(reinterpret_cast<const char*>(&header), sizeof(header));

    // Convert and write samples
    for (float s : data)
    {
        float clamped = std::clamp(s, -1.0f, 1.0f);
        int16_t pcm = static_cast<int16_t>(clamped * 32767.0f);
        uint8_t mu = linear_to_mulaw(pcm);
        oss.write(reinterpret_cast<const char*>(&mu), 1);
    }

    std::string wav_data = oss.str();
    return kcpp_base64_encode(wav_data);
}

std::string save_wav16_base64(const std::vector<float> &data, int sample_rate) {
    std::ostringstream oss;
    wav16_header header;

    // Fill header fields
    header.sample_rate = sample_rate;
    header.byte_rate = header.sample_rate * header.num_channels * (header.bits_per_sample / 8);
    header.block_align = header.num_channels * (header.bits_per_sample / 8);
    header.data_size = data.size() * (header.bits_per_sample / 8);
    header.chunk_size = 36 + header.data_size;

    // Write header
    oss.write(reinterpret_cast<const char*>(&header), sizeof(header));

    // Write samples
    for (const auto &sample : data) {
        int16_t pcm_sample = static_cast<int16_t>(std::clamp(sample * 32767.0, -32768.0, 32767.0));
        oss.write(reinterpret_cast<const char*>(&pcm_sample), sizeof(pcm_sample));
    }

    // Get binary WAV data
    std::string wav_data = oss.str();
    return kcpp_base64_encode(wav_data); //return as base64 string
}

//assumes planar stereo input from acestep
std::string save_stereo_wav16_base64(const std::vector<float> & raw_audio, int T_audio, int sample_rate) {
    std::ostringstream oss(std::ios::binary);
    const int n_channels = 2;
    const int bits = 16;
    const int byte_rate = sample_rate * n_channels * (bits / 8);
    const int block_align = n_channels * (bits / 8);
    const int data_size = T_audio * n_channels * (bits / 8);
    const int file_size = 36 + data_size;
    oss.write("RIFF", 4);
    oss.write(reinterpret_cast<const char*>(&file_size), 4);
    oss.write("WAVE", 4);
    oss.write("fmt ", 4);
    int32_t fmt_size = 16;
    oss.write(reinterpret_cast<const char*>(&fmt_size), 4);
    int16_t audio_fmt = 1; // PCM
    oss.write(reinterpret_cast<const char*>(&audio_fmt), 2);
    int16_t nc = n_channels;
    oss.write(reinterpret_cast<const char*>(&nc), 2);
    oss.write(reinterpret_cast<const char*>(&sample_rate), 4);
    oss.write(reinterpret_cast<const char*>(&byte_rate), 4);
    int16_t ba = block_align;
    oss.write(reinterpret_cast<const char*>(&ba), 2);
    int16_t bp = bits;
    oss.write(reinterpret_cast<const char*>(&bp), 2);
    oss.write("data", 4);
    oss.write(reinterpret_cast<const char*>(&data_size), 4);

    // EXPECTS PLANAR INPUT:
    // raw_audio[0 ... T_audio-1]           = Left
    // raw_audio[T_audio ... 2*T_audio-1]   = Right
    for (int t = 0; t < T_audio; ++t) {
        for (int c = 0; c < 2; ++c) {
            float s = raw_audio[c * T_audio + t];
            s = std::max(-1.0f, std::min(1.0f, s));  // clamp to [-1, 1]
            int16_t v = static_cast<int16_t>(s * 32767.0f);
            oss.write(reinterpret_cast<const char*>(&v), 2);
        }
    }
    std::string wav_data = oss.str();
    return kcpp_base64_encode(wav_data);
}

std::string save_stereo_mp3_base64(const std::vector<float> & raw_audio,int T_audio,int sample_rate) {
    const float * enc_audio = raw_audio.data();
    int enc_T  = T_audio;
    int enc_sr = sample_rate;
    std::vector<float> resampled;

    // resample to 44100 if sr is not a valid MPEG1 rate
    if (sample_rate != 32000 && sample_rate != 44100 && sample_rate != 48000) {
        resampled = resample_wav(2,raw_audio,sample_rate,44100);
        enc_audio = resampled.data();
        enc_sr    = 44100;
    }

    const int kbps = 128;
    mp3enc_t * enc = mp3enc_init(enc_sr, 2, kbps);
    if (!enc) {
        fprintf(stderr, "[Audio] mp3enc_init failed\n");
        return "";
    }

    std::string mp3_data;
    mp3_data.reserve(enc_T); // rough preallocation

    int chunk = enc_sr;

    // reusable buffer (replaces malloc inside loop)
    std::vector<float> buf((size_t)chunk * 2);

    for (int pos = 0; pos < enc_T; pos += chunk) {
        int n = (pos + chunk <= enc_T) ? chunk : (enc_T - pos);
        // build planar chunk
        memcpy(buf.data(),     enc_audio + pos,          (size_t)n * sizeof(float));
        memcpy(buf.data() + n, enc_audio + enc_T + pos,  (size_t)n * sizeof(float));
        int out_size = 0;
        const uint8_t * mp3 = mp3enc_encode(enc, buf.data(), n, &out_size);
        if (out_size > 0) {
            mp3_data.append((const char *) mp3, out_size);
        }
    }
    int flush_size = 0;
    const uint8_t * flush_data = mp3enc_flush(enc, &flush_size);
    if (flush_size > 0) {
        mp3_data.append((const char *) flush_data, flush_size);
    }
    mp3enc_free(enc);
    return kcpp_base64_encode(mp3_data);
}

//a very rudimentary all in one sampling function which has no dependencies
int32_t kcpp_quick_sample(float * logits, const int n_logits, const std::vector<int32_t> & last_n_tokens, float rep_pen, float top_p, int top_k, float temp, std::mt19937 & rng)
{
    if (temp <= 0) {
        // select the token with the highest logit directly
        float max_logit = logits[0];
        int32_t max_id = 0;
        for (int i = 1; i < n_logits; ++i) {
            if (logits[i] > max_logit) {
                max_logit = logits[i];
                max_id = i;
            }
        }
        return max_id;
    }

    top_k = (top_k<=0 || top_k>300)?300:top_k;
    top_k = std::min(top_k, n_logits);

    std::vector<std::pair<float, int32_t>> logits_id;
    logits_id.reserve(n_logits);

    //temperature sample
    const float scale = 1.0f/temp;

    //sample rep pen
    for (int i = 0; i < n_logits; ++i) {
        if (rep_pen>1.0f && std::find(last_n_tokens.begin(), last_n_tokens.end(), i) != last_n_tokens.end()) {
            // if score < 0 then repetition penalty has to multiplied to reduce the previous token probability
            if (logits[i] < 0.0f) {
                logits_id.push_back(std::make_pair((logits[i]*scale)*rep_pen, i));
            } else {
                logits_id.push_back(std::make_pair((logits[i]*scale)/rep_pen, i));
            }
        } else {
            logits_id.push_back(std::make_pair(logits[i]*scale, i));
        }
    }

    //sample top_k
    std::partial_sort(
            logits_id.begin(),
            logits_id.begin() + top_k, logits_id.end(),
            [](const std::pair<float, int32_t> & a, const std::pair<float, int32_t> & b) {
        return a.first > b.first;
    });
    logits_id.resize(top_k);

    // compute probs for the top k tokens
    std::vector<float> probs;
    probs.reserve(logits_id.size());
    float maxl = logits_id[0].first;
    double sum = 0.0;
    for (const auto & kv : logits_id) {
        const float p = expf(kv.first - maxl);
        probs.push_back(p);
        sum += p;
    }

    // normalize the probs
    for (auto & p : probs) {
        p /= sum;
    }

    //apply top p
    if (top_p < 1.0) {
        double cumsum = 0.0;
        for (int i = 0; i < (int) probs.size(); i++) {
            cumsum += probs[i];
            if (cumsum >= top_p) {
                probs.resize(i + 1);
                logits_id.resize(i + 1);
                break;
            }
        }
    }

    // normalize the probs
    for (auto & p : probs) {
        p /= sum;
    }

    std::discrete_distribution<> dist(probs.begin(), probs.end());
    int idx = dist(rng);

    return logits_id[idx].second;
}


std::vector<std::string> split_string(const std::string& input, const std::string& separator) {
    std::vector<std::string> result;
    size_t start = 0;
    size_t end = input.find(separator);

    while (end != std::string::npos) {
        result.push_back(input.substr(start, end - start));
        start = end + separator.length();
        end = input.find(separator, start);
    }

    // Add the remaining part after the last separator
    result.push_back(input.substr(start));

    return result;
}


static bool buf_is_audio_file(const char * buf, size_t len) {
    if (len < 12) {
        return false;
    }

    // RIFF ref: https://en.wikipedia.org/wiki/Resource_Interchange_File_Format
    // WAV ref: https://www.mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html
    bool is_wav = memcmp(buf, "RIFF", 4) == 0 && memcmp(buf + 8, "WAVE", 4) == 0;
    bool is_mp3 = len >= 3 && (
        memcmp(buf, "ID3", 3) == 0 ||
        // Check for MPEG sync word (simplified check)
        ((unsigned char)buf[0] == 0xFF && ((unsigned char)buf[1] & 0xE0) == 0xE0)
    );
    bool is_flac = memcmp(buf, "fLaC", 4) == 0;

    return is_wav || is_mp3 || is_flac;
}

// returns true if the buffer is a valid audio file
bool kcpp_decode_audio_from_buf(const unsigned char * buf_in, size_t len, int target_sampler_rate, std::vector<float> & pcmf32_mono) {
    if (!buf_is_audio_file((const char *)buf_in, len))
    {
        return false;
    }

    ma_result result;
    const int channels = 1;
    ma_decoder_config decoder_config = ma_decoder_config_init(ma_format_f32, channels, target_sampler_rate);
    ma_decoder decoder;

    result = ma_decoder_init_memory(buf_in, len, &decoder_config, &decoder);
    if (result != MA_SUCCESS) {
        return false;
    }

    ma_uint64 frame_count;
    ma_uint64 frames_read;
    result = ma_decoder_get_length_in_pcm_frames(&decoder, &frame_count);
    if (result != MA_SUCCESS) {
        ma_decoder_uninit(&decoder);
        return false;
    }

    pcmf32_mono.resize(frame_count);
    result = ma_decoder_read_pcm_frames(&decoder, pcmf32_mono.data(), frame_count, &frames_read);
    if (result != MA_SUCCESS) {
        ma_decoder_uninit(&decoder);
        return false;
    }

    ma_decoder_uninit(&decoder);
    return true;
}

//this version is specifically required for ace-step
bool kcpp_decode_audio_to_f32_stereo_48k(const uint8_t * data, size_t data_size, std::vector<float> & pcm, int & T_audio) {
    ma_result result;

    // Force the exact format expected by the VAE
    ma_decoder_config config =
        ma_decoder_config_init(ma_format_f32, 2, 48000);

    ma_decoder decoder;

    result = ma_decoder_init_memory(data, data_size, &config, &decoder);
    if (result != MA_SUCCESS)
        return false;

    ma_uint64 frame_count = 0;

    result = ma_decoder_get_length_in_pcm_frames(&decoder, &frame_count);
    if (result != MA_SUCCESS) {
        ma_decoder_uninit(&decoder);
        return false;
    }

    // allocate stereo
    pcm.resize(frame_count * 2);

    ma_uint64 frames_read = 0;

    result = ma_decoder_read_pcm_frames(
        &decoder,
        pcm.data(),
        frame_count,
        &frames_read
    );

    ma_decoder_uninit(&decoder);

    if (result != MA_SUCCESS)
        return false;

    pcm.resize(frames_read * 2);
    T_audio = (int)frames_read;

    return true;
}

static std::vector<std::string> kcpp_string_split(const std::string & input, char separator)
{
    std::vector<std::string> parts;
    size_t begin_pos = 0;
    size_t separator_pos = input.find(separator);
    while (separator_pos != std::string::npos) {
        std::string part = input.substr(begin_pos, separator_pos - begin_pos);
        parts.emplace_back(part);
        begin_pos = separator_pos + 1;
        separator_pos = input.find(separator, begin_pos);
    }
    parts.emplace_back(input.substr(begin_pos, separator_pos - begin_pos));
    return parts;
}

//for llama.cpp style device overrides e.g. --device Vulkan0,Vulkan1
std::vector<ggml_backend_dev_t> kcpp_parse_device_list(const std::string & value) {
    std::vector<ggml_backend_dev_t> devices;
    auto dev_names = kcpp_string_split(value, ',');
    if (dev_names.empty()) {
        printf("\nkcpp_parse_device_list error: no devices specified\n");
        return std::vector<ggml_backend_dev_t>();
    }
    if (dev_names.size() == 1 && dev_names[0] == "none") {
        return std::vector<ggml_backend_dev_t>();
    } else {
        for (const auto & device : dev_names) {
            auto * dev = ggml_backend_dev_by_name(device.c_str());
            if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
                printf("\nkcpp_parse_device_list error: invalid device: %s\n",device.c_str());
                return std::vector<ggml_backend_dev_t>();
            }
            devices.push_back(dev);
        }
        devices.push_back(nullptr);
    }
    return devices;
}

bool kcpp_string_ends_with(const std::string& str, const std::string& suffix) {
    return str.size() >= suffix.size() &&
           str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0;
}

std::string kcpp_rstrip(const std::string& s) {
    size_t end = s.find_last_not_of(" \t\n\r\f\v");
    return (end == std::string::npos) ? "" : s.substr(0, end + 1);
}

//counts the number of matching prefix tokens between two sequences
int ComputeSharedPrefixLength(const std::vector<int> &tokens_a,const std::vector<int> &tokens_b)
{
    size_t min_length = std::min(tokens_a.size(), tokens_b.size());

    int match_count = 0;
    for (size_t i = 0; i < min_length; ++i) {
        if (tokens_a[i] != tokens_b[i]) {
            break;
        }
        match_count++;
    }

    return match_count;
}

//counts the number of matching prefix tokens between two sequences, returns percentage matched 0.0 to 1.0
float ComputePrefixMatchPercent(const std::vector<int> &tokens_a,const std::vector<int> &tokens_b)
{
    size_t min_length = std::min(tokens_a.size(), tokens_b.size());

    if (min_length == 0) {
        return 0.0f;
    }

    int match_count = ComputeSharedPrefixLength(tokens_a, tokens_b);
    return static_cast<float>(match_count) / static_cast<float>(min_length);
}

//returns true if and only if sequence 1 is fully contained within the starting of sequence 2
bool FullyContainedPrefix(std::vector<int> &sequence1, std::vector<int> &sequence2)
{
    if (sequence1.size() > sequence2.size() || sequence1.size()==0 || sequence2.size()==0) {
        return false;
    }
    for (size_t i = 0; i < sequence1.size(); ++i) {
        if (sequence1[i] != sequence2[i]) {
            return false;
        }
    }
    return true;
}