/* dragon_read.c — dragon.c PHASE 1: safetensors reader.
 * mmap a shard, parse the header, locate a tensor by byte offset, print values.
 * VALIDATION: must match python ground truth (safetensors_ground_truth.json, embed row 962).
 * Build: gcc -O2 -o dragon_read dragon_read.c
 * Run:   ./dragon_read <shard.safetensors>
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

/* bf16 -> float: bf16 is the top 16 bits of an f32 */
static float bf16_to_f32(uint16_t b) {
    uint32_t u = ((uint32_t)b) << 16;
    float f; memcpy(&f, &u, 4); return f;
}

/* minimal JSON scan: find "<name>":{..."data_offsets":[a,b]...} and the dtype/shape */
static const char* find_tensor(const char* hdr, size_t hlen, const char* name,
                               long long* off0, long long* off1, char* dtype, long long* d0, long long* d1) {
    char key[512]; snprintf(key, sizeof key, "\"%s\":", name);
    const char* p = memmem(hdr, hlen, key, strlen(key));
    if (!p) return NULL;
    const char* q = strstr(p, "\"dtype\":\"");
    if (q) sscanf(q + 9, "%15[^\"]", dtype);
    q = strstr(p, "\"shape\":[");
    if (q) sscanf(q + 9, "%lld,%lld", d0, d1);
    q = strstr(p, "\"data_offsets\":[");
    if (!q) return NULL;
    sscanf(q + 16, "%lld,%lld", off0, off1);
    return p;
}

int main(int argc, char** argv) {
    if (argc < 2) { fprintf(stderr, "usage: %s <shard.safetensors>\n", argv[0]); return 1; }
    int fd = open(argv[1], O_RDONLY);
    if (fd < 0) { perror("open"); return 1; }
    struct stat st; fstat(fd, &st);
    /* mmap the whole shard — read weights off disk, never load into heap RAM */
    const uint8_t* base = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    if (base == MAP_FAILED) { perror("mmap"); return 1; }

    uint64_t hlen; memcpy(&hlen, base, 8);            /* u64 LE header length */
    const char* hdr = (const char*)(base + 8);
    printf("shard: %s (%.1f GB)\nheader: %llu bytes, data starts at byte %llu\n",
           argv[1], st.st_size / 1e9, (unsigned long long)hlen, (unsigned long long)(8 + hlen));

    long long o0=-1, o1=-1, d0=-1, d1=-1; char dtype[16] = {0};
    const char* name = "model.embed_tokens.weight";
    if (!find_tensor(hdr, hlen, name, &o0, &o1, dtype, &d0, &d1)) {
        fprintf(stderr, "tensor %s not in this shard\n", name); return 2;
    }
    printf("tensor %s: dtype=%s shape=[%lld,%lld] offsets=[%lld,%lld]\n", name, dtype, d0, d1, o0, o1);

    /* row 962, first 8 values — bf16, row-major: offset = (962*d1 + col) * 2 bytes */
    const uint8_t* data = base + 8 + hlen + o0;
    printf("row 962, first 8 values (C mmap read):\n  ");
    for (int c = 0; c < 8; c++) {
        uint16_t b; memcpy(&b, data + ((long long)962 * d1 + c) * 2, 2);
        printf("%.6f ", bf16_to_f32(b));
    }
    printf("\nCompare against safetensors_ground_truth.json -> first_8_values_row962\n");
    munmap((void*)base, st.st_size); close(fd);
    return 0;
}