#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#define PS_ALLOW_MALLOC			// we're allowed to call malloc()
#include "psLib.h"
#include "Private/p_psMemory.h"

#undef psAlloc				// don't get the __FILE__ versions
#undef psRealloc
#undef psFree

static int bad_memblock(const psMemBlock *m, int quiet);
/*
 * This strawman allocation package uses an array to track calls to psAlloc;
 * not nearly good enough for production work!
 */
static int nMemBlock = 0;		// size of memBlocks
static psMemBlock **memBlocks = NULL;

/*****************************************************************************/
/*
 * Unique ID for allocated blocks
 */
static unsigned long memid = 0;

/*****************************************************************************/

void *psAlloc(size_t size)
{
    return p_psAlloc(size, "(unknown)", 0);
}

void *psRealloc(void *ptr, size_t size)
{
    return p_psRealloc(ptr, size, "(unknown)", 0);
}

void psFree(void *ptr)
{
    p_psFree(ptr, "(unknown)", 0);
}

/*****************************************************************************/
/*
 * Callbacks
 */
/*
 * First the I-can't-get-the-memory callback
 */
static void *memExhaustedCallback0(size_t size)
{
    return NULL;
}
static psMemExhaustedCallback memExhaustedCallback = memExhaustedCallback0;

psMemExhaustedCallback psMemExhaustedCallbackSet(psMemExhaustedCallback func)
{
    psMemExhaustedCallback old = memExhaustedCallback;

    memExhaustedCallback = (func != NULL) ? func : memExhaustedCallback0;

    return old;
}
/*
 * then the I-have-detected-a-problem callback
 */
static void memProblemCallback0(const psMemBlock *ptr,
			  const char *file, int lineno)
{
    if (bad_memblock(ptr, 0)) {
	fprintf(stderr, "Memory corruption or an attempt to manipulate memory "
		"not allocated with psAlloc at %s", file);
	if (lineno > 0) {
	    fprintf(stderr, ":%d", lineno);
	}
	fprintf(stderr, "\n");
    } else if (ptr->refCounter <= 0) {
	psError(__func__, PS_ERR_BADFREE, 1,
		"Block %ld allocated at %s:%d freed more than once at %s:%d\n",
		ptr->id, ptr->file, ptr->lineno, file, lineno);
    } else if (ptr->refCounter > 1) {
	psError(__func__,  PS_ERR_BADFREE, 1,
		"Block %ld allocated at %s:%d freed while still referenced at %s:%d\n",
		ptr->id, ptr->file, ptr->lineno, file, lineno);
    }

    if (lineno > 0) {
	psAbort(__func__, "Detected a problem in the memory system at %s:%d", file, lineno);
    }
}
static psMemProblemCallback memProblemCallback = memProblemCallback0;

psMemProblemCallback psMemProblemCallbackSet(psMemProblemCallback func)
{
    psMemProblemCallback old = memProblemCallback;

    memProblemCallback = (func != NULL) ? func : memProblemCallback0;

    return old;
}
/*
 * And now the I-want-to-be-informed callbacks
 *
 * Call the callbacks when these IDs are allocated/freed
 */
long p_psMemAllocateID = 0;		// notify user this block is allocated
long p_psMemFreeID = 0;			// notify user this block is freed

long psMemAllocateIDSet(long id)	// set p_psMemAllocateID to id
{
    long old = p_psMemAllocateID;
    p_psMemAllocateID = id;

    return old;
}

long psMemFreeIDSet(long id)		// set p_psMemFreeID to id
{
    long old = p_psMemFreeID;
    p_psMemFreeID = id;

    return old;
}

/*
 * Default callback for both allocate and free. Note that the
 * value of p_psMemAllocateID/p_psMemFreeID is incremented
 * by the return value (so returning 0 means that the callback
 * isn't resignalled)
 */
static long memAllocateCallback0(const psMemBlock *ptr)
{
    static int incr = 0;		// "p_psMemAllocateID += incr"

    assert (ptr != NULL);		// prevent compiler warnings
    
    return incr;
}

static long memFreeCallback0(const psMemBlock *ptr)
{
    static int incr = 0;		// "p_psMemFreeID += incr"

    assert (ptr != NULL);		// prevent compiler warnings
    
    return incr;
}

