#define _GNU_SOURCE
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <sys/syscall.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <errno.h>
#include <signal.h>
#include <stdlib.h>
#include "kem.h"

#define EVE_CPU 0
#define BOB_CPU 3

#define N LWE_N
#define K MODULE_RANK

static int pipe_eve_to_bob[2];
static int pipe_bob_to_eve[2];
static int pipe_bob_to_judge[2];
static int pipe_eve_to_judge[2];

static void writeall(int fd, const void *x, long long n) {
  while (n > 0) {
    ssize_t w = n < 65536 ? n : 65536;
    w = write(fd, x, w);
    if (w < 1) _exit(111);
    x = w + (char *)x;
    n -= w;
  }
}

static void readall(int fd, void *x, long long n) {
  while (n > 0) {
    ssize_t r = n < 65536 ? n : 65536;
    r = read(fd, x, r);
    if (r < 1) _exit(111);
    x = r + (char *)x;
    n -= r;
  }
}

static long long ticks(void) {
  unsigned long long result;
  asm volatile(".byte 15;.byte 49;shlq $32,%%rdx;orq %%rdx,%%rax"
    : "=a" (result) ::  "%rdx");
  return result;
}

// ----- bob

uint8_t bob_pk[PUBLICKEY_BYTES];
uint8_t bob_sk_bytes[KEM_SECRETKEY_BYTES];

static void bob_keypair(void) {
  crypto_kem_keypair(bob_pk, bob_sk_bytes);
  printf("pk ");
  for (long long i = 0; i < PUBLICKEY_BYTES; ++i) printf("%02x", bob_pk[i]);
  printf("\n");
  fflush(stdout);
}

uint8_t bob_ct[CIPHERTEXT_BYTES];
uint8_t bob_ss[CRYPTO_BYTES];

static void bob(void) {
  bob_keypair();
  writeall(pipe_bob_to_judge[1], (void *)&bob_sk_bytes, sizeof bob_sk_bytes);
  writeall(pipe_bob_to_eve[1], (void *)bob_pk, sizeof bob_pk);

  for (;;) {
    long long t = ticks();
    writeall(pipe_bob_to_eve[1], (void *)&t, sizeof t);
    readall(pipe_eve_to_bob[0], (void *)bob_ct, sizeof bob_ct);
    crypto_kem_decap(bob_ss, bob_sk_bytes, bob_pk, bob_ct);
  }
}

// ----- eve

static long long timestamp(void) {
  long long result;
  readall(pipe_bob_to_eve[0], (void *)&result, sizeof result);
  return result;
}

static void send_ciphertext(const uint8_t *ct) { writeall(pipe_eve_to_bob[1], (void *)ct, CIPHERTEXT_BYTES); }

public_key bob_pk_internal;

static void eve_unpackpk(void) {
  load_from_string_pk(&bob_pk_internal, bob_pk);
}

static void int64_minmax(long long *a, long long *b) {
  int64_t ab = *b ^ *a, c = *b - *a;
  c = ((((c ^ *b) & ab) ^ c) >> 63) & ab;
  *a ^= c;
  *b ^= c;
}

static void sort(long long *x, long long n) {
  long long top = 1, p, q, r, i;
  if (n < 2) return;
  while (top < n - top) top += top;
  for (i = 0; i < n; ++i) x[i] ^= 0x8000000000000000;
  for (p = top; p > 0; p >>= 1) {
    for (i = 0; i < n - p; ++i)
      if (!(i & p)) int64_minmax(&x[i], &x[i + p]);
    i = 0;
    for (q = top; q > p; q >>= 1) {
      for (; i < n - q; ++i) {
        if (!(i & p)) {
          long long a = x[i + p];
          for (r = q; r > p; r >>= 1) int64_minmax(&a, &x[i + r]);
          x[i + p] = a;
        }
      }
    }
  }
  for (i = 0; i < n; ++i) x[i] ^= 0x8000000000000000;
}

#define IQM_SCALE 16

