#include "testing.h"

#include "../src/llama-longhaul.h"

#include "ggml-backend.h"

#include <cstring>
#include <cstdio>
#include <memory>
#include <vector>

struct longhaul_fixture {
    static constexpr size_t expert_size = 32;

    FILE * file = nullptr;
    ggml_backend_ptr backend;
    ggml_context_ptr ctx;
    ggml_backend_buffer_ptr buffer;
    ggml_tensor * weights_a = nullptr;
    ggml_tensor * weights_b = nullptr;
    ggml_tensor * ids = nullptr;
    std::unique_ptr<llama_longhaul_cache> cache;

    longhaul_fixture(
            size_t n_slots,
            uint32_t n_experts,
            bool two_sources = false,
            uint32_t written_experts = 0,
            ggml_backend_dev_t dev = nullptr) {
        file = tmpfile();
        GGML_ASSERT(file != nullptr);
        write_experts(written_experts == 0 ? n_experts : written_experts, two_sources);

        backend.reset(dev
            ? ggml_backend_dev_init(dev, nullptr)
            : ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr));
        GGML_ASSERT(backend);

        ggml_init_params params = {
            /*.mem_size   =*/ 16 * 1024,
            /*.mem_buffer =*/ nullptr,
            /*.no_alloc   =*/ true,
        };
        ctx.reset(ggml_init(params));
        GGML_ASSERT(ctx);

        weights_a = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots);
        if (two_sources) {
            weights_b = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots);
        }
        ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, n_slots);
        buffer.reset(ggml_backend_alloc_ctx_tensors(ctx.get(), backend.get()));
        GGML_ASSERT(buffer);

        llama_files files;
        files.emplace_back(std::make_unique<llama_file>(file));

        std::vector<llama_model_loader::longhaul_source> sources = {
            { weights_a, 0, 0, expert_size, 0 },
        };
        if (two_sources) {
            sources.push_back({ weights_b, 0, n_experts * expert_size, expert_size, 0 });
        }
        cache = std::make_unique<llama_longhaul_cache>(
            std::move(files), std::move(sources), n_slots, n_experts, 1);
    }

    ~longhaul_fixture() {
        cache.reset();
        if (file) {
            fclose(file);
        }
    }

    void write_experts(uint32_t n_experts, bool two_sources) {
        GGML_ASSERT(fseek(file, 0, SEEK_END) == 0);
        for (int source = 0; source < (two_sources ? 2 : 1); ++source) {
            for (uint32_t expert = 0; expert < n_experts; ++expert) {
                std::vector<uint8_t> data(expert_size, uint8_t(1 + expert + source * 32));
                GGML_ASSERT(fwrite(data.data(), 1, data.size(), file) == data.size());
            }
        }
        GGML_ASSERT(fflush(file) == 0);
    }

    std::vector<int32_t> remap(std::vector<int32_t> values) {
        GGML_ASSERT(values.size() <= (size_t) ids->ne[0]);
        ggml_backend_tensor_set(ids, values.data(), 0, values.size() * sizeof(int32_t));
        ids->ne[0] = values.size();
        const bool ok = cache->remap(0, ids);
        if (ok) {
            ggml_backend_tensor_get(ids, values.data(), 0, values.size() * sizeof(int32_t));
        }
        cache->release(0);
        return ok ? values : std::vector<int32_t>{};
    }

    uint8_t slot_value(ggml_tensor * tensor, int32_t slot) {
        uint8_t value = 0;
        ggml_backend_tensor_get(tensor, &value, size_t(slot) * tensor->nb[2], 1);
        return value;
    }
};

