#include "llama-longhaul.h"

#include "llama-impl.h"

#include <algorithm>
#include <cstring>
#include <inttypes.h>

llama_longhaul_cache::llama_longhaul_cache(
        llama_files files,
        std::vector<llama_model_loader::longhaul_source> sources,
        size_t n_slots,
        uint32_t n_experts,
        uint32_t n_layers) :
    files(std::move(files)),
    sources(std::move(sources)),
    layers(n_layers),
    n_slots(n_slots),
    n_experts(n_experts) {
    for (auto & file : this->files) {
        file->set_no_cache();
    }
    for (auto & layer : layers) {
        layer.expert_ids.assign(n_slots, -1);
        layer.last_used.assign(n_slots, 0);
    }
}

llama_longhaul_cache::~llama_longhaul_cache() {
    LLAMA_LOG_INFO("%s: hits = %" PRIu64 ", misses = %" PRIu64 ", read = %.2f MiB\n",
        __func__, n_hits, n_misses, bytes_read / 1024.0 / 1024.0);
}

int llama_longhaul_cache::find_slot(int layer, int32_t expert_id) {
    auto & state = layers.at(layer);
    for (size_t i = 0; i < n_slots; ++i) {
        if (state.expert_ids[i] == expert_id) {
            state.last_used[i] = ++tick;
            ++n_hits;
            return i;
        }
    }

    size_t slot = 0;
    for (size_t i = 0; i < n_slots; ++i) {
        if (state.expert_ids[i] == -1) {
            slot = i;
            break;
        }
        if (state.last_used[i] < state.last_used[slot]) {
            slot = i;
        }
    }
    if (!load_slot(layer, slot, expert_id)) {
        return -1;
    }
    state.expert_ids[slot] = expert_id;
    state.last_used[slot] = ++tick;
    ++n_misses;
    return slot;
}

bool llama_longhaul_cache::load_slot(int layer, int slot, int32_t expert_id) {
    try {
        std::vector<uint8_t> buffer;
        for (const auto & source : sources) {
            if (source.layer != layer) {
                continue;
            }
            buffer.resize(source.expert_size);
            files.at(source.file_idx)->read_at(
                buffer.data(), buffer.size(), source.offset + size_t(expert_id) * source.expert_size);
            ggml_backend_tensor_set(source.tensor, buffer.data(), size_t(slot) * source.tensor->nb[2], buffer.size());
            bytes_read += buffer.size();
        }
        return true;
    } catch (const std::exception & e) {
        last_error = e.what();
        return false;
    }
}

bool llama_longhaul_cache::remap(int layer, ggml_tensor * ids) {
    last_error.clear();
    mutex.lock();
    locked = true;
    locked_layer = layer;

    std::vector<int32_t> values(ggml_nbytes(ids) / sizeof(int32_t));
    ggml_backend_tensor_get(ids, values.data(), 0, ggml_nbytes(ids));
    for (int32_t & id : values) {
        if (id < 0 || id >= (int32_t) n_experts) {
            last_error = "router produced an out-of-range expert ID";
            release(layer);
            return false;
        }
        const int slot = find_slot(layer, id);
        if (slot < 0) {
            release(layer);
            return false;
        }
        id = slot;
    }
    ggml_backend_tensor_set(ids, values.data(), 0, ggml_nbytes(ids));
    return true;
}

void llama_longhaul_cache::release(int layer) {
    if (locked && (locked_layer == layer || layer < 0)) {
        locked = false;
        locked_layer = -1;
        mutex.unlock();
    }
}

uint32_t llama_longhaul_cache::capacity() const {
    return n_slots;
}

uint32_t llama_longhaul_cache::max_ubatch(uint32_t n_expert_used) const {
    return std::max<uint32_t>(1, n_slots / n_expert_used);
}

const std::string & llama_longhaul_cache::error() const {
    return last_error;
}

bool llama_longhaul_cache::failed() const {
    return !last_error.empty();
}