#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/un.h>
#include <time.h>
#include <unistd.h>

enum {
  HEADER_SIZE = 4096,
  SUB_TAIL = 64,
  SUB_HEAD = 128,
  COMP_TAIL = 192,
  COMP_HEAD = 256,
  AGENT_PARKED = 320,
  CLIENT_PARKED = 384,
  RECORD_HEADER = 16,
  MAX_SUBMISSION = 16403,
  MAX_DETAIL = 65535,
  KIND_TRANSACTION = 1,
  FLAG_REPLY = 0x01,
};

struct completion {
  uint8_t code;
  uint64_t id;
  uint32_t length;
  uint8_t payload[MAX_DETAIL];
};

static int sock = -1;
static uint8_t *base;
static size_t mapping_length;
static uint64_t sub_cap;
static uint64_t comp_cap;

static void die(const char *message) {
  fprintf(stderr, "%s\n", message);
  exit(1);
}

static void die_errno(const char *what) {
  perror(what);
  exit(1);
}

static uint32_t get_u32(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 uint64_t get_u64(const uint8_t *p) {
  return (uint64_t)get_u32(p) | (uint64_t)get_u32(p + 4) << 32;
}

static void put_u32(uint8_t *p, uint32_t value) {
  for (int i = 0; i < 4; i++) p[i] = (uint8_t)(value >> (8 * i));
}

static void put_u64(uint8_t *p, uint64_t value) {
  for (int i = 0; i < 8; i++) p[i] = (uint8_t)(value >> (8 * i));
}

static uint64_t *word64(size_t offset) { return (uint64_t *)(base + offset); }

static uint32_t *word32(size_t offset) { return (uint32_t *)(base + offset); }

static uint64_t record_size(uint64_t length) {
  return (RECORD_HEADER + length + 15) & ~(uint64_t)15;
}

static int valid_capacity(uint64_t capacity) {
  return capacity >= 16 && (capacity & (capacity - 1)) == 0;
}

static void connect_agent(const char *path) {
  struct sockaddr_un address = {.sun_family = AF_UNIX};
  size_t path_length = strlen(path);
  if (path_length >= sizeof address.sun_path) die("socket path is too long");
  memcpy(address.sun_path, path, path_length + 1);
  sock = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
  if (sock < 0) die_errno("socket");
  if (connect(sock, (struct sockaddr *)&address, sizeof address) != 0)
    die_errno("connect");

  struct timeval limit = {.tv_sec = 2};
  if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &limit, sizeof limit) != 0 ||
      setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &limit, sizeof limit) != 0)
    die_errno("setsockopt");
  const uint8_t hello[5] = {'Z', 'C', 'I', '1', 0};
  size_t sent = 0;
  while (sent < sizeof hello) {
    ssize_t n = send(sock, hello + sent, sizeof hello - sent, MSG_NOSIGNAL);
    if (n < 0 && errno == EINTR) continue;
    if (n < 0) die_errno("send hello");
    if (n == 0) die("the agent closed the socket during the handshake");
    sent += (size_t)n;
  }

  uint8_t reply[9];
  size_t received = 0;
  int memfd = -1;
  while (received < sizeof reply) {
    union {
      struct cmsghdr align;
      char bytes[CMSG_SPACE(sizeof(int))];
    } control;
    struct iovec chunk = {.iov_base = reply + received,
                          .iov_len = sizeof reply - received};
    struct msghdr message = {
        .msg_iov = &chunk,
        .msg_iovlen = 1,
        .msg_control = control.bytes,
        .msg_controllen = sizeof control.bytes,
    };
    ssize_t n = recvmsg(sock, &message, MSG_CMSG_CLOEXEC);
    if (n < 0 && errno == EINTR) continue;
    if (n < 0) die_errno("receive handshake reply");
    if (n == 0) die("the agent closed the socket during the handshake");
    if ((message.msg_flags & (MSG_CTRUNC | MSG_TRUNC)) != 0)
      die("truncated handshake reply");
    for (struct cmsghdr *c = CMSG_FIRSTHDR(&message); c != NULL;
         c = CMSG_NXTHDR(&message, c)) {
      if (c->cmsg_level == SOL_SOCKET && c->cmsg_type == SCM_RIGHTS) {
        if (c->cmsg_len != CMSG_LEN(sizeof(int)) || memfd >= 0)
          die("unexpected file descriptors in handshake reply");
        memcpy(&memfd, CMSG_DATA(c), sizeof memfd);
      }
    }
    received += (size_t)n;
  }
  if (reply[0] != 0) {
    fprintf(stderr, "the agent refused the session with code %u\n",
            (unsigned)reply[0]);
    exit(1);
  }
  if (memfd < 0) die("the handshake reply carried no shared memory");
  sub_cap = get_u32(reply + 1);
  comp_cap = get_u32(reply + 5);
  if (!valid_capacity(sub_cap) || !valid_capacity(comp_cap))
    die("invalid ring capacity");

  mapping_length = (size_t)(HEADER_SIZE + sub_cap + comp_cap);
  struct stat status;
  if (fstat(memfd, &status) != 0) die_errno("fstat");
  if (status.st_size < 0 || (uint64_t)status.st_size < mapping_length)
    die("shared memory is too small");
  void *mapped = mmap(NULL, mapping_length, PROT_READ | PROT_WRITE, MAP_SHARED,
                      memfd, 0);
  if (mapped == MAP_FAILED) die_errno("mmap");
  base = mapped;
  close(memfd);
  if (memcmp(base, "ZCR1", 4) != 0 || get_u32(base + 4) != 1 ||
      get_u32(base + 8) != sub_cap || get_u32(base + 12) != comp_cap)
    die("unexpected shared memory header");

  int flags = fcntl(sock, F_GETFL);
  if (flags < 0 || fcntl(sock, F_SETFL, flags | O_NONBLOCK) != 0)
    die_errno("fcntl");
}

