/*
  Copyright (c) 2021 Alex Diener
  
  This software is provided 'as-is', without any express or implied
  warranty. In no event will the authors be held liable for any damages
  arising from the use of this software.
  
  Permission is granted to anyone to use this software for any purpose,
  including commercial applications, and to alter it and redistribute it
  freely, subject to the following restrictions:
  
  1. The origin of this software must not be misrepresented; you must not
     claim that you wrote the original software. If you use this software
     in a product, an acknowledgment in the product documentation would be
     appreciated but is not required.
  2. Altered source versions must be plainly marked as such, and must not be
     misrepresented as being the original software.
  3. This notice may not be removed or altered from any source distribution.
  
  Alex Diener alex@ludobloom.com
*/

#include "pcmaudio/WAVAudioIO.h"
#include "utilities/IOUtilities.h"
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <unistd.h>

#define WAV_HEADER_SIZE 44

static uint16_t parseUInt16LE(const char * value) {
	return ((uint16_t) value[0] & 0xFF) |
	       (((uint16_t) value[1] & 0xFF) << 8);
}

static uint32_t parseUInt32LE(const char * value) {
	return ((uint32_t) value[0] & 0xFF) |
	       (((uint32_t) value[1] & 0xFF) << 8) |
	       (((uint32_t) value[2] & 0xFF) << 16) |
	       (((uint32_t) value[3] & 0xFF) << 24);
}

static void packUInt16LE(uint16_t value, char * outBuffer) {
	outBuffer[0] = value & 0xFF;
	outBuffer[1] = value >> 8 & 0xFF;
}

static void packUInt32LE(uint32_t value, char * outBuffer) {
	outBuffer[0] = value & 0xFF;
	outBuffer[1] = value >> 8 & 0xFF;
	outBuffer[2] = value >> 16 & 0xFF;
	outBuffer[3] = value >> 24 & 0xFF;
}

static uint32_t readChunkHeader(size_t (* readCallback)(void * outData, size_t length, void * context), void * context, char * outChunkType) {
	char chunkSize[4];
	readCallback(outChunkType, 4, context);
	readCallback(chunkSize, 4, context);
	return parseUInt32LE(chunkSize);
}

static bool isChunkType(const char * type, const char * expectedType) {
	return type[0] == expectedType[0] &&
	       type[1] == expectedType[1] &&
	       type[2] == expectedType[2] &&
	       type[3] == expectedType[3];
}

PCMAudio * readWAVInternal(size_t (* readCallback)(void * outData, size_t length, void * context), void (* seekForwardCallback)(size_t offset, void * context), bool (* eofCallback)(void * context), void * context) {
	char chunkType[4];
	unsigned int chunkSize = readChunkHeader(readCallback, context, chunkType);
	if (!isChunkType(chunkType, "RIFF")) {
#ifdef DEBUG
		fprintf(stderr, "readWAVInternal error: File doesn't look like a WAV file (no RIFF chunk at beginning)\n");
#endif
		return NULL;
	}
	
	readCallback(chunkType, 4, context);
	if (!isChunkType(chunkType, "WAVE")) {
#ifdef DEBUG
		fprintf(stderr, "readWAVInternal error: File is in an unknown format (RIFF chunk's type is something other than WAVE; found \"%c%c%c%c\")\n", chunkType[0], chunkType[1], chunkType[2], chunkType[3]);
#endif
		return NULL;
	}
	
	chunkSize = readChunkHeader(readCallback, context, chunkType);
	while (!isChunkType(chunkType, "fmt ") && !eofCallback(context)) {
		seekForwardCallback(chunkSize, context);
		chunkSize = readChunkHeader(readCallback, context, chunkType);
	}
	if (eofCallback(context)) {
#ifdef DEBUG
		fprintf(stderr, "readWAVInternal error: Reached EOF without finding 'fmt ' chunk\n");
#endif
		return NULL;
	}
	if (chunkSize == 18) {
		seekForwardCallback(2, context);
	} else if (chunkSize != 16) {
#ifdef DEBUG
		fprintf(stderr, "readWAVInternal error: Unexpected WAV header format ('fmt ' chunk not 16 or 18 bytes in size)\n");
#endif
		return NULL;
	}
	
	char fmtHeader[16];
	readCallback(fmtHeader, 16, context);
	
	chunkSize = readChunkHeader(readCallback, context, chunkType);
	while (!isChunkType(chunkType, "data") && !eofCallback(context)) {
		seekForwardCallback(chunkSize, context);
		chunkSize = readChunkHeader(readCallback, context, chunkType);
	}
	if (eofCallback(context)) {
#ifdef DEBUG
		fprintf(stderr, "readWAVInternal error: Reached EOF without finding 'data' chunk\n");
#endif
		return NULL;
	}
	
	void * samples = malloc(chunkSize);
	readCallback(samples, chunkSize, context);
	
	unsigned int channelCount = parseUInt16LE(fmtHeader + 2);
	unsigned int sampleRate = parseUInt32LE(fmtHeader + 4);
	unsigned int bytesPerSample = parseUInt16LE(fmtHeader + 14) / 8;
	
	return PCMAudio_create(bytesPerSample, channelCount, sampleRate, chunkSize / bytesPerSample / channelCount, samples, false);
}

