/* shake256-fanout.cu — reference cuda fan-out worker for lumbda.
 *
 * Reads an S-expression input portal listing N inputs, hashes each
 * with SHAKE256 on its own GPU thread, writes an output portal with
 * N hashes.  The "fan-out primitive" lumbda would expose via its
 * spawn-process primitive across Python / C / asm tiers — every tier
 * inherits this binary without dragging the CUDA toolchain into
 * lumbda's core.
 *
 * Wire format (input portal — written by the caller tier):
 *
 *   (cuda-shake-fanout
 *     (output-bytes 32)
 *     (inputs
 *       "hex-bytes-1"
 *       "hex-bytes-2"
 *       ...))
 *
 * Wire format (output portal — written by this binary):
 *
 *   (cuda-shake-fanout-result
 *     (n 17)
 *     (output-bytes 32)
 *     (hashes
 *       "hex-output-1"
 *       "hex-output-2"
 *       ...)
 *     (timing-ms (load 0.4) (launch 0.1) (kernel 12.7) (write 0.3)))
 *
 * Usage: shake256-fanout <input.portal> <output.portal>
 *
 * License: AGPLv3.  Keccak permutation derived from FIPS 202 reference;
 * the tiny-sha3 distillation (Markku-Juhani O. Saarinen, CC0 / public
 * domain) was used as the starting point and re-licensed here under
 * AGPLv3 to align with lumbda's own license.
 */

#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdint.h>

#define CUDA_CHECK(stmt) do { \
    cudaError_t err = (stmt); \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \
                #stmt, __FILE__, __LINE__, cudaGetErrorString(err)); \
        exit(1); \
    } \
} while (0)

#define KECCAK_RATE 136   /* SHAKE256 rate (1600 - 2*256) / 8 */

/* ── Keccak permutation, __host__ __device__ — runs on either ── */

__host__ __device__ static const uint64_t RC[24] = {
    0x0000000000000001ULL, 0x0000000000008082ULL,
    0x800000000000808AULL, 0x8000000080008000ULL,
    0x000000000000808BULL, 0x0000000080000001ULL,
    0x8000000080008081ULL, 0x8000000000008009ULL,
    0x000000000000008AULL, 0x0000000000000088ULL,
    0x0000000080008009ULL, 0x000000008000000AULL,
    0x000000008000808BULL, 0x800000000000008BULL,
    0x8000000000008089ULL, 0x8000000000008003ULL,
    0x8000000000008002ULL, 0x8000000000000080ULL,
    0x000000000000800AULL, 0x800000008000000AULL,
    0x8000000080008081ULL, 0x8000000000008080ULL,
    0x0000000080000001ULL, 0x8000000080008008ULL,
};
__host__ __device__ static const int R[24] = {
    1,  3,  6, 10, 15, 21, 28, 36, 45, 55, 2, 14,
   27, 41, 56,  8, 25, 43, 62, 18, 39, 61, 20, 44,
};
__host__ __device__ static const int PI[24] = {
   10,  7, 11, 17, 18,  3,  5, 16,  8, 21, 24, 4,
   15, 23, 19, 13, 12,  2, 20, 14, 22,  9,  6, 1,
};

__host__ __device__ static inline uint64_t rotl64(uint64_t x, int n) {
    return (x << n) | (x >> (64 - n));
}

__host__ __device__ static void keccakf(uint64_t s[25]) {
    uint64_t t, bc[5];
    for (int r = 0; r < 24; r++) {
        for (int i = 0; i < 5; i++)
            bc[i] = s[i] ^ s[i + 5] ^ s[i + 10] ^ s[i + 15] ^ s[i + 20];
        for (int i = 0; i < 5; i++) {
            t = bc[(i + 4) % 5] ^ rotl64(bc[(i + 1) % 5], 1);
            for (int j = 0; j < 25; j += 5) s[j + i] ^= t;
        }
        t = s[1];
        for (int i = 0; i < 24; i++) {
            int j = PI[i];
            bc[0] = s[j];
            s[j] = rotl64(t, R[i]);
            t = bc[0];
        }
        for (int j = 0; j < 25; j += 5) {
            for (int i = 0; i < 5; i++) bc[i] = s[j + i];
            for (int i = 0; i < 5; i++)
                s[j + i] ^= (~bc[(i + 1) % 5]) & bc[(i + 2) % 5];
        }
        s[0] ^= RC[r];
    }
}