// assumes x is sorted
static long long scaled_iqm(long long *x, long long n) {
  long long left = 0;
  long long right = n;
  long long total = 0;
  while ((right - left - 2) * 2 >= n) {
    ++left;
    --right;
  }
  for (long long i = left; i < right; ++i) total += x[i];
  return (total * IQM_SCALE) / (right - left);
}

int eve_guess[K * N];
int eve_guess2[K * N];
long long eve_confidence[K * N];

static int sk_lookup(secret_key *s, long long segment, long long pos)
{
  sppoly *op = s->sp_vec + segment;

  for (long long j = 0; j < op->neg_start; ++j)
    if (op->sx[j] == pos) return 1;
  for (long long j = op->neg_start; j < op->cnt; ++j)
    if (op->sx[j] == pos) return -1;
  return 0;
}

#define FUZZ 3

static int eve_happy(void) {
  uint8_t res[K * N];
  secret_key eve_sk_internal;

  // XXX: can use lattice attacks to speed up search

  for (long long search = 0;search < (1LL << (2 * FUZZ));++search) {
    long long searchbits = search;
    int happy = 1;

    for (long long i = 0;i < K * N;++i) {
      res[i] = eve_guess[i];
      if (eve_guess2[i] != eve_guess[i]) {
        if (searchbits & 1) res[i] = eve_guess2[i];
        searchbits >>= 1;
      }
    }

    long long total = 0;

    for (long long i = 0; i < K; ++i) {
      long long count = 0;
      for (long long j = 0;j < N;++j)
        count += res[i*N+j] != 0;
      total += count;
    }

    if (total != HS) continue;

    for (long long i = 0; i < K; ++i) {
      long long count = 0;
      for (long long j = 0;j < N;++j)
        count += res[i*N+j] != 0;
      eve_sk_internal.sp_vec[i].cnt = count;
      eve_sk_internal.sp_vec[i].sx = calloc(count, 1);
      if (!eve_sk_internal.sp_vec[i].sx) abort();
      eve_sk_internal.sp_vec[i].neg_start = convToIdx(eve_sk_internal.sp_vec[i].sx, count, res + i * N, N);
    }

    printf("trying %lld\n",search);
    fflush(stdout);
    polyvec y = bob_pk_internal.b;
    matrix_vec_mult_add(&y, bob_pk_internal.A, eve_sk_internal.sp_vec, 0);
    for (long long i = 0; i < K; ++i)
      free(eve_sk_internal.sp_vec[i].sx);

    for (long long i = 0; i < K; ++i)
      for (long long j = 0; j < N; ++j) {
        int c = y.vec[i].coeffs[j];
        c &= 65535;
        if (c > 32768) c -= 65536;
        if (c < -256) happy = 0;
        if (c > 256) happy = 0;
      }

    if (happy) {
      for (long long i = 0;i < K * N;++i)
        eve_guess[i] = (int) (char) res[i];
      return 1;
    }
  }

  return 0;
}

#define MULTIPLIERS 3
long long multipliers[MULTIPLIERS] = {0,50,50};
long long Csel[MULTIPLIERS] = {0,3,32-3};

#define TIMINGS 16
#define BATCH (N * K * MULTIPLIERS)

long long order[BATCH];

uint8_t eve_ct[CIPHERTEXT_BYTES];
uint16_t eve_whichsegment[BATCH];
uint8_t eve_whichpos[BATCH];
uint8_t eve_whichm[BATCH];
long long timings[TIMINGS + 1];
long long results[BATCH];

#define SAMPLES 1024
long long data[K][N][MULTIPLIERS][SAMPLES];

