/* radix-sort.cu — bend form G: batched 64-bit key sort via NVIDIA CUB.
 *
 * Wraps cub::DeviceRadixSort::SortKeys on a u64 key stream. One-shot
 * sort kernel; the dispatcher reserves op_id slots for descending
 * order, u32 keys, & key-value pairs but only 0x01 (sort-u64-asc) is
 * wired day-1.
 *
 * CATALOG.md form G entry: 1.4 G keys/sec on Titan; 20-50x over CPU
 * merge sort; 257x over Intel Xeon Phi for scan.
 *
 * Wire (BSRT input):
 *   "BSRT" (4 B magic)
 *   u32 op_id            (0x01 = sort-u64-asc)
 *   u32 n_items
 *   u64[n_items]         (little-endian keys)
 *
 * Wire (BSRR output):
 *   "BSRR" (4 B magic)
 *   u32 status           (0 = ok)
 *   u32 n_items
 *   u64[n_items]         (sorted ascending)
 *
 * Reserved op slots (NOT day-1):
 *   0x02 sort-u64-desc
 *   0x03 sort-u32-asc / desc
 *   0x04 sort-key-value pairs
 *
 * Build:
 *   make radix-sort
 *
 * License: AGPLv3. CUB is BSD-3-Clause (NVIDIA), header-only, so the
 * binary is a thin wrapper around their device sort.
 */

#include <cuda_runtime.h>
#include <cub/device/device_radix_sort.cuh>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.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)); \
        return -1; \
    } \
} while (0)

/* Op codes. Only OP_SORT_U64_ASC is wired today; the rest hold their
 * slot so a future build can grow the dispatch without renumbering. */
enum {
    OP_SORT_U64_ASC  = 0x01,
    OP_SORT_U64_DESC = 0x02,
    OP_SORT_U32      = 0x03,   /* asc/desc selected by low bit later */
    OP_SORT_KV       = 0x04,   /* key-value pairs */
};

/* ── small endian helpers ─────────────────────────────────────────── */

static uint32_t rd32_le(const uint8_t *p) {
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8)
         | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void wr32_le(uint8_t *p, uint32_t v) {
    p[0] = v & 0xff; p[1] = (v >> 8) & 0xff;
    p[2] = (v >> 16) & 0xff; p[3] = (v >> 24) & 0xff;
}

static int read_file_to_buffer(const char *path, uint8_t **out_buf, size_t *out_len) {
    FILE *f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "open %s failed\n", 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 (!buf) { fclose(f); return -1; }
    if (fread(buf, 1, sz, f) != (size_t)sz) {
        free(buf); fclose(f); return -1;
    }
    fclose(f);
    *out_buf = buf;
    *out_len = sz;
    return 0;
}

static int write_buffer_to_file(const char *path, const uint8_t *buf, size_t len) {
    FILE *f = fopen(path, "wb");
    if (!f) { fprintf(stderr, "create %s failed\n", path); return -1; }
    if (fwrite(buf, 1, len, f) != len) { fclose(f); return -1; }
    fclose(f);
    return 0;
}

