#include "gem/cpu.h"
#include "gem/metal.h"
#include "gem/metal_layer.h"

#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>

static size_t append_region(size_t *cursor, size_t bytes)
{
    *cursor = (*cursor + 255) & ~(size_t)255;
    size_t result = *cursor;
    *cursor += bytes;
    return result;
}

static size_t q4_size(size_t input, size_t output)
{
    return input / 32 * 18 * output;
}

static void fill_q4(uint8_t *matrix, size_t input, size_t output, unsigned seed)
{
    size_t row_bytes = input / 32 * 18;
    for (size_t row = 0; row < output; ++row)
        for (size_t block = 0; block < input / 32; ++block) {
            uint8_t *destination = matrix + row * row_bytes + block * 18;
            destination[0] = 0x00;
            destination[1] = 0x24; /* binary16 0.015625 */
            for (size_t i = 0; i < 16; ++i) {
                unsigned low = (unsigned)(row * 3 + block * 5 + i * 7 + seed) & 15;
                unsigned high = (unsigned)(row * 11 + block * 3 + i * 5 + seed * 2) & 15;
                destination[2 + i] = (uint8_t)(low | high << 4);
            }
        }
}

int main(void)
{
    enum { embedding = 64, query_heads = 2, kv_heads = 1, head_dimension = 32,
           query_size = 64, kv_size = 32, feed_forward = 64, cache_capacity = 2 };
    long page_size = sysconf(_SC_PAGESIZE);
    if (page_size <= 0)
        return 2;
    size_t cursor = 0;
    size_t q = append_region(&cursor, q4_size(embedding, query_size));
    size_t k = append_region(&cursor, q4_size(embedding, kv_size));
    size_t v = append_region(&cursor, q4_size(embedding, kv_size));
    size_t o = append_region(&cursor, q4_size(query_size, embedding));
    size_t gate = append_region(&cursor, q4_size(embedding, feed_forward));
    size_t up = append_region(&cursor, q4_size(embedding, feed_forward));
    size_t down = append_region(&cursor, q4_size(feed_forward, embedding));
    size_t attn_norm = append_region(&cursor, embedding * sizeof(float));
    size_t q_norm = append_region(&cursor, head_dimension * sizeof(float));
    size_t k_norm = append_region(&cursor, head_dimension * sizeof(float));
    size_t post_attn_norm = append_region(&cursor, embedding * sizeof(float));
    size_t ffn_norm = append_region(&cursor, embedding * sizeof(float));
    size_t post_ffn_norm = append_region(&cursor, embedding * sizeof(float));
    size_t mapping_size = (cursor + (size_t)page_size - 1) / (size_t)page_size *
                          (size_t)page_size;
    uint8_t *mapping = mmap(NULL, mapping_size, PROT_READ | PROT_WRITE,
                            MAP_PRIVATE | MAP_ANON, -1, 0);
    if (mapping == MAP_FAILED)
        return 2;
    fill_q4(mapping + q, embedding, query_size, 1);
    fill_q4(mapping + k, embedding, kv_size, 2);
    fill_q4(mapping + v, embedding, kv_size, 3);
    fill_q4(mapping + o, query_size, embedding, 4);
    fill_q4(mapping + gate, embedding, feed_forward, 5);
    fill_q4(mapping + up, embedding, feed_forward, 6);
    fill_q4(mapping + down, feed_forward, embedding, 7);
    size_t norm_offsets[] = {attn_norm, q_norm, k_norm, post_attn_norm, ffn_norm,
                             post_ffn_norm};
    size_t norm_counts[] = {embedding, head_dimension, head_dimension, embedding, embedding,
                            embedding};
    for (size_t n = 0; n < sizeof(norm_offsets) / sizeof(norm_offsets[0]); ++n) {
        float *weight = (float *)(mapping + norm_offsets[n]);
        for (size_t i = 0; i < norm_counts[n]; ++i)
            weight[i] = 0.75f + (float)((i + n) % 9) * 0.03125f;
    }

    struct gem_cpu_layer layer = {
        .embedding_size = embedding,
        .query_heads = query_heads,
        .kv_heads = kv_heads,
        .head_dimension = head_dimension,
        .rotary_dimension = head_dimension,
        .feed_forward_size = feed_forward,
        .rms_epsilon = 1e-6f,
        .rope_frequency_base = 10000.0f,
        .rope_frequency_scale = 1.0f,
        .attention_scale = 1.0f,
        .attention_query = mapping + q,
        .attention_key = mapping + k,
        .attention_value = mapping + v,
        .attention_output = mapping + o,
        .ffn_gate = mapping + gate,
        .ffn_up = mapping + up,
        .ffn_down = mapping + down,
        .attention_norm = (float *)(mapping + attn_norm),
        .query_norm = (float *)(mapping + q_norm),
        .key_norm = (float *)(mapping + k_norm),
        .post_attention_norm = (float *)(mapping + post_attn_norm),
        .ffn_norm = (float *)(mapping + ffn_norm),
        .post_ffn_norm = (float *)(mapping + post_ffn_norm),
        .layer_output_scale = 0.8125f,
    };
    float input[embedding];
    float expected[embedding];
    float actual[embedding];
    float initial_key_cache[cache_capacity * kv_size];
    float initial_value_cache[cache_capacity * kv_size];
    float cpu_key_cache[cache_capacity * kv_size];
    float cpu_value_cache[cache_capacity * kv_size];
    for (size_t i = 0; i < embedding; ++i)
        input[i] = ((int)(i % 19) - 9) * 0.035f;
    for (size_t i = 0; i < cache_capacity * kv_size; ++i) {
        initial_key_cache[i] = ((int)(i % 13) - 6) * 0.025f;
        initial_value_cache[i] = ((int)(i % 11) - 5) * 0.04f;
    }
    memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
    memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
    char error[512];
    if (gem_cpu_layer_decode(&layer, input, 2, cpu_key_cache, cpu_value_cache,
                             cache_capacity, 2, cache_capacity, expected,
                             error, sizeof(error)) != 0) {
        fprintf(stderr, "CPU layer: %s\n", error);
        return 1;
    }

    struct gem_metal *metal = NULL;
    if (gem_metal_open(&metal, error, sizeof(error)) != 0) {
        fprintf(stderr, "SKIP: %s\n", error);
        return 77;
    }
    struct gem_metal_layer_workspace *overflow_workspace = NULL;
    size_t overflowing_query_heads = SIZE_MAX / (32 * sizeof(float)) + 1;
    if (gem_metal_layer_workspace_open(&overflow_workspace, metal, 1, 1, 1, 1,
                                       overflowing_query_heads, 1,
                                       error, sizeof(error)) == 0 || overflow_workspace) {
        fprintf(stderr, "Metal score scratch byte overflow was accepted\n");
        gem_metal_layer_workspace_close(overflow_workspace);
        return 1;
    }
    struct gem_metal_buffer *model_buffer = NULL;
    struct gem_metal_buffer *cache_buffer = NULL;
    struct gem_metal_layer_workspace *workspace = NULL;
    size_t one_cache_bytes = sizeof(initial_key_cache);
    if (gem_metal_buffer_wrap(metal, mapping, mapping_size, &model_buffer,
                              error, sizeof(error)) != 0 ||
        gem_metal_buffer_new(metal, one_cache_bytes * 4, &cache_buffer,
                             error, sizeof(error)) != 0 ||
        gem_metal_layer_workspace_open(&workspace, metal, embedding, query_size, kv_size,
                                       feed_forward, query_heads, cache_capacity,
                                       error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal setup: %s\n", error);
        return 1;
    }
    size_t initial_score_bytes = query_heads * cache_capacity * 32 * sizeof(float);
    size_t grown_score_bytes = query_heads * cache_capacity * 64 * sizeof(float);
    if (gem_metal_layer_workspace_score_bytes(workspace) != initial_score_bytes) {
        fprintf(stderr, "Metal initial score scratch is %zu, expected %zu\n",
                gem_metal_layer_workspace_score_bytes(workspace), initial_score_bytes);
        return 1;
    }
    if (gem_metal_layer_workspace_reserve_score_rows(workspace, 64,
                                                     error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal score scratch reserve: %s\n", error);
        return 1;
    }
    if (gem_metal_layer_workspace_score_bytes(workspace) != grown_score_bytes) {
        fprintf(stderr, "Metal grown score scratch is %zu, expected %zu\n",
                gem_metal_layer_workspace_score_bytes(workspace), grown_score_bytes);
        return 1;
    }
    if (gem_metal_layer_workspace_reserve_score_rows(workspace, 1,
                                                     error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal score scratch no-shrink reserve: %s\n", error);
        return 1;
    }
    if (gem_metal_layer_workspace_score_bytes(workspace) != grown_score_bytes) {
        fprintf(stderr, "Metal score scratch shrank to %zu from %zu\n",
                gem_metal_layer_workspace_score_bytes(workspace), grown_score_bytes);
        return 1;
    }
    memcpy(gem_metal_buffer_contents(cache_buffer), initial_key_cache, one_cache_bytes);
    memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + one_cache_bytes,
           initial_value_cache, one_cache_bytes);
    if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
        gem_metal_layer_decode(workspace, &layer, mapping, mapping_size, model_buffer,
                               cache_buffer, 0, one_cache_bytes, cache_capacity, 2,
                               cache_capacity, 2, error, sizeof(error)) != 0 ||
        gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal layer: %s\n", error);
        return 1;
    }
    for (size_t i = 0; i < embedding; ++i)
        if (fabsf(actual[i] - expected[i]) > 2e-3f * (1.0f + fabsf(expected[i]))) {
            fprintf(stderr, "layer %zu: Metal %.9g CPU %.9g\n", i, actual[i], expected[i]);
            return 1;
        }
    /* Encode two complete layers into one command buffer. This checks the candidate token-span
     * scheduler boundary without changing the promoted layer-span default. */
    layer.attention_value = mapping + v;
    float middle[embedding];
    memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
    memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
    if (gem_cpu_layer_decode(&layer, input, 0, cpu_key_cache, cpu_value_cache,
                             cache_capacity, 0, cache_capacity, middle,
                             error, sizeof(error)) != 0) {
        fprintf(stderr, "CPU spanned layer one: %s\n", error);
        return 1;
    }
    memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
    memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
    if (gem_cpu_layer_decode(&layer, middle, 0, cpu_key_cache, cpu_value_cache,
                             cache_capacity, 0, cache_capacity, expected,
                             error, sizeof(error)) != 0) {
        fprintf(stderr, "CPU spanned layer two: %s\n", error);
        return 1;
    }
    for (size_t layer_index = 0; layer_index < 2; ++layer_index) {
        size_t base = layer_index * one_cache_bytes * 2;
        memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + base,
               initial_key_cache, one_cache_bytes);
        memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + base + one_cache_bytes,
               initial_value_cache, one_cache_bytes);
    }
    struct gem_metal_batch *batch = NULL;
    if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
        gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 ||
        gem_metal_layer_encode_decode(workspace, batch, &layer, mapping, mapping_size,
                                      model_buffer, cache_buffer, 0, one_cache_bytes,
                                      cache_capacity, 0, cache_capacity, 0,
                                      error, sizeof(error)) != 0 ||
        gem_metal_layer_encode_decode(workspace, batch, &layer, mapping, mapping_size,
                                      model_buffer, cache_buffer, one_cache_bytes * 2,
                                      one_cache_bytes * 3, cache_capacity, 0,
                                      cache_capacity, 0, error, sizeof(error)) != 0 ||
        gem_metal_batch_submit(batch, error, sizeof(error)) != 0 ||
        gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal spanned layers: %s\n", error);
        gem_metal_batch_close(batch);
        return 1;
    }
    gem_metal_batch_close(batch);
    for (size_t i = 0; i < embedding; ++i)
        if (fabsf(actual[i] - expected[i]) > 4e-3f * (1.0f + fabsf(expected[i]))) {
            fprintf(stderr, "spanned layer %zu: Metal %.9g CPU %.9g\n",
                    i, actual[i], expected[i]);
            return 1;
        }
    layer.attention_value = NULL;
    memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
    memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
    if (gem_cpu_layer_decode(&layer, input, 1, cpu_key_cache, cpu_value_cache,
                             cache_capacity, 1, 0, expected,
                             error, sizeof(error)) != 0) {
        fprintf(stderr, "CPU K-as-V layer: %s\n", error);
        return 1;
    }
    memcpy(gem_metal_buffer_contents(cache_buffer), initial_key_cache, one_cache_bytes);
    memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + one_cache_bytes,
           initial_value_cache, one_cache_bytes);
    if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
        gem_metal_layer_decode(workspace, &layer, mapping, mapping_size, model_buffer,
                               cache_buffer, 0, one_cache_bytes, cache_capacity, 1,
                               0, 1, error, sizeof(error)) != 0 ||
        gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
        fprintf(stderr, "Metal K-as-V layer: %s\n", error);
        return 1;
    }
    for (size_t i = 0; i < embedding; ++i)
        if (fabsf(actual[i] - expected[i]) > 2e-3f * (1.0f + fabsf(expected[i]))) {
            fprintf(stderr, "K-as-V layer %zu: Metal %.9g CPU %.9g\n",
                    i, actual[i], expected[i]);
            return 1;
        }
    gem_metal_layer_workspace_close(workspace);
    gem_metal_buffer_close(cache_buffer);
    gem_metal_buffer_close(model_buffer);
    gem_metal_close(metal);
    munmap(mapping, mapping_size);
    return 0;
}