bool writeWAVInternal(PCMAudio * audio, bool (* writeCallback)(const void * data, size_t length, void * context), void * context) {
	char wavHeader[WAV_HEADER_SIZE] = "RIFF\x00\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00\x22\x56\x00\x00\x44\xAC\x00\x00\x02\x00\x10\x00""data\x00\x00\x00\x00";
	packUInt32LE(audio->frameCount * audio->channelCount * audio->bytesPerSample + 36, wavHeader + 4); // RIFF chunk length
	packUInt16LE(audio->channelCount, wavHeader + 22); // wChannels
	packUInt32LE(audio->sampleRate, wavHeader + 24); // dwSamplesPerSec
	packUInt32LE(audio->sampleRate * audio->channelCount * audio->bytesPerSample, wavHeader + 28); // dwAvgBytesPerSec
	packUInt16LE(audio->channelCount * audio->bytesPerSample, wavHeader + 32); // wBlockAlign
	packUInt16LE(audio->bytesPerSample * 8, wavHeader + 34); // wBitsPerSample
	packUInt32LE(audio->frameCount * audio->channelCount * audio->bytesPerSample, wavHeader + 40); // data chunk length
	return writeCallback(wavHeader, WAV_HEADER_SIZE, context) &&
	       writeCallback(audio->samples, audio->frameCount * audio->channelCount * audio->bytesPerSample, context);
}

static size_t readCallback_fread(void * outData, size_t length, void * context) {
	return fread(outData, 1, length, context);
}

static void seekForwardCallback_fseek(size_t offset, void * context) {
	fseek(context, offset, SEEK_CUR);
}

static bool eofCallback_feof(void * context) {
	return feof((FILE *) context);
}

PCMAudio * WAVAudioIO_readWAVFile(const char * filePath) {
	FILE * file;
	if (filePath == NULL) {
		file = fdopen(STDIN_FILENO, "rb");
		if (file == NULL) {
#ifdef DEBUG
			fprintf(stderr, "WAVAudioIO_readWAVFile couldn't open stdout for reading (errno = %d)\n", errno);
#endif
			return NULL;
		}
	} else {
		file = fopen(filePath, "rb");
		if (file == NULL) {
#ifdef DEBUG
			fprintf(stderr, "WAVAudioIO_readWAVFile couldn't open \"%s\" for reading (errno = %d)\n", filePath, errno);
#endif
			return NULL;
		}
	}
	PCMAudio * audio = readWAVInternal(readCallback_fread, seekForwardCallback_fseek, eofCallback_feof, file);
	if (filePath != NULL) {
		fclose(file);
	}
	return audio;
}

static size_t readCallback_memread(void * outData, size_t length, void * context) {
	struct memreadContext * memreadContext = context;
	bool success = memread(memreadContext, length, outData);
	if (!success) {
		size_t remainingLength = memreadContext->length - memreadContext->position;
		memread(memreadContext, remainingLength, outData);
		return remainingLength;
	}
	return length;
}

static void seekForwardCallback_memread(size_t offset, void * context) {
	struct memreadContext * memreadContext = context;
	memreadContext->position += offset;
}

static bool eofCallback_memread(void * context) {
	struct memreadContext * memreadContext = context;
	return memreadContext->position >= memreadContext->length;
}

PCMAudio * WAVAudioIO_readWAVData(const void * data, size_t size) {
	struct memreadContext memreadContext = memreadContextInit(data, size);
	return readWAVInternal(readCallback_memread, seekForwardCallback_memread, eofCallback_memread, &memreadContext);
}

static bool writeCallback_fwrite(const void * data, size_t length, void * context) {
	return fwrite(data, 1, length, context) == length;
}

bool WAVAudioIO_writeWAVFile(PCMAudio * audio, const char * filePath) {
	FILE * file;
	if (filePath == NULL) {
		file = fdopen(STDOUT_FILENO, "wb");
		if (file == NULL) {
#ifdef DEBUG
			fprintf(stderr, "WAVAudioIO_writeWAVFile error: Couldn't open stdout for writing (errno = %d)\n", errno);
#endif
			return false;
		}
		
	} else {
		file = fopen(filePath, "wb");
		if (file == NULL) {
#ifdef DEBUG
			fprintf(stderr, "WAVAudioIO_writeWAVFile error: Couldn't open \"%s\" for writing (errno = %d)\n", filePath, errno);
#endif
			return false;
		}
	}
	bool success = writeWAVInternal(audio, writeCallback_fwrite, file);
	if (filePath != NULL) {
		fclose(file);
	}
	return success;
}

static bool writeCallback_memwrite(const void * data, size_t length, void * context) {
	return memwrite(context, length, data);
}

void * WAVAudioIO_writeWAVData(PCMAudio * audio, size_t * outSize) {
	size_t allocatedSize = WAV_HEADER_SIZE + audio->frameCount * audio->channelCount * audio->bytesPerSample;
	struct memwriteContext memwriteContext = memwriteContextInit(malloc(allocatedSize), 0, allocatedSize, false);
	bool success = writeWAVInternal(audio, writeCallback_memwrite, &memwriteContext);
	if (!success) {
		free(memwriteContext.data);
		return false;
	}
	*outSize = memwriteContext.position;
	return memwriteContext.data;
}