static void eve_main(void) {
  ciphertext eve_ct_internal;
  long long last_timestamp, this_timestamp;
  long long stats_dec = 0;

  readall(pipe_bob_to_eve[0], (void *)bob_pk, sizeof bob_pk);
  eve_unpackpk();
  last_timestamp = timestamp();

  for (long long sample = 0; sample < SAMPLES;) {
    randombytes((void *)order, sizeof order);
    for (int b = 0; b < BATCH; ++b) {
      order[b] &= 0x0fffffffffffffff;
      order[b] = (order[b] / BATCH) * BATCH + b;
    }
    sort(order, BATCH);
    for (int b = 0; b < BATCH; ++b) {
      int index = order[b] % BATCH;
      eve_whichm[b] = index % MULTIPLIERS;
      index /= MULTIPLIERS;
      eve_whichpos[b] = index % N;
      eve_whichsegment[b] = index / N;
    }

    for (int b = 0; b < BATCH; ++b) {
      int segment = eve_whichsegment[b];
      int pos = eve_whichpos[b];
      int m = multipliers[eve_whichm[b]];

      memset(&eve_ct_internal, 0, sizeof(ciphertext));
      for (long long i = 0; i < MODULE_RANK; ++i)
        for (long long j = 0; j < LWE_N; ++j)
          eve_ct_internal.c1.vec[i].coeffs[j] = m * (i == segment) * (j + pos == N - 1);
      for (long long i = 0; i < N; ++i)
        eve_ct_internal.c2.coeffs[i] = i == N-1 ? Csel[eve_whichm[b]] : 0;
      save_to_string(eve_ct, &eve_ct_internal);

      for (int t = 0; t < TIMINGS; ++t) {
        send_ciphertext(eve_ct);
        this_timestamp = timestamp();
        timings[t] = this_timestamp - last_timestamp;
        last_timestamp = this_timestamp;
      }
      for (int t = 0; t < TIMINGS; ++t) timings[t] &= 0xffffffff;
      sort(timings, TIMINGS);
      results[b] = timings[TIMINGS / 2];
      stats_dec += TIMINGS;
    }

    for (int j = 0; j < BATCH; ++j) {
      int segment = eve_whichsegment[j];
      int pos = eve_whichpos[j];
      int whichm = eve_whichm[j];
      int insert = 0;
      while (insert < sample && data[segment][pos][whichm][insert] < results[j]) ++insert;
      for (int i = sample; i > insert; --i) data[segment][pos][whichm][i] = data[segment][pos][whichm][i - 1];
      data[segment][pos][whichm][insert] = results[j];
    }

    ++sample;

    long long iqm[K][N][MULTIPLIERS];

    for (int segment = 0; segment < K; ++segment) {
      for (int pos = 0; pos < N; ++pos) {

        for (int whichm = 0; whichm < MULTIPLIERS; ++whichm) iqm[segment][pos][whichm] = scaled_iqm(data[segment][pos][whichm], sample);

        // XXX: ad-hoc decision, assuming MULTIPLIERS is 3 etc.
        long long offset1 = llabs(iqm[segment][pos][1] - iqm[segment][pos][0]);
        long long offset2 = llabs(iqm[segment][pos][2] - iqm[segment][pos][0]);
        eve_guess[segment * N + pos] = offset1 > offset2 ? 1 : -1;
        eve_confidence[segment * N + pos] = offset1 > offset2 ? offset1 : offset2;
      }
    }

    for (int j = 0;j < K * N;++j)
      eve_guess2[j] = eve_guess[j];
    for (int j = 0;j < K * N;++j)
      eve_confidence[j] = eve_confidence[j]*K*N + j;
    sort(eve_confidence,K*N);
    for (int j = 0;j < K * N - HS;++j)
      eve_guess[eve_confidence[j]%(K*N)] = 0;
    for (int j = 0;j < K * N - HS - FUZZ;++j)
      eve_guess2[eve_confidence[j]%(K*N)] = 0;
    for (int j = K * N - HS;j < K * N - HS + FUZZ;++j)
      eve_guess2[eve_confidence[j]%(K*N)] = 0;

    for (int segment = 0; segment < K; ++segment) {
      for (int pos = 0; pos < N; ++pos) {
        printf("pos %d", segment * N + pos);
        printf(" guess %d", eve_guess[segment * N + pos]);
        printf(" guess2 %d", eve_guess2[segment * N + pos]);
        printf(" samples/m %lld", sample);
        printf(" iqms");
        for (int whichm = 0; whichm < MULTIPLIERS; ++whichm) printf(" %lld,%lld:%lld", multipliers[whichm], Csel[whichm], iqm[segment][pos][whichm]);
        for (int whichm = 0; whichm < MULTIPLIERS; ++whichm) printf(" %lld", iqm[segment][pos][whichm]-iqm[segment][pos][0]);
        printf("\n");
      }
    }

    printf("totaldec %lld", stats_dec);
    printf(" samples/m/pos %lld", sample);
    printf(" limit %d", SAMPLES);
    printf("\n");
    fflush(stdout);
    if (eve_happy()) {
      printf("eve declares success\n");
      fflush(stdout);
      return;
    }
  }
  printf("eve giving up ... for now\n");
}