static int submit(uint64_t id, const uint8_t *payload, uint32_t length) {
  uint8_t *ring = base + HEADER_SIZE;
  uint64_t tail = __atomic_load_n(word64(SUB_TAIL), __ATOMIC_RELAXED);
  uint64_t head = __atomic_load_n(word64(SUB_HEAD), __ATOMIC_ACQUIRE);
  if (((head | tail) & 15) != 0 || head > tail || tail - head > sub_cap)
    return -1;
  uint64_t size = record_size(length);
  uint64_t offset = tail & (sub_cap - 1);
  uint64_t pad = size > sub_cap - offset ? sub_cap - offset : 0;
  if (pad + size > sub_cap - (tail - head)) return -1;
  if (pad != 0) {
    memset(ring + offset, 0, RECORD_HEADER);
    put_u32(ring + offset, (uint32_t)(pad - RECORD_HEADER));
    offset = 0;
  }
  uint8_t *record = ring + offset;
  put_u32(record, length);
  record[4] = KIND_TRANSACTION;
  record[5] = FLAG_REPLY;
  record[6] = 0;
  record[7] = 0;
  put_u64(record + 8, id);
  memcpy(record + RECORD_HEADER, payload, length);
  memset(record + RECORD_HEADER + length, 0, size - RECORD_HEADER - length);
  __atomic_store_n(word64(SUB_TAIL), tail + pad + size, __ATOMIC_RELEASE);

  __atomic_thread_fence(__ATOMIC_SEQ_CST);
  if (__atomic_load_n(word32(AGENT_PARKED), __ATOMIC_ACQUIRE) == 1 &&
      __atomic_exchange_n(word32(AGENT_PARKED), 0, __ATOMIC_SEQ_CST) == 1) {
    const uint8_t doorbell = 0;
    ssize_t rung = send(sock, &doorbell, 1, MSG_NOSIGNAL | MSG_DONTWAIT);
    (void)rung;
  }
  return 0;
}