/* per-thread SHAKE256: absorb input, finalize, squeeze out_len bytes. */
__device__ static void shake256_oneshot(const uint8_t *in, uint32_t in_len,
                                        uint8_t *out, uint32_t out_len)
{
    uint8_t state[200] = {0};
    uint32_t cur = 0;
    for (uint32_t i = 0; i < in_len; i++) {
        state[cur++] ^= in[i];
        if (cur == KECCAK_RATE) { keccakf((uint64_t*)state); cur = 0; }
    }
    /* SHAKE pad */
    state[cur] ^= 0x1F;
    state[KECCAK_RATE - 1] ^= 0x80;
    keccakf((uint64_t*)state);
    cur = 0;
    for (uint32_t i = 0; i < out_len; i++) {
        if (cur == KECCAK_RATE) { keccakf((uint64_t*)state); cur = 0; }
        out[i] = state[cur++];
    }
}

/* kernel: one thread per input. inputs are stored back-to-back with
 * a per-input offset[] + len[] index. outputs are flat output_bytes ×
 * N. block = 64 threads (~enough warps for typical fan-out widths). */
__global__ void shake_fanout_kernel(const uint8_t *inputs,
                                    const uint32_t *offsets,
                                    const uint32_t *lengths,
                                    uint8_t *outputs,
                                    uint32_t output_bytes,
                                    uint32_t n)
{
    uint32_t t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n) return;
    shake256_oneshot(inputs + offsets[t], lengths[t],
                     outputs + (size_t)t * output_bytes, output_bytes);
}

/* ── host: portal I/O + driver ────────────────────────────────────── */

static char *slurp(const char *path, size_t *len_out) {
    FILE *f = fopen(path, "rb");
    if (!f) { perror(path); exit(1); }
    fseek(f, 0, SEEK_END); long n = ftell(f); fseek(f, 0, SEEK_SET);
    char *buf = (char*)malloc(n + 1);
    if (fread(buf, 1, n, f) != (size_t)n) { perror(path); exit(1); }
    buf[n] = 0; fclose(f);
    if (len_out) *len_out = n;
    return buf;
}

static int hex_to_bytes(const char *hex, size_t hex_len, uint8_t *out) {
    if (hex_len % 2 != 0) return 0;
    for (size_t i = 0; i < hex_len; i += 2) {
        unsigned v;
        if (sscanf(hex + i, "%2x", &v) != 1) return 0;
        out[i / 2] = (uint8_t)v;
    }
    return 1;
}

static void bytes_to_hex(const uint8_t *bytes, size_t n, char *out) {
    static const char *H = "0123456789abcdef";
    for (size_t i = 0; i < n; i++) {
        out[2 * i]     = H[bytes[i] >> 4];
        out[2 * i + 1] = H[bytes[i] & 0xF];
    }
    out[2 * n] = 0;
}

/* parse: extract `(output-bytes N)` & all "..."-quoted strings inside
 * `(inputs ...)`.  Tolerant of whitespace; assumes inputs are simple
 * hex strings. */
static int parse_input_portal(const char *src, uint32_t *out_bytes,
                              char ***hex_inputs, uint32_t *n_inputs)
{
    *out_bytes = 32;
    const char *p = strstr(src, "(output-bytes");
    if (p) sscanf(p, "(output-bytes %u)", out_bytes);
    p = strstr(src, "(inputs");
    if (!p) return 0;
    char **list = NULL;
    uint32_t count = 0, cap = 0;
    p += strlen("(inputs");
    while ((p = strchr(p, '"')) != NULL) {
        const char *end = strchr(p + 1, '"');
        if (!end) break;
        if (count == cap) {
            cap = cap ? cap * 2 : 16;
            list = (char**)realloc(list, cap * sizeof(char*));
        }
        size_t L = end - (p + 1);
        list[count] = (char*)malloc(L + 1);
        memcpy(list[count], p + 1, L); list[count][L] = 0;
        count++;
        p = end + 1;
    }
    *hex_inputs = list;
    *n_inputs = count;
    return count > 0;
}

