#pragma once

#include "ggml-backend.h"

#ifdef  __cplusplus
extern "C" {
#endif

#define RPC_PROTO_MAJOR_VERSION    4
#define RPC_PROTO_MINOR_VERSION    0
#define RPC_PROTO_PATCH_VERSION    4

#ifdef  __cplusplus
static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
#endif

#define GGML_RPC_MAX_SERVERS       16

// backend API
GGML_BACKEND_API ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device);
GGML_BACKEND_API bool ggml_backend_is_rpc(ggml_backend_t backend);

GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device);

GGML_BACKEND_API void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total);

GGML_BACKEND_API void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir,
                                                    size_t n_threads, size_t n_devices, ggml_backend_dev_t * devices);

struct ggml_backend_rpc_longhaul_shard {
    const char * name;
    uint64_t size;
    uint64_t metadata_size;
    uint8_t metadata_sha256[32];
};

struct ggml_backend_rpc_longhaul_source {
    struct ggml_tensor * tensor;
    uint32_t shard;
    uint64_t offset;
    uint64_t expert_size;
    int32_t layer;
};

struct ggml_backend_rpc_longhaul_params {
    uint32_t n_layers;
    uint32_t n_experts;
    uint32_t n_slots;
    size_t n_shards;
    const struct ggml_backend_rpc_longhaul_shard * shards;
    size_t n_sources;
    const struct ggml_backend_rpc_longhaul_source * sources;
};

GGML_BACKEND_API bool ggml_backend_rpc_longhaul_register(
        const struct ggml_backend_rpc_longhaul_params * params,
        uint64_t * registration_id,
        char * error,
        size_t error_size);

GGML_BACKEND_API void ggml_backend_rpc_longhaul_unregister(
        struct ggml_tensor * tensor,
        uint64_t registration_id);

GGML_BACKEND_API void ggml_backend_rpc_start_server_with_options(
        const char * endpoint,
        const char * cache_dir,
        const char * longhaul_root,
        size_t n_threads,
        size_t n_devices,
        ggml_backend_dev_t * devices);

GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_reg(void);
GGML_BACKEND_API ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint);

#ifdef  __cplusplus
}
#endif