#include <assert.h>
#include <gmp.h>
#include <math.h>
#include <errno.h>
#include <error.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <asm/unistd.h>
#include <sys/ptrace.h>
#include <sys/reg.h>
#include <sys/wait.h>
#include <sys/uio.h>
#include <string.h>

#include "./ssv.h"
#include "trace.h"
#include <dlfcn.h>

#ifdef ENABLE_VERBOSE
#define VPRINT(fmt, ...) fprintf(stderr, "%s: " fmt, __func__, __VA_ARGS__)
#define VERBOSE(stmt)                                                          \
  { stmt; }
#else
#define VPRINT(...)
#define VERBOSE(stmt)
#endif

void pstatus(int status) {
  if (WIFEXITED(status)) printf("EXITED:%d ", WEXITSTATUS(status));
  if (WIFSIGNALED(status)) {
    printf("SIGNALED:%d:%d ",WTERMSIG(status),WCOREDUMP(status));
  }
  if (WIFSTOPPED(status)) printf("STOPPED:%d ", WSTOPSIG(status));
  if (WIFCONTINUED(status)) printf("CONTINUED ");
  printf("\n");
}

const char * stateStr(enum STATE s) {
  switch (s) {
  case INIT:
    return "INIT";
  case RUN:
    return "RUN";
  case EXIT:
    return "EXIT";
  case PANIC:
    return "PANIC";
  default:
    return "ERROR -- UNKOWN STATE\n";
  }
}

void 
processStep(struct ProcessDesc *pd) 
{
  pid_t id;

  VPRINT("%s: %ld:0x%x: %s\n",  __func__,  
	 pd->count, pd->status, stateStr(pd->state));

  if (ptrace(PTRACE_SYSEMU_SINGLESTEP, pd->pid, 0, 0) < 0) {
    pd->state=PANIC; return;
  }

  if ((id = waitpid(-1, &(pd->pid_status), 0)) < 0) {
    pd->state = PANIC; return;
  }

  VERBOSE(pstatus(pd->pid_status));

  if (id != pd->pid) {pd->state = PANIC; return;}

  if (!WIFSTOPPED(pd->pid_status)) {pd->state=PANIC; return;}

  pd->count++;

}

void 
processUpdate(struct ProcessDesc *pd) {
    switch (WSTOPSIG(pd->pid_status)) {
    case SIGTRAP:
      pd->state=RUN;
      break;
    case SIGTRAP | 0x80: // system call
      pd->nr = ptrace(PTRACE_PEEKUSER, pd->pid, 8 * ORIG_RAX, 0);
      if (errno) {pd->state=PANIC; break;}
      switch (pd->nr) {
      case __NR_exit: pd->state=EXIT; break;
      default:
	fprintf(stderr, "%ld: NYI: SYSCALL: %ld\n", pd->count, pd->nr);
	pd->state = PANIC;
      }
      break;
    default:
      pd->state=PANIC;
    }
}

void processInit(struct ProcessDesc *pd, int argc, char **argv)
{
  pd->count=0; pd->state = INIT;
  if ((pd->pid = asc_ptrace_start(argc, argv)) < 0) pd->state = PANIC;
  else pd->state=RUN;
}

ssv_t *
processGetState(struct ProcessDesc *pd, ssv_t *s)
{
  return asc_ssv_gather(s, pd->pid);
}

FILE *
openTrace(char *basename)
{
  char ext[]=".trc";
  int n=strlen(basename);
  char *path=malloc(n+strlen(ext)+1);
  FILE *f;
  
  strcpy(path,basename);
  strcpy(&path[n],ext);
  f=fopen(path, "w+");
  
  if (f==NULL) perror(path);
  return f;
}

struct TraceGlobals {
  char *actionFile;
  void *actionLibHdl;
  struct TraceAction action;
} globals;


int
loadTraceAction(void)
{
  char *error;
  int (*setup)(struct TraceAction *);
  
  globals.actionLibHdl = dlopen("/Users/jonathan/Work/Projects/PSML/apps/counters/RawLog.so", RTLD_LAZY);
  if (!globals.actionLibHdl)  {
    fprintf(stderr, "0. %s\n", dlerror());
    return -1;
  }
  
  setup = dlsym(globals.actionLibHdl, "SetupAction");
  if ((error = dlerror()) != NULL) {
    fprintf(stderr, "1. %s\n", error);
    return -1;
  }  
  return setup(&globals.action);
}

void
usage(void)
{
  fprintf(stderr, "trace <action> [action args] -- pgm [pgm args]\n");
}

int
main(int argc, char *argv[])
{
  struct ProcessDesc pd;
  ssv_t *state=NULL;
  size_t count=0;

  // load the trace action and obtain a pointer to the trace action object
  if (loadTraceAction()<0) {
    usage();
    fprintf(stderr, "ERROR: Unable to load action");
    return -1;
  }
  
  // setup process
  processInit(&pd, argc-1, argv+1);

  VPRINT("START: %s\n",stateStr(info.state));
 
  // call action specific init
  globals.action.init(&globals.action, argc, argv, &pd);

  // get initial state of process prior to first instruction execution
  state = processGetState(&pd, state); count++;

  // call action on initial state
  globals.action.action(&globals.action, state, &pd);

  while (pd.state == RUN) {
    processStep(&pd); if (pd.state == PANIC) break;
    
    // one step done get state of process
    state = processGetState(&pd, state); count++;

    // invoke action on the current state
    globals.action.action(&globals.action, state, &pd);
    
    processUpdate(&pd);
  }

  // Do Action specific cleanup passing last state of process as well
  globals.action.cleanup(&globals.action, state, &pd);

  // Do common cleanup
  asc_ssv_free(state);
  dlclose(globals.actionLibHdl);
  
  printf("EXIT: %ld: %s\n", pd.count, stateStr(pd.state));
  printf("states gathered:%ld\n", count);
  if (pd.state == PANIC) fprintf(stderr, "PANIC: run aborted\n");
  
  exit(errno);
}