#pragma once

#include "llama-mmap.h"
#include "llama-model-loader.h"

#include <cstdint>
#include <memory>
#include <mutex>
#include <vector>

struct 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);
    ~llama_longhaul_cache();

    bool remap(int layer, ggml_tensor * ids);
    void release(int layer);

    uint32_t max_ubatch(uint32_t n_expert_used) const;
    const std::string & error() const;
    bool failed() const;

private:
    struct layer_state {
        std::vector<int32_t> expert_ids;
        std::vector<uint64_t> last_used;
    };

    llama_files files;
    std::vector<llama_model_loader::longhaul_source> sources;
    std::vector<layer_state> layers;
    size_t n_slots;
    uint32_t n_experts;
    uint64_t tick = 0;
    uint64_t n_hits = 0;
    uint64_t n_misses = 0;
    uint64_t bytes_read = 0;
    std::mutex mutex;
    bool locked = false;
    int locked_layer = -1;
    std::string last_error;

    int find_slot(int layer, int32_t expert_id);
    bool load_slot(int layer, int slot, int32_t expert_id);
};
