# include <stdio.h>
# include <stdlib.h>
# include <stdarg.h>
# include <stdint.h>
# include <pthread.h>

# undef OHANA_MEMORY
# include "ohana_allocate.h"

/* need an internal version that does not use ohana_memory functions */
# define FALSE 0 
# define TRUE 1
# define MAX(X,Y) ((X) > (Y) ? (X) : (Y))

# define OHANA_MEMMAGIC (uint32_t) 0xdeadbeef

typedef struct Memblock {
  uint32_t startblock;	      // endpost marker
  struct Memblock *prevBlock; // previously allocated memory
  struct Memblock *nextBlock; // next allocated memory
  size_t size;		      // size of memory
  const char *file;	      // file (re)allocated 
  const char *func;	      // func (re)allocated 
  int line;		      // line (re)allocated
  uint32_t endblock;	      // endpost marker
} Memblock;

static Memblock *lastBlock = NULL;

static pthread_mutex_t memBlockListMutex = PTHREAD_MUTEX_INITIALIZER;

void ohana_meminit () {
  return;
}

void ohana_memabort (char *format, ...) {
  va_list argp;  

  va_start (argp, format);
  vfprintf (stderr, format, argp);
  va_end (argp);
  abort();
}

void *ohana_malloc (const char *file, int line, const char *func, size_t Nelem, size_t esize) {

  void *ptr;		      // actual user memory allocated
  Memblock *new;	      // new memblock created to track the user memory
  size_t size;		      // total number of bytes requested (not items)

  Nelem = MAX (1, Nelem);
  size = Nelem * esize;

  // total size is : memblock + data + endpost
  new = (Memblock *) malloc (sizeof(Memblock) + size + sizeof(void *)); 
  if (new == NULL) ohana_memabort ("failed to allocate memory (%s, %d, %s)\n", file, line, func);

  // pointer to the start of the user memory
  ptr = (char *)(new + 1);

  // set the end-post values
  new->startblock = OHANA_MEMMAGIC;
  new->endblock = OHANA_MEMMAGIC;
  *(uint32_t *)(ptr + size) = OHANA_MEMMAGIC;

  // set the memory metadata
  new->size = size;
  new->file = file;
  new->line = line;
  new->func = func;

  // new memblock becomes the 'lastBlock':
  // lastBlock = new
  // new->prev = (new - 1)
  // (new - 1)->next = new

  // for memory allocated in the order m0, m1, m2
  // m0->next = m1, m1->next = m2, m2->next = NULL
  // m2->prev = m1, m1->prev = m0, m0->prev = NULL
  // lastBlock = m2

  // set up the pointers
  new->nextBlock = NULL;

  // protect the list during this operation
  pthread_mutex_lock(&memBlockListMutex);

  if (lastBlock) {
    lastBlock->nextBlock = new;
  }
  new->prevBlock = lastBlock;
  lastBlock = new;

  pthread_mutex_unlock(&memBlockListMutex);

  return (ptr);
}

void *ohana_realloc (const char *file, int line, const char *func, void *in, size_t Nelem, size_t esize) {

  void *ptr;		      // actual user memory allocated
  Memblock *old;	      // original memblock associated with user memory
  Memblock *new;	      // new memblock associated with user memory
  size_t size;

  // just allocate if not previously allocated
  if (!in) {
    ptr = ohana_malloc (file, line, func, Nelem, esize);
    return ptr;
  }

  // memblock of supplied pointer
  old = (Memblock *) in - 1;

  if (old->startblock != OHANA_MEMMAGIC) ohana_memabort ("corrupt memory (%s, %d, %s)\n", file, line, func);
  if (old->endblock != OHANA_MEMMAGIC) ohana_memabort ("corrupt memory (%s, %d, %s)\n", file, line, func);

  Nelem = MAX (1, Nelem);
  size = Nelem * esize;

  // requested same size as current allocation
  if (size == old->size) {
    return in;
  }

  pthread_mutex_lock(&memBlockListMutex);

  Memblock *nextBlock = old->nextBlock;
  Memblock *prevBlock = old->prevBlock;

  int isLast = (old == lastBlock);

  // ask for new memory
  // total size is : memblock + data + endpost
  new = (Memblock *) realloc (old, sizeof(Memblock) + size + sizeof(void *));
  if (new == NULL) ohana_memabort ("failed to reallocate memory (%s, %d, %s)\n", file, line, func);
  ptr = (char *) (new + 1);
  
  // give the realloc location:
  new->size = size; 
  new->file = file;
  new->line = line;
  new->func = func;

  // update the endpost (the others were set originally
  *(uint32_t *)(ptr + size) = OHANA_MEMMAGIC;

  // need to reset lastBlock in case we moved
  if (isLast) {
    lastBlock = new;
  }

  // need to adjust the neighbors' pointers:
  if (nextBlock) {
    nextBlock->prevBlock = new;
  }
  if (prevBlock) {
    prevBlock->nextBlock = new;
  }
  
  pthread_mutex_unlock(&memBlockListMutex);

  return (ptr);
}