/*
 * The default callbacks, and the routines to change them
 */
static psMemAllocateCallback memAllocateCallback = memAllocateCallback0;
static psMemFreeCallback memFreeCallback = memFreeCallback0;

psMemAllocateCallback psMemAllocateCallbackSet(psMemAllocateCallback func)
{
    psMemAllocateCallback old = memAllocateCallback;

    memAllocateCallback = (func != NULL) ? func : memAllocateCallback0;

    return old;
}

psMemFreeCallback psMemFreeCallbackSet(psMemFreeCallback func)
{
    psMemFreeCallback old = memFreeCallback;

    memFreeCallback = (func != NULL) ? func : memFreeCallback0;

    return old;
}

/*****************************************************************************/
/*
 * Return memory ID counter for next block to be allocated
 */
int psMemGetId(void)
{
    return memid + 1;
}

/*****************************************************************************/
/*
 * Routines to check the consistency of the allocated and/or free memory arena
 *
 * N.b. If the block wasn't allocated by psAlloc, it will appear corrupted
 */
#define ALIGNED(P) ((void *)((long)(P) & ~03) == (P))

static int
bad_memblock(const psMemBlock *m,	// block to check
	     int quiet)			// be quiet?
{
    if (m == NULL) {
	if (!quiet) {
	    psError(__func__,  PS_ERR_MEMORY_CORRUPTION, 1,
		    "psMemCheckCorruption: NULL memory block\n");
	}
	return(1);
    }
    
    if (!ALIGNED(m)) {
	if (!quiet) {
	    psError(__func__,  PS_ERR_MEMORY_CORRUPTION, 1,
		    "psMemCheckCorruption: non-aligned memory block\n");
	}
	return(1);
    }
    
    if (m->startblock != P_PS_MEMMAGIC || m->endblock != P_PS_MEMMAGIC) {
	if (!quiet) {
	    psError(__func__, PS_ERR_MEMORY_CORRUPTION, 1,
		    "psMemCheckCorruption: memory block %ld is corrupted\n", m->id);
	}
	return(1);
    }

    return(0);
}

int psMemCheckCorruption(int abort_on_error)
{
    int nbad = 0;			// number of bad blocks

    if (memBlocks == NULL) {
	return 0;			// no memory is allocated to be corrupted
    }

    for (int i = 1; i <= memid; i++) {
	if (memBlocks[i] != NULL) {
	    if (bad_memblock(memBlocks[i], 1)) {
		nbad++;

		memProblemCallback(memBlocks[i], "(psMemCheckCorruption)", 0);
		
		if (abort_on_error) {
		    psAbort(__func__, "Detected memory corruption");
		}
	    }
	}
    }

    return nbad;
}

/*****************************************************************************/

void *p_psAlloc(size_t size, const char *file, int lineno)
{
    psMemBlock *ptr = malloc(sizeof(psMemBlock) + size);

    if (ptr == NULL) {
	ptr = memExhaustedCallback(size);
	if (ptr == NULL) {
	    psAbort(__func__, "Failed to allocate %ld bytes at %s:%d",
		    size, file, lineno);
	}
    }

    *(unsigned long *)&ptr->id = ++memid;
    ptr->file = file;
    *(int *)&ptr->lineno = lineno;
    *(void **)&ptr->startblock = *(void **)&ptr->endblock = P_PS_MEMMAGIC;
    ptr->endpost = &ptr->endblock;	// should point just beyond allocated memory XXX

    ptr->refCounter = 1;		// one user so far
    /*
     * Did the user ask to be informed about this allocation?
     */
    
    if (ptr->id == p_psMemAllocateID && p_psMemAllocateID > 0) {
	p_psMemAllocateID += memAllocateCallback(ptr);	
    }
    /*
     * Register the allocation.  This is not a production quality algorithm!
     */
    if (memid >= nMemBlock) {
	nMemBlock = (nMemBlock == 0) ? 10000 : 2*nMemBlock;
	memBlocks = realloc(memBlocks, nMemBlock*sizeof(psMemBlock *)); // NOT psRealloc
    }
    memBlocks[ptr->id] = ptr;
    /*
     * And return the user the memory that they allocated
     */
    return ptr + 1;			// usr memory
}