static void test_batch_planning(testing & t) {
    longhaul_fixture fixture(3, 4);

    const auto initial = fixture.remap({0, 1, 2});
    t.assert_equal(3u, initial.size());
    t.assert_equal(3u, fixture.cache->misses());

    // Expert 0 is the oldest cache entry. A sequential miss for expert 3 used
    // to evict it before the later 0 in this same routed batch was observed.
    const auto remapped = fixture.remap({3, 0, 3});
    t.assert_equal(3u, remapped.size());
    t.assert_equal("only expert 3 is loaded", 4u, fixture.cache->misses());
    t.assert_true("duplicates map to the same slot", remapped[0] == remapped[2]);
    t.assert_true("different experts map to different slots", remapped[0] != remapped[1]);
    t.assert_equal(uint8_t(4), fixture.slot_value(fixture.weights_a, remapped[0]));
    t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[1]));
}

static void test_multiple_sources(testing & t) {
    longhaul_fixture fixture(2, 4, true);

    const auto remapped = fixture.remap({2, 1});
    t.assert_equal(2u, remapped.size());
    t.assert_equal(uint64_t(4 * longhaul_fixture::expert_size), fixture.cache->bytes_read_count());
    t.assert_equal(uint8_t(3),  fixture.slot_value(fixture.weights_a, remapped[0]));
    t.assert_equal(uint8_t(35), fixture.slot_value(fixture.weights_b, remapped[0]));
    t.assert_equal(uint8_t(2),  fixture.slot_value(fixture.weights_a, remapped[1]));
    t.assert_equal(uint8_t(34), fixture.slot_value(fixture.weights_b, remapped[1]));
}

static void test_invalid_ids(testing & t) {
    longhaul_fixture fixture(2, 4);

    const auto remapped = fixture.remap({0, 4});
    t.assert_true(remapped.empty());
    t.assert_true(fixture.cache->failed());
    t.assert_equal(0u, fixture.cache->misses());
}

static void test_read_failure_recovery(testing & t) {
    longhaul_fixture fixture(1, 4, false, 3);

    const auto failed = fixture.remap({3});
    t.assert_true(failed.empty());
    t.assert_true(fixture.cache->failed());
    t.assert_equal(0u, fixture.cache->misses());

    fixture.write_experts(1, false);
    const auto remapped = fixture.remap({3});
    t.assert_equal(1u, remapped.size());
    t.assert_equal(1u, fixture.cache->misses());
    t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[0]));
}

static void test_staged_device_upload(testing & t, ggml_backend_dev_t dev) {
    longhaul_fixture fixture(2, 4, true, 0, dev);

    const auto remapped = fixture.remap({3, 1});
    t.assert_equal(2u, remapped.size());
    t.assert_equal(uint64_t(4 * longhaul_fixture::expert_size), fixture.cache->bytes_read_count());
    t.assert_equal(uint8_t(4),  fixture.slot_value(fixture.weights_a, remapped[0]));
    t.assert_equal(uint8_t(36), fixture.slot_value(fixture.weights_b, remapped[0]));
    t.assert_equal(uint8_t(2),  fixture.slot_value(fixture.weights_a, remapped[1]));
    t.assert_equal(uint8_t(34), fixture.slot_value(fixture.weights_b, remapped[1]));
}

int main(int argc, char ** argv) {
    testing t;

    const char * verbose = getenv("LLAMA_TEST_VERBOSE");
    if (verbose) {
        t.verbose = std::string(verbose) == "1";
    }
    if (!t.verbose) {
        llama_log_set([](ggml_log_level, const char *, void *) {}, nullptr);
    }
    ggml_backend_load_all();

    if (argc > 1) {
        t.set_filter(argv[1]);
    }

    t.test("batch_planning",  test_batch_planning);
    t.test("multiple_sources", test_multiple_sources);
    t.test("invalid_ids",     test_invalid_ids);
    t.test("read_failure",    test_read_failure_recovery);

    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
        const char * backend = ggml_backend_reg_name(ggml_backend_dev_backend_reg(dev));
        if (strcmp(backend, "Vulkan") == 0) {
            t.test("vulkan_staged_upload", [dev](testing & current) {
                test_staged_device_upload(current, dev);
            });
            break;
        }
    }

    return t.summary();
}