void real_free (void *in) {

  if (!in) return;
  free (in);
  return;
}

// this is very slow.  should we speed this up by indexing on the ptr?
void ohana_free (const char *file, int line, const char *func, void *in) {

  Memblock *ref;

  if (!in) return;

  ref = (Memblock *) in - 1;

  pthread_mutex_lock(&memBlockListMutex);
  
  Memblock *nextBlock = ref->nextBlock;
  Memblock *prevBlock = ref->prevBlock;

  if (nextBlock) {
    nextBlock->prevBlock = prevBlock;
  }
  if (prevBlock) {
    prevBlock->nextBlock = nextBlock;
  }
  if (lastBlock == ref) {
    lastBlock = prevBlock;
  }
  
  pthread_mutex_unlock(&memBlockListMutex);
  
  free (ref);

  return;
}

void ohana_memcheck_func (int allmemory) {

  if (!lastBlock) {
    fprintf (stderr, "no memory allocated\n");
    return;
  }

  Memblock *thisBlock = lastBlock;

  size_t Ngood  = 0;
  size_t Nbad   = 0;
  size_t Ntotal = 0;
  size_t Nbytes = 0;

  while (thisBlock) {

    int good = TRUE;
    
    if (thisBlock->startblock != OHANA_MEMMAGIC) good = FALSE;
    if (thisBlock->endblock != OHANA_MEMMAGIC) good = FALSE;

    // pointer to the start of the user memory
    char *ptr = (char *)(thisBlock + 1) + thisBlock->size;
    uint32_t endpost = *(uint32_t *)ptr;
    if (endpost != OHANA_MEMMAGIC) good = FALSE;

    // XXX keep checking even if memory is corrupted?
    if (!good) {
      if (Nbad < 1) {
	fprintf (stderr, "memory corruption\n");
      }
      if (Nbad < 100) {
	fprintf (stderr, "  file: %s, line: %d, func: %s\n", thisBlock->file, thisBlock->line, thisBlock->func);
      }
      Nbad ++;
    } else {
      Ngood ++;
    }
    Ntotal ++;
    Nbytes += thisBlock->size;

    thisBlock = thisBlock->prevBlock;
  }
 
  fprintf (stderr, "%zd memory blocks allocated (%zd bytes total), %zd good, %zd bad\n", Ntotal, Nbytes, Ngood, Nbad);

  return;
}

void ohana_memdump_func (int allmemory) {

  if (!lastBlock) {
    fprintf (stderr, "no memory allocated\n");
    return;
  }

  Memblock *thisBlock = lastBlock;

  size_t Ntotal = 0;
  size_t Nbytes = 0;

  fprintf (stderr, " entry | bytes | cumulative | STATUS | file | line | function\n");

  while (thisBlock) {

    int good = TRUE;
    
    if (thisBlock->startblock != OHANA_MEMMAGIC) good = FALSE;
    if (thisBlock->endblock != OHANA_MEMMAGIC) good = FALSE;

    // pointer to the start of the user memory
    char *ptr = (char *)(thisBlock + 1) + thisBlock->size;
    uint32_t endpost = *(uint32_t *)ptr;
    if (endpost != OHANA_MEMMAGIC) good = FALSE;

    // XXX keep checking even if memory is corrupted?
    Ntotal ++;
    Nbytes += thisBlock->size;

    if (Ntotal < 100) {
      if (good) {
	fprintf (stderr, "  %zd  %zd  %zd  GOOD  %s %d, func: %s\n", Ntotal, thisBlock->size, Nbytes, thisBlock->file, thisBlock->line, thisBlock->func);
      } else {
	fprintf (stderr, "  %zd  %zd  %zd  BAD   %s %d, func: %s\n", Ntotal, thisBlock->size, Nbytes, thisBlock->file, thisBlock->line, thisBlock->func);
      }
    }

    thisBlock = thisBlock->prevBlock;
  }
  fprintf (stderr, "%zd memory blocks allocated (%zd bytes total)\n", Ntotal, Nbytes);

  return;
}

OhanaMemstats ohana_memstats (int allmemory) {

  OhanaMemstats memstats;

  memstats.exists = 0;
  memstats.Ntotal = 0;
  memstats.Nbytes = 0;
  memstats.Ngood = 0;
  memstats.Nbad = 0;
  
  Memblock *thisBlock = lastBlock;

  while (thisBlock) {

    memstats.exists = TRUE;

    int good = TRUE;
    
    if (thisBlock->startblock != OHANA_MEMMAGIC) good = FALSE;
    if (thisBlock->endblock != OHANA_MEMMAGIC) good = FALSE;

    // pointer to the start of the user memory
    char *ptr = (char *)(thisBlock + 1) + thisBlock->size;
    uint32_t endpost = *(uint32_t *)ptr;
    if (endpost != OHANA_MEMMAGIC) good = FALSE;

    memstats.Ntotal ++;
    memstats.Nbytes += thisBlock->size;

    if (good) {
      memstats.Ngood ++;
    } else {
      memstats.Nbad ++;
    }

    thisBlock = thisBlock->prevBlock;
  }

  return memstats;
}
