#include "test.h"

static gex_Client_t      myclient;
static gex_EP_t    myep;
static gex_TM_t myteam;
static gex_Segment_t     mysegment;

int myproc;
int numprocs;

uint8_t *myseg;      /* my segment */
uint8_t *peerseg;

uintptr_t maxlongreq, maxlongrep;
uintptr_t maxsz = 32*1024;

int peer;
int iters = 1;
size_t req_counter = 0, rep_counter = 0;
int wait_counter = 0;

/* Test handlers */
#define hidx_sleep_req_shorthandler   200
#define hidx_sleep_rep_shorthandler   201
#define hidx_done_req_shorthandler    202
#define hidx_flood_req_longhandler    203
#define hidx_flood_rep_longhandler    204

void sleep_req_shorthandler(gex_Token_t token, void *buf, size_t nbytes, gex_AM_Arg_t delay) {
  //MSG("sleep %d usec", delay);
  gex_Token_Info_t info;

  usleep(delay);
  gex_Token_Info(token, &info, GEX_TI_SRCRANK);
  uint8_t *peerseg = TEST_SEG(info.gex_srcrank);
  gex_AM_ReplyLong0(token, hidx_sleep_rep_shorthandler, buf, nbytes, peerseg,
                    GEX_EVENT_NOW, 0);
}

void sleep_rep_shorthandler(gex_Token_t token, void *buf, size_t nbytes) {
  //MSG("rep");
  wait_counter++;
}

void done_req_shorthandler(gex_Token_t token) {
  //MSG("done");
  iters = 0;
}

void flood_req_longhandler(gex_Token_t token, void *buf, size_t nbytes) {
  gex_Token_Info_t info;

  gex_Token_Info(token, &info, GEX_TI_SRCRANK);
  uint8_t *peerseg = TEST_SEG(info.gex_srcrank);
  req_counter++;
  gex_AM_ReplyLong0(token, hidx_flood_rep_longhandler, buf, nbytes, peerseg,
                    GEX_EVENT_NOW, 0);
}

void flood_rep_longhandler(gex_Token_t token, void *buf, size_t nbytes) {
  //MSG("flood_rep_longhandler");
  rep_counter++;
}

int main(int argc, char **argv)
{
  int delay = 1000000;
  gex_AM_Entry_t htable[] = {
    { hidx_sleep_req_shorthandler, sleep_req_shorthandler, GEX_FLAG_AM_REQREP|GEX_FLAG_AM_LONG, 1, NULL, NULL },
    { hidx_sleep_rep_shorthandler, sleep_rep_shorthandler, GEX_FLAG_AM_REQREP|GEX_FLAG_AM_LONG, 0, NULL, NULL },
    { hidx_done_req_shorthandler,  done_req_shorthandler,  GEX_FLAG_AM_REQUEST|GEX_FLAG_AM_SHORT, 0, NULL, NULL },
    { hidx_flood_req_longhandler,  flood_req_longhandler,  GEX_FLAG_AM_REQUEST|GEX_FLAG_AM_LONG, 0,  NULL, NULL },
    { hidx_flood_rep_longhandler,  flood_rep_longhandler,  GEX_FLAG_AM_REQREP|GEX_FLAG_AM_LONG, 0,   NULL, NULL },
  };

  if (argc > 3) {
    fprintf(stderr, "Use:\n%s (iters) (delay)\n", argv[0]);
    return 1;
  }
  if (argc >= 2) {
    iters = atoi(argv[1]);
  }
  if (argc == 3) {
    delay = atoi(argv[2]);
  }

  GASNET_Safe(gex_Client_Init(&myclient, &myep, &myteam, "testflood", &argc, &argv, 0));
  GASNET_Safe(gex_Segment_Attach(&mysegment, myteam, TEST_SEGSZ_REQUEST));
  GASNET_Safe(gex_EP_RegisterHandlers(myep, htable, sizeof(htable)/sizeof(gex_AM_Entry_t)));

  myproc = gex_TM_QueryRank(myteam);
  numprocs = gex_TM_QuerySize(myteam);

  //maxlongreq = MIN(maxsz, gex_AM_MaxRequestLong  (myteam,peer,NULL,0,0));
  //maxlongrep = MIN(maxsz, gex_AM_MaxReplyLong    (myteam,peer,GEX_EVENT_NOW,0,0));

  TEST_PRINT_CONDUITINFO();
  MSG0("iters=%u, delay=%u (usec)", iters, delay);

  myseg = TEST_MYSEG();
  switch (myproc) {
    case 0:
    case 1:
      peer = myproc ^ 1;
      break;
    default:
      peer = myproc & 1;
      break;
  }
  peerseg = TEST_SEG(peer);
  GASNET_Safe(gasnet_barrier(0,GASNET_BARRIERFLAG_UNNAMED));

  while (iters) {
    if (myproc == 0 || myproc == 1) {
      iters--;
      wait_counter = 0;
    }
    switch (myproc) {
      case 0:
      case 1:
        //gex_AM_RequestShort1(myteam, peer, hidx_sleep_req_shorthandler, 0, delay);
        gex_AM_RequestLong1(myteam, peer, hidx_sleep_req_shorthandler,
                            myseg, maxsz, peerseg, GEX_EVENT_NOW, 0, delay);

        while (!wait_counter) {
          gasnet_AMPoll();
        }
        MSG0("%d", iters);
        break;
      default:
        gex_AM_RequestLong0(myteam, peer, hidx_flood_req_longhandler,
                            myseg, maxsz, peerseg, GEX_EVENT_NOW, 0);
        break;
    }
    gasnet_AMPoll();
  }

  if (myproc == 0) {
    for (int i = 2; i < numprocs; i++) {
      gex_AM_RequestShort0(myteam, i, hidx_done_req_shorthandler, 0);
    }
  }

  GASNET_Safe(gasnet_barrier(0,GASNET_BARRIERFLAG_UNNAMED));

  MSG("long req/rep %lu/%lu", req_counter, rep_counter);
  gasnet_exit(test_errs > 0 ? 1 : 0);

  return 0;
}