#include "aes.h"
#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define BUFFER_SIZE 8192
#define PACKAGE_NUM 12
#define GLOG_ON false

#if GLOG_ON
#include <glog/logging.h>
// pipeline macros
#define PIPELINE_SEND(idx)                                                     \
  LOG(INFO) << "Proc" << rank << ": Isend started, idx:" << idx;               \
  MPI_Isend(send_buffer[(idx) % 2], BUFFER_SIZE, MPI_UINT8_T, rank + 1, idx,   \
            MPI_COMM_WORLD, &send_request);                                    \
  LOG(INFO) << "Proc" << rank << ": Isend finished, idx:" << idx;

#define PIPELINE_SEND_WAIT()                                                   \
  LOG(INFO) << "Proc" << rank << ": send wait started";                        \
  MPI_Wait(&send_request, MPI_STATUS_IGNORE);                                  \
  LOG(INFO) << "Proc" << rank << ": send wait finished";

#define PIPELINE_RECV(idx)                                                     \
  LOG(INFO) << "Proc" << rank << ": Irecv started, idx:" << idx;               \
  MPI_Irecv(recv_buffer[(idx) % 2], BUFFER_SIZE, MPI_UINT8_T, rank - 1, idx,   \
            MPI_COMM_WORLD, &recv_request);                                    \
  LOG(INFO) << "Proc" << rank << ": Irecv finished, idx:" << idx;

#define PIPELINE_RECV_WAIT()                                                   \
  LOG(INFO) << "Proc" << rank << ": recv wait started";                        \
  MPI_Wait(&recv_request, MPI_STATUS_IGNORE);                                  \
  LOG(INFO) << "Proc" << rank << ": recv wait finished";
#define CALCULATE(idx)                                                         \
  LOG(INFO) << "Proc" << rank << "cal " << idx << " started";                  \
  for (int cali = 0; cali < BUFFER_SIZE; cali += AES_ONCE) {                   \
    aes_cipher(&recv_buffer[(idx) % 2][cali], &send_buffer[(idx) % 2][cali],   \
               w);                                                             \
  }                                                                            \
  LOG(INFO) << "Proc" << rank << "cal " << idx << " finished";
// end pipeline macros

#else
// pipeline macros
#define PIPELINE_SEND(idx)                                                     \
  MPI_Isend(send_buffer[(idx) % 2], BUFFER_SIZE, MPI_UINT8_T, rank + 1, idx,   \
            MPI_COMM_WORLD, &send_request);

#define PIPELINE_SEND_WAIT() MPI_Wait(&send_request, MPI_STATUS_IGNORE);

#define PIPELINE_RECV(idx)                                                     \
  MPI_Irecv(recv_buffer[(idx) % 2], BUFFER_SIZE, MPI_UINT8_T, rank - 1, idx,   \
            MPI_COMM_WORLD, &recv_request);

#define PIPELINE_RECV_WAIT()                                                   \
  MPI_Wait(&recv_request, MPI_STATUS_IGNORE);                                  \

#define CALCULATE(idx)                                                         \
  for (int cali = 0; cali < BUFFER_SIZE; cali += AES_ONCE) {                   \
    aes_cipher(&recv_buffer[(idx) % 2][cali], &send_buffer[(idx) % 2][cali],   \
               w);                                                             \
  }
// end pipeline macros
#endif

// aes variables
#define AES_ONCE 16

uint8_t key[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
                 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
                 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f};
uint8_t *w;
uint8_t in[AES_ONCE], out[AES_ONCE];
// aes variables

int main(int argc, char *argv[]) {
  int rank, size;
  MPI_Request send_request, recv_request;
  uint8_t send_buffer[2][BUFFER_SIZE], recv_buffer[2][BUFFER_SIZE];
#if GLOG_ON
  google::InitGoogleLogging(argv[0]);
#endif
  // MPI init
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  // end MPI init

  // aes init
  w = aes_init(sizeof(key));
  aes_key_expansion(key, w);
  // end aes init

  if (rank == 0) {
    CALCULATE(0);
    PIPELINE_SEND(0);
    for (int i = 1; i < PACKAGE_NUM; ++i) {
      CALCULATE(i);
      PIPELINE_SEND_WAIT();
      PIPELINE_SEND(i);
    }
    PIPELINE_SEND_WAIT();
  } else if (rank == 1) {
    PIPELINE_RECV(0);
    PIPELINE_RECV_WAIT();
    CALCULATE(0);
    PIPELINE_SEND(0);
    PIPELINE_RECV(1);
    for (int i = 1; i < PACKAGE_NUM - 1; ++i) {
      PIPELINE_RECV_WAIT();
      PIPELINE_RECV(i + 1);
      CALCULATE(i);
      PIPELINE_SEND_WAIT();
      PIPELINE_SEND(i);
    }
    PIPELINE_RECV_WAIT();
    CALCULATE(PACKAGE_NUM - 1);
    PIPELINE_SEND_WAIT();
    PIPELINE_SEND(PACKAGE_NUM - 1);
    PIPELINE_SEND_WAIT();
  } else if (rank == 2) {
    PIPELINE_RECV(0);
    for (int i = 1; i < PACKAGE_NUM; ++i) {
      PIPELINE_RECV_WAIT();
      PIPELINE_RECV(i);
      CALCULATE(i - 1)
    }
    PIPELINE_RECV_WAIT();
    CALCULATE(PACKAGE_NUM - 1);
  }
  MPI_Finalize();
}