static void eve(void) {
  eve_main();
  writeall(pipe_eve_to_judge[1], (void *)&eve_guess, sizeof eve_guess);
}

// ----- judge

static void judge(void) {
  int match = 1;
  printf("checking whether eve's secret key matches bob's secret key...\n");

  secret_key bob_sk_internal;

  load_from_string_sk(&bob_sk_internal, bob_sk_bytes);

  for (long long i = 0; i < K; ++i)
    for (long long j = 0; j < N; ++j) {
      int c = eve_guess[i*N+j];
      int d = sk_lookup(&bob_sk_internal,i,j);
      printf("pos %lld eve %d bob %d matches", i * N + j, c, d);
      if (c == d)
        printf(" True");
      else {
        printf(" False");
        match = 0;
      }
      printf("\n");
    }
  if (match)
    printf("yes, eve succeeded\n");
  else
    printf("no, eve failed\n");
  fflush(stdout);
}

int main() {
  pid_t pid_bob;
  pid_t pid_eve;
  cpu_set_t mask;

  if (pipe(pipe_eve_to_bob) == -1) {
    fprintf(stderr, "pipe failed: %s\n", strerror(errno));
    return 111;
  }
  if (pipe(pipe_bob_to_eve) == -1) {
    fprintf(stderr, "pipe failed: %s\n", strerror(errno));
    return 111;
  }
  if (pipe(pipe_bob_to_judge) == -1) {
    fprintf(stderr, "pipe failed: %s\n", strerror(errno));
    return 111;
  }
  if (pipe(pipe_eve_to_judge) == -1) {
    fprintf(stderr, "pipe failed: %s\n", strerror(errno));
    return 111;
  }

  switch (pid_bob = fork()) {
    case -1:
      fprintf(stderr, "fork failed: %s\n", strerror(errno));
      return 111;
    case 0:
      CPU_ZERO(&mask);
      CPU_SET(BOB_CPU, &mask);
      sched_setaffinity(0, sizeof mask, &mask);
      close(pipe_eve_to_bob[1]);
      close(pipe_bob_to_eve[0]);
      close(pipe_bob_to_judge[0]);
      close(pipe_eve_to_judge[0]);
      close(pipe_eve_to_judge[1]);
      bob();
      return 0;
    default:
      close(pipe_eve_to_bob[0]);
      close(pipe_bob_to_eve[1]);
      close(pipe_bob_to_judge[1]);
  }

  readall(pipe_bob_to_judge[0], (void *)&bob_sk_bytes, sizeof bob_sk_bytes);
  close(pipe_bob_to_judge[0]);

  switch (pid_eve = fork()) {
    case -1:
      fprintf(stderr, "fork failed: %s\n", strerror(errno));
      return 111;
    case 0:
      CPU_ZERO(&mask);
      CPU_SET(EVE_CPU, &mask);
      sched_setaffinity(0, sizeof mask, &mask);
      close(pipe_eve_to_judge[0]);
      eve();
      return 0;
    default:
      close(pipe_eve_to_bob[1]);
      close(pipe_bob_to_eve[0]);
      close(pipe_eve_to_judge[1]);
  }

  readall(pipe_eve_to_judge[0], (void *)&eve_guess, sizeof eve_guess);
  judge();

  waitpid(pid_bob, 0, 0);
  waitpid(pid_eve, 0, 0);
  return 0;
}