void *p_psRealloc(void *vptr, size_t size, const char *file, int lineno)
{
    if (vptr == NULL) {
	return p_psAlloc(size, file, lineno);
    } else {
	psMemBlock *ptr = ((psMemBlock *)vptr) - 1;
	ptr = (psMemBlock *)realloc(ptr, sizeof(psMemBlock) + size);
	memBlocks[ptr->id] = ptr;
	
	return ptr + 1;			// usr memory
    }
}

void p_psFree(void *vptr, const char *file, int lineno)
{
    if (vptr == NULL) {
	return;
    }

    psMemBlock *ptr = ((psMemBlock *)vptr) - 1;
    /*
     * Did the user ask to be informed about this deallocation?
     */
    if (ptr->id == p_psMemFreeID && p_psMemFreeID > 0) {
	p_psMemFreeID += memFreeCallback(ptr);	
    }

    if (bad_memblock(ptr, 0)) {
	memProblemCallback(ptr, file, lineno); // we may not own this block; don't free it
    } else {
	if (ptr->refCounter != 1) {
	    memProblemCallback(ptr, file, lineno);
	}
	ptr->refCounter--;
#if 1
	memBlocks[ptr->id] = NULL;	// XXX sub-optimal! Can't check
					// freed blocks for corruption/double frees
	free(ptr);
#endif
    }
}

/*****************************************************************************/
/*
 * Check for memory leaks. Not production quality code
 */
int psMemCheckLeaks(
    int id0,				// don't list blocks with id < id0
    psMemBlock ***arr,			// pointer to array of pointers to
					// leaked blocks, or NULL
    FILE *fd)				// print list of leaks to fd (or NULL)
{
    int nleak = 0;

    if (memBlocks == NULL) {		// nothing was allocated
	return 0;
    }

    for (int i = id0; i <= memid; i++) {
	if (memBlocks[i] != NULL) {
	    if (memBlocks[i]->refCounter > 0) {
		nleak++;

		if (fd != NULL) {
		    if (nleak == 1) {
			fprintf(fd, "   %20s:line ID\n", "file");
		    }

		    fprintf(fd, "   %20s:%-4d %ld\n",
			    memBlocks[i]->file, memBlocks[i]->lineno,
			    memBlocks[i]->id);
		}
	    }
	}
    }

    if (nleak == 0 || arr == NULL) {
	return nleak;
    }

    assert (*arr == NULL);		// Don't generate a memory leak!

    *arr = p_psAlloc(nleak*sizeof(psMemBlock), __FILE__, __LINE__);
    for (int i = id0, j = 0; i < memid; i++) {
	if (memBlocks[i] != NULL) {
	    if (memBlocks[i]->refCounter > 0) {
		(*arr)[j++] = memBlocks[i];
		if (j == nleak) {	// found them all
		    break;
		}
	    }
	}
    }

    return nleak;
}

/*****************************************************************************/
/*
 * Reference counting APIs
 */
int psMemGetRefCounter(void *vptr)	// return refCounter
{
    if (vptr == NULL) {
	return 0;
    }

    psMemBlock *ptr = ((psMemBlock *)vptr) - 1;

    if (bad_memblock(ptr, 0)) {
	memProblemCallback(ptr, "(psMemGetRefCounter)", -1);
    }

    return ptr->refCounter;
}

void *psMemIncrRefCounter(void *vptr)	// increment and return refCounter
{
    if (vptr == NULL) {
	return vptr;
    }
		 
    psMemBlock *ptr = ((psMemBlock *)vptr) - 1;

    if (bad_memblock(ptr, 0)) {
	memProblemCallback(ptr, "(psMemIncrRefCounter)", -1);
    }

    ptr->refCounter++;

    return vptr;
}

void *psMemDecrRefCounter(void *vptr)	// decrement and return refCounter
{
    if (vptr == NULL) {
	return vptr;
    }

    psMemBlock *ptr = ((psMemBlock *)vptr) - 1;

    if (bad_memblock(ptr, 0)) {
	memProblemCallback(ptr, "(psMemDecrRefCounter)", -1);
    }

    ptr->refCounter--;

    return vptr;
}