static int process_one(const char *in_path, const char *out_path);
static int process_one_bin(const char *in_path, const char *out_path);

int main(int argc, char **argv) {
    /* daemon mode: stay alive across many calls, kill the 200 ms
     * cuda-ctx-init overhead per spawn.
     *
     * Protocol (stdin, one command per line):
     *   process <in.portal> <out.portal>
     *   quit
     * Stdout response (one line per command):
     *   done <out.portal>      → success
     *   error <reason>         → failure
     *   bye                    → quit acknowledged
     */
    if (argc >= 2 && !strcmp(argv[1], "--daemon")) {
        /* Touch the CUDA context once at startup so first request is fast. */
        int dev_count = 0;
        cudaGetDeviceCount(&dev_count);
        if (dev_count == 0) {
            fprintf(stderr, "no CUDA devices\n"); return 1;
        }
        cudaFree(0);  /* forces context init */
        fprintf(stdout, "ready\n"); fflush(stdout);

        char line[4096];
        while (fgets(line, sizeof(line), stdin)) {
            char *nl = strchr(line, '\n');
            if (nl) *nl = 0;
            if (!strcmp(line, "quit")) {
                fprintf(stdout, "bye\n"); fflush(stdout);
                break;
            }
            int is_bin = 0;
            char *cmd_args = NULL;
            if (!strncmp(line, "process ", 8))         { cmd_args = line + 8; is_bin = 0; }
            else if (!strncmp(line, "process-bin ", 12)) { cmd_args = line + 12; is_bin = 1; }
            if (cmd_args) {
                char *in_path = cmd_args;
                char *sep = strchr(in_path, ' ');
                if (!sep) {
                    fprintf(stdout, "error bad-args\n"); fflush(stdout);
                    continue;
                }
                *sep = 0;
                char *out_path = sep + 1;
                int r = is_bin ? process_one_bin(in_path, out_path)
                               : process_one(in_path, out_path);
                if (r == 0) {
                    fprintf(stdout, "done %s\n", out_path);
                } else {
                    fprintf(stdout, "error process-failed\n");
                }
                fflush(stdout);
            } else if (line[0]) {
                fprintf(stdout, "error unknown-command\n"); fflush(stdout);
            }
        }
        return 0;
    }

    if (argc >= 4 && !strcmp(argv[1], "--binary")) {
        return process_one_bin(argv[2], argv[3]);
    }
    if (argc < 3) {
        fprintf(stderr,
            "usage: %s <input.portal> <output.portal>\n"
            "       %s --binary <input.bin> <output.bin>\n"
            "       %s --daemon  (read commands on stdin)\n",
            argv[0], argv[0], argv[0]);
        return 1;
    }
    return process_one(argv[1], argv[2]);
}

/* Binary format:
 *   in:  u32 out_bytes | u32 n | (u32 len | len bytes) × n
 *   out: u32 n | u32 out_bytes | (out_bytes bytes) × n
 *
 * Eliminates the hex-string parse + emit overhead. At 262 MB workloads
 * we measured 421 sec parsing 524 MB of hex text; binary skips that
 * entirely. */