static int read_completion(struct completion *out) {
  const uint8_t *ring = base + HEADER_SIZE + sub_cap;
  for (;;) {
    uint64_t head = __atomic_load_n(word64(COMP_HEAD), __ATOMIC_RELAXED);
    uint64_t tail = __atomic_load_n(word64(COMP_TAIL), __ATOMIC_ACQUIRE);
    uint64_t used = tail - head;
    if (((head | tail) & 15) != 0 || head > tail || used > comp_cap) return -1;
    if (used == 0) return 0;
    uint64_t offset = head & (comp_cap - 1);
    uint64_t remaining = comp_cap - offset;
    const uint8_t *record = ring + offset;
    uint32_t length = get_u32(record);
    if (record[6] != 0 || record[7] != 0) return -1;
    if (record[4] == 0) {
      if (length != remaining - RECORD_HEADER || remaining > used ||
          record[5] != 0 || get_u64(record + 8) != 0)
        return -1;
      __atomic_store_n(word64(COMP_HEAD), head + remaining, __ATOMIC_RELEASE);
      continue;
    }
    uint64_t size = record_size(length);
    if (record[4] != KIND_TRANSACTION || length > MAX_DETAIL ||
        size > remaining || size > used)
      return -1;
    out->code = record[5];
    out->id = get_u64(record + 8);
    out->length = length;
    memcpy(out->payload, record + RECORD_HEADER, length);
    __atomic_store_n(word64(COMP_HEAD), head + size, __ATOMIC_RELEASE);
    return 1;
  }
}

static int socket_closed(void) {
  uint8_t bytes[128];
  for (;;) {
    ssize_t n = recv(sock, bytes, sizeof bytes, 0);
    if (n > 0) continue;
    if (n == 0) return 1;
    if (errno == EINTR) continue;
    return errno != EAGAIN && errno != EWOULDBLOCK;
  }
}

static int64_t now_ms(void) {
  struct timespec now;
  if (clock_gettime(CLOCK_MONOTONIC, &now) != 0) die_errno("clock_gettime");
  return (int64_t)now.tv_sec * 1000 + now.tv_nsec / 1000000;
}

static int wait_for(uint64_t id, struct completion *out, int timeout_ms) {
  int64_t deadline = now_ms() + timeout_ms;
  for (;;) {
    int found = read_completion(out);
    if (found < 0) return -1;
    if (found == 1) {
      if (out->id == id) return 1;
      continue;
    }
    int64_t left = deadline - now_ms();
    if (left <= 0) return 0;

    __atomic_store_n(word32(CLIENT_PARKED), 1, __ATOMIC_RELEASE);
    __atomic_thread_fence(__ATOMIC_SEQ_CST);
    if (__atomic_load_n(word64(COMP_TAIL), __ATOMIC_ACQUIRE) !=
        __atomic_load_n(word64(COMP_HEAD), __ATOMIC_RELAXED)) {
      __atomic_store_n(word32(CLIENT_PARKED), 0, __ATOMIC_RELEASE);
      continue;
    }
    struct pollfd watch = {.fd = sock, .events = POLLIN};
    int ready = poll(&watch, 1, (int)left);
    __atomic_store_n(word32(CLIENT_PARKED), 0, __ATOMIC_RELEASE);
    if (ready < 0 && errno != EINTR) return -1;
    if (ready > 0 &&
        ((watch.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0 || socket_closed())) {
      while ((found = read_completion(out)) == 1) {
        if (out->id == id) return 1;
      }
      return -1;
    }
  }
}

int main(int argc, char **argv) {
  if (argc != 3) {
    fprintf(stderr, "usage: %s SOCKET_PATH TRANSACTION_FILE\n", argv[0]);
    return 2;
  }
  static uint8_t transaction[MAX_SUBMISSION + 1];
  FILE *file = fopen(argv[2], "rb");
  if (file == NULL) die_errno(argv[2]);
  size_t length = fread(transaction, 1, sizeof transaction, file);
  if (ferror(file)) die_errno("read transaction");
  if (fclose(file) != 0) die_errno("close transaction");
  if (length == 0 || length > MAX_SUBMISSION)
    die("the transaction payload must hold 1 to 16403 bytes");

  connect_agent(argv[1]);
  const uint64_t id = 1;
  if (submit(id, transaction, (uint32_t)length) != 0)
    die("the submission ring is full or invalid");

  static struct completion reply;
  int found = wait_for(id, &reply, 5000);
  if (found == 0) die("no completion within 5 seconds; check the signature before retrying");
  if (found < 0)
    die("the session ended or its ring is invalid; check the signature before retrying");
  printf("code %u\n", (unsigned)reply.code);
  if (reply.length > 0) {
    fwrite(reply.payload, 1, reply.length, stdout);
    fputc('\n', stdout);
  }
  munmap(base, mapping_length);
  close(sock);
  return reply.code == 0 ? 0 : 1;
}