/* ── core: BSRT in → BSRR out ─────────────────────────────────────── */

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);

    uint8_t *in_buf = NULL;
    size_t in_len = 0;
    if (read_file_to_buffer(in_path, &in_buf, &in_len) != 0) return -1;

    if (in_len < 12 || memcmp(in_buf, "BSRT", 4) != 0) {
        fprintf(stderr, "bad BSRT magic / short input (%zu bytes)\n", in_len);
        free(in_buf); return -1;
    }
    uint32_t op_id = rd32_le(in_buf + 4);
    uint32_t n     = rd32_le(in_buf + 8);

    if (op_id != OP_SORT_U64_ASC) {
        fprintf(stderr, "unsupported op_id 0x%02x (only 0x01 sort-u64-asc wired)\n",
                op_id);
        free(in_buf); return -1;
    }

    size_t expected = 12 + (size_t)n * sizeof(uint64_t);
    if (in_len != expected) {
        fprintf(stderr, "BSRT payload size %zu != expected %zu (op=0x%02x n=%u)\n",
                in_len, expected, op_id, n);
        free(in_buf); return -1;
    }

    const uint64_t *h_keys_in = (const uint64_t*)(in_buf + 12);

    /* Allocate device buffers (input + output ping-pong CUB needs).      */
    uint64_t *d_keys_in = NULL;
    uint64_t *d_keys_out = NULL;
    size_t n_bytes = (size_t)n * sizeof(uint64_t);
    /* Always allocate at least 1 byte to keep cudaMalloc happy at n=0.   */
    size_t alloc_bytes = n_bytes ? n_bytes : 1;
    CUDA_CHECK(cudaMalloc(&d_keys_in,  alloc_bytes));
    CUDA_CHECK(cudaMalloc(&d_keys_out, alloc_bytes));
    if (n > 0) {
        CUDA_CHECK(cudaMemcpy(d_keys_in, h_keys_in, n_bytes,
                              cudaMemcpyHostToDevice));
    }
    clock_gettime(CLOCK_MONOTONIC, &t1);

    /* CUB DeviceRadixSort: two-call pattern. First call with
     * d_temp_storage=NULL returns the temp-storage size; second call
     * does the actual sort. End-bit defaults to sizeof(KeyT)*8 (64).   */
    void *d_temp = NULL;
    size_t temp_bytes = 0;
    cub::DeviceRadixSort::SortKeys(
        d_temp, temp_bytes,
        d_keys_in, d_keys_out,
        (int)n);
    CUDA_CHECK(cudaGetLastError());
    if (temp_bytes == 0) temp_bytes = 1;
    CUDA_CHECK(cudaMalloc(&d_temp, temp_bytes));
    clock_gettime(CLOCK_MONOTONIC, &t2);

    cub::DeviceRadixSort::SortKeys(
        d_temp, temp_bytes,
        d_keys_in, d_keys_out,
        (int)n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    clock_gettime(CLOCK_MONOTONIC, &t3);

    /* Pull sorted keys back, frame BSRR. */
    size_t out_size = 12 + n_bytes;
    uint8_t *out_buf = (uint8_t*)malloc(out_size ? out_size : 1);
    if (!out_buf) {
        cudaFree(d_keys_in); cudaFree(d_keys_out); cudaFree(d_temp);
        free(in_buf); return -1;
    }
    memcpy(out_buf, "BSRR", 4);
    wr32_le(out_buf + 4, 0);   /* status = ok */
    wr32_le(out_buf + 8, n);
    if (n > 0) {
        CUDA_CHECK(cudaMemcpy(out_buf + 12, d_keys_out, n_bytes,
                              cudaMemcpyDeviceToHost));
    }
    clock_gettime(CLOCK_MONOTONIC, &t4);

    cudaFree(d_keys_in);
    cudaFree(d_keys_out);
    cudaFree(d_temp);
    free(in_buf);

    int wr = write_buffer_to_file(out_path, out_buf, out_size);
    free(out_buf);

    double ld_ms     = (t1.tv_sec - t0.tv_sec) * 1e3 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
    double probe_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;
    double gkeys = (kern_ms > 0.0) ? ((double)n / kern_ms / 1e6) : 0.0;
    fprintf(stderr,
        "radix-sort op=0x%02x n=%u load=%.2f probe=%.2f kernel=%.3f write=%.2f ms (%.3f Gkeys/s)\n",
        op_id, n, ld_ms, probe_ms, kern_ms, write_ms, gkeys);
    return wr;
}

/* ── main: --binary one-shot, --daemon long-lived ─────────────────── */

int main(int argc, char **argv) {
    if (argc >= 2 && !strcmp(argv[1], "--daemon")) {
        int dev_count = 0;
        cudaGetDeviceCount(&dev_count);
        if (dev_count == 0) {
            fprintf(stderr, "no CUDA devices\n");
            return 1;
        }
        cudaFree(0);  /* eager 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;
            }
            if (!strncmp(line, "process-bin ", 12)) {
                char *args = line + 12;
                char *sep = strchr(args, ' ');
                if (!sep) {
                    fprintf(stdout, "error bad-args\n"); fflush(stdout);
                    continue;
                }
                *sep = 0;
                char *in_path  = args;
                char *out_path = sep + 1;
                int r = process_one_bin(in_path, out_path);
                fprintf(stdout, r == 0 ? "done %s\n" : "error process-failed\n",
                        out_path);
                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]);
    }

    fprintf(stderr,
        "usage: %s --binary <input.bin> <output.bin>\n"
        "       %s --daemon  (read commands on stdin)\n",
        argv[0], argv[0]);
    return 1;
}