static int process_one_bin(const char *in_path, const char *out_path) {
    struct timespec t0, t1, t2, t3, t4;
    clock_gettime(CLOCK_MONOTONIC, &t0);

    FILE *f = fopen(in_path, "rb");
    if (!f) { perror(in_path); return 1; }
    fseek(f, 0, SEEK_END); long sz = ftell(f); fseek(f, 0, SEEK_SET);
    uint8_t *buf = (uint8_t*)malloc(sz);
    if (fread(buf, 1, sz, f) != (size_t)sz) { perror(in_path); return 1; }
    fclose(f);

    uint32_t out_bytes, n;
    memcpy(&out_bytes, buf, 4);
    memcpy(&n, buf + 4, 4);
    /* slice into offsets / lengths */
    uint32_t *offsets = (uint32_t*)malloc(n * sizeof(uint32_t));
    uint32_t *lengths = (uint32_t*)malloc(n * sizeof(uint32_t));
    size_t cur = 8;
    size_t data_off = 0;
    /* one pass: collect lengths, compute offsets in a separate flat array */
    size_t total = 0;
    for (uint32_t i = 0; i < n; i++) {
        if (cur + 4 > (size_t)sz) {
            fprintf(stderr, "%s: truncated at input %u\n", in_path, i); return 1;
        }
        uint32_t L; memcpy(&L, buf + cur, 4); cur += 4;
        if (cur + L > (size_t)sz) {
            fprintf(stderr, "%s: bad length at input %u\n", in_path, i); return 1;
        }
        offsets[i] = (uint32_t)data_off;
        lengths[i] = L;
        data_off += L;
        cur += L;
    }
    total = data_off;
    /* flat array — just point at the buf bytes via offsets relative to start */
    uint8_t *flat = (uint8_t*)malloc(total ? total : 1);
    cur = 8;
    size_t fcur = 0;
    for (uint32_t i = 0; i < n; i++) {
        uint32_t L = lengths[i];
        cur += 4;
        memcpy(flat + fcur, buf + cur, L);
        cur += L; fcur += L;
    }
    free(buf);
    clock_gettime(CLOCK_MONOTONIC, &t1);

    uint8_t *d_in; uint32_t *d_off, *d_len; uint8_t *d_out;
    CUDA_CHECK(cudaMalloc(&d_in,  total ? total : 1));
    CUDA_CHECK(cudaMalloc(&d_off, n * sizeof(uint32_t)));
    CUDA_CHECK(cudaMalloc(&d_len, n * sizeof(uint32_t)));
    CUDA_CHECK(cudaMalloc(&d_out, (size_t)n * out_bytes));
    CUDA_CHECK(cudaMemcpy(d_in,  flat,    total ? total : 1,    cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_off, offsets, n * sizeof(uint32_t), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_len, lengths, n * sizeof(uint32_t), cudaMemcpyHostToDevice));
    clock_gettime(CLOCK_MONOTONIC, &t2);

    int threads = 64;
    int blocks  = (n + threads - 1) / threads;
    shake_fanout_kernel<<<blocks, threads>>>(d_in, d_off, d_len, d_out, out_bytes, n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    clock_gettime(CLOCK_MONOTONIC, &t3);

    uint8_t *h_out = (uint8_t*)malloc((size_t)n * out_bytes);
    CUDA_CHECK(cudaMemcpy(h_out, d_out, (size_t)n * out_bytes, cudaMemcpyDeviceToHost));

    FILE *fo = fopen(out_path, "wb");
    if (!fo) { perror(out_path); return 1; }
    fwrite(&n, 4, 1, fo);
    fwrite(&out_bytes, 4, 1, fo);
    fwrite(h_out, 1, (size_t)n * out_bytes, fo);
    fclose(fo);
    clock_gettime(CLOCK_MONOTONIC, &t4);

    double ld_ms     = (t1.tv_sec - t0.tv_sec) * 1e3 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
    double launch_ms = (t2.tv_sec - t1.tv_sec) * 1e3 + (t2.tv_nsec - t1.tv_nsec) / 1e6;
    double kern_ms   = (t3.tv_sec - t2.tv_sec) * 1e3 + (t3.tv_nsec - t2.tv_nsec) / 1e6;
    double write_ms  = (t4.tv_sec - t3.tv_sec) * 1e3 + (t4.tv_nsec - t3.tv_nsec) / 1e6;
    fprintf(stderr, "  bin: load=%.1f launch=%.1f kernel=%.2f write=%.1f ms (n=%u, %.0f MB)\n",
            ld_ms, launch_ms, kern_ms, write_ms, n, total / 1e6);

    cudaFree(d_in); cudaFree(d_off); cudaFree(d_len); cudaFree(d_out);
    free(offsets); free(lengths); free(flat); free(h_out);
    return 0;
}

static int process_one(const char *in_path, const char *out_path) {
    struct timespec t0, t1, t2, t3, t4;
    clock_gettime(CLOCK_MONOTONIC, &t0);

    size_t portal_len;
    char *portal = slurp(in_path, &portal_len);
    uint32_t out_bytes = 32;
    char **hex_in = NULL; uint32_t n = 0;
    if (!parse_input_portal(portal, &out_bytes, &hex_in, &n)) {
        fprintf(stderr, "no (inputs ...) found in %s\n", in_path);
        return 1;
    }

    /* lay out flat input bytes with offsets+lengths */
    uint32_t *offsets = (uint32_t*)malloc(n * sizeof(uint32_t));
    uint32_t *lengths = (uint32_t*)malloc(n * sizeof(uint32_t));
    size_t total = 0;
    for (uint32_t i = 0; i < n; i++) {
        size_t hex_len = strlen(hex_in[i]);
        if (hex_len % 2 != 0) {
            fprintf(stderr, "input %u has odd hex length\n", i); return 1;
        }
        offsets[i] = total;
        lengths[i] = hex_len / 2;
        total += hex_len / 2;
    }
    uint8_t *flat = (uint8_t*)malloc(total ? total : 1);
    for (uint32_t i = 0; i < n; i++) {
        if (!hex_to_bytes(hex_in[i], strlen(hex_in[i]), flat + offsets[i])) {
            fprintf(stderr, "bad hex at input %u\n", i); return 1;
        }
    }
    clock_gettime(CLOCK_MONOTONIC, &t1);

    /* upload */
    uint8_t  *d_in;  uint32_t *d_off, *d_len; uint8_t *d_out;
    CUDA_CHECK(cudaMalloc(&d_in,  total ? total : 1));
    CUDA_CHECK(cudaMalloc(&d_off, n * sizeof(uint32_t)));
    CUDA_CHECK(cudaMalloc(&d_len, n * sizeof(uint32_t)));
    CUDA_CHECK(cudaMalloc(&d_out, (size_t)n * out_bytes));
    CUDA_CHECK(cudaMemcpy(d_in,  flat,    total ? total : 1,    cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_off, offsets, n * sizeof(uint32_t), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_len, lengths, n * sizeof(uint32_t), cudaMemcpyHostToDevice));
    clock_gettime(CLOCK_MONOTONIC, &t2);

    int threads = 64;
    int blocks  = (n + threads - 1) / threads;
    shake_fanout_kernel<<<blocks, threads>>>(d_in, d_off, d_len, d_out, out_bytes, n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    clock_gettime(CLOCK_MONOTONIC, &t3);

    uint8_t *h_out = (uint8_t*)malloc((size_t)n * out_bytes);
    CUDA_CHECK(cudaMemcpy(h_out, d_out, (size_t)n * out_bytes, cudaMemcpyDeviceToHost));

    /* write portal */
    FILE *f = fopen(out_path, "w");
    if (!f) { perror(out_path); return 1; }
    fprintf(f, "(cuda-shake-fanout-result\n");
    fprintf(f, "  (n %u)\n", n);
    fprintf(f, "  (output-bytes %u)\n", out_bytes);
    fprintf(f, "  (hashes\n");
    char *hex_buf = (char*)malloc(out_bytes * 2 + 1);
    for (uint32_t i = 0; i < n; i++) {
        bytes_to_hex(h_out + (size_t)i * out_bytes, out_bytes, hex_buf);
        fprintf(f, "    \"%s\"\n", hex_buf);
    }
    fprintf(f, "  )\n");
    clock_gettime(CLOCK_MONOTONIC, &t4);
    double ld_ms     = (t1.tv_sec - t0.tv_sec) * 1e3 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
    double launch_ms = (t2.tv_sec - t1.tv_sec) * 1e3 + (t2.tv_nsec - t1.tv_nsec) / 1e6;
    double kern_ms   = (t3.tv_sec - t2.tv_sec) * 1e3 + (t3.tv_nsec - t2.tv_nsec) / 1e6;
    double write_ms  = (t4.tv_sec - t3.tv_sec) * 1e3 + (t4.tv_nsec - t3.tv_nsec) / 1e6;
    fprintf(f, "  (timing-ms (load %.3f) (launch %.3f) (kernel %.3f) (write %.3f)))\n",
            ld_ms, launch_ms, kern_ms, write_ms);
    fclose(f);

    cudaFree(d_in); cudaFree(d_off); cudaFree(d_len); cudaFree(d_out);
    for (uint32_t i = 0; i < n; i++) free(hex_in[i]);
    free(hex_in); free(offsets); free(lengths); free(flat); free(h_out);
    free(hex_buf); free(portal);
    return 0;
}

