#include "sfl.h"
#include "sflIoObject.h"
#include "sflSoundMP3.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <malloc.h>
#include <limits.h>
#include <assert.h>
#include <pspkernel.h>
#include <pspaudio.h>
#include <pspaudiocodec.h>
#include <psputility_avmodules.h>

// シーク用のルックアップテーブルの要素数
#define SEEK_TABLE_SIZE 1024

// MP3はデコードに過去の入力データも使うので
// シークしたい場所の少し前から再生を始めないと正しいデコード結果にならないよ
#define SEEK_IGNORE_FRAMES 2

// デコーダが曲の先頭に追加する無音サンプル数
// 調べてみたら、 MPEG1-Layer3 と MPEG2-Layer3 では、他のデコーダと同様 529 だった
// 他の version, layer でも同じかどうかは謎
// というかlayer3以外が再生できるかどうかが謎
#define MP3_DECODER_DELAY 529

// 文字列をunsigned longで比較する用
#define STR_ID3  0x00334449
#define STR_XING 0x676e6958
#define STR_INFO 0x6f666e49

struct sflSoundMP3_t;
typedef struct sflSoundMP3_t sflSoundMP3;

struct sflSeekTable_t {
	int data[SEEK_TABLE_SIZE];
	int fill;
	int step;
};

struct sflSoundMP3_t {
	union {
		struct {
			sflSound super; 
			sflIoObject *io;
			int id3v2tagSize;
			int streamSize;
			long totalSamples;

			long decodedSamples;
			int frameNum;
			int firstOffset;
			int mixOutput;
			int mixBufCount;

			short vbrHeaderSize;
			short accurateLength;
			short encoderDelay;
			short encoderPadding;
			unsigned short samplerate;
			short samplePerFrame;
			short bitrate;
			short channelMode;
			short version;
			short layer;

			char getEdram;
		};
		unsigned char dummy[64];
	};

	short mixBuffer[1152 * 2];
	unsigned char dataBuffer[1472];
	unsigned long codecBuffer[65];
	struct sflSeekTable_t index;
};

static int sflSoundMP3Close(sflSound *self);
static size_t sflSoundMP3Read(sflSound *self, void *outbuf, size_t size);
static int sflSoundMP3Seek(sflSound *self, int offset, int whence);
static long sflSoundMP3Tell(sflSound *self);
static sflSound *sflSoundMP3Clone(sflSound *self);

static int sflSoundMP3SeekSample(sflSoundMP3 *music, long pos);
static int sflSoundMP3Decode(sflSoundMP3 *music);

static int readMpegAudioHeader(sflSoundMP3 *music);
static int getID3v2TagSize(sflIoObject *io);
static int getEncoderDelayAndPadding(sflSoundMP3 *music);
static int seekNextFrame(sflSoundMP3 *music);
static void updateSeekTable(sflSoundMP3 *music);

static const unsigned short g_samplerates[4][3] = {
    {11025, 12000, 8000,},  //mpeg 2.5
    {0, 0, 0,},             //reserved
    {22050, 24000, 16000,}, //mpeg 2
    {44100, 48000, 32000}   //mpeg 1
};
static const short g_bitrates[2][4][16] = {
	{ // MPEG1
		{0, },                                                                    // reserved
		{0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0},    // layer 3
		{0, 32, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 384, 0},   // layer 2
		{0, 32, 64, 96, 128, 160, 192, 224, 256, 288, 320, 352, 384, 416, 448, 0} // layer 1
	},
	{ // MPEG2
		{0, },                                                                 // reserved
		{0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0},     // layer 3
		{0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0},     // layer 2
		{0, 32, 48, 56, 64, 80, 96, 112, 128, 144, 160, 176, 192, 224, 256, 0} // layer 1
	}
};
static const short g_sample_per_frame[4][4] = {
	{0, 576, 1152, 384}, // mpeg 2.5
	{0, },               // reserved
	{0, 576, 1152, 384}, // mpeg 2
	{0, 1152, 1152, 384} // mpeg 1
};

sflSound *sflSoundMP3Open(const char *filename)
{
	static void *sflSoundMP3Func[] = {
		(void*)sflSoundMP3Read,
		(void*)sflSoundMP3Seek,
		(void*)sflSoundMP3Tell,
		(void*)sflSoundMP3Close,
		(void*)sflSoundMP3Clone,
	};

	sflSoundMP3 *self = memalign(64, sizeof(sflSoundMP3));
	if (!self) {
		DEBUG_PRINTF("malloc failed\n");
		goto error;
	}
	DEBUG_PRINTF("sizeof(sflSoundMP3) = %d, mixBuffer address = %p\n", sizeof(sflSoundMP3), self->mixBuffer);

	// ゼロで初期化
	memset(self, 0, sizeof(sflSoundMP3));

	int res;
	//self->io = sflIoObjectDefaultOpen(filename, SFL_IO_READ_ONLY, 0777, &res);
	self->io = sflIoObjectFullBufferedOpen(filename, SFL_IO_READ_ONLY, 0777, &res);
	//self->io = sflIoObjectBufferedOpen(filename, SFL_IO_READ_ONLY, 0777, &res, 32768);
	if (!self->io) {
		DEBUG_PRINTF("file open failed 0x%08x\n", res);
		goto error;
	}

	self->id3v2tagSize = getID3v2TagSize(self->io);
	DEBUG_PRINTF("id3v2tagSize = %d\n", self->id3v2tagSize);

	if (getEncoderDelayAndPadding(self) < 0) {
		DEBUG_PRINTF("getEncoderDelayAndPadding failed\n");
		goto error;
	}

	self->firstOffset = self->encoderDelay + MP3_DECODER_DELAY;
	self->super.vptr = sflSoundMP3Func;

	return (sflSound*)self;

error:
	if (self) {
		if (self->io)
			self->io->close(self->io);

		free(self);
	}
	return 0;
}

static int sflSoundMP3Close(sflSound *self)
{
	sflSoundMP3 *music = (sflSoundMP3*)self;
	int ret;
	if (music->getEdram && (ret = sceAudiocodecReleaseEDRAM(music->codecBuffer)) < 0) {
		DEBUG_PRINTF("sceAudiocodecReleaseEDRAM failed 0x%08x\n", ret);
	}
	if ((ret = music->io->close(music->io)) < 0) {
		DEBUG_PRINTF("io->close failed 0x%08x\n", ret);
	}
	free(music);
	return ret;
}

static size_t sflSoundMP3Read(sflSound *self, void *outbuf, size_t size)
{
	sflSoundMP3 *music = (sflSoundMP3*)self;
	short *data = outbuf;
	int bufferFill = 0;
	if (music->mixOutput < music->mixBufCount) {
		int copySamples = music->mixBufCount - music->mixOutput;
		if (copySamples + bufferFill > size)
			copySamples = size - bufferFill;
		
		memcpy(&data[bufferFill * 2], &music->mixBuffer[music->mixOutput * 2], copySamples * 2 * sizeof(short));
		music->decodedSamples += copySamples;
		music->mixOutput += copySamples;
		bufferFill += copySamples;
	}
	while(bufferFill < size) {
		if (music->mixOutput == music->mixBufCount) {
			if (sflSoundMP3Decode(music) < 0) {
				if (seekNextFrame(music) > 0)
					continue;
				
				DEBUG_PRINTF("seekNextFrame failed, mixOutput = %d, mixBufCount = %d\n", music->mixOutput, music->mixBufCount);
				break;
			}

			music->mixBufCount = music->samplePerFrame;
			music->mixOutput = 0;

			// 末尾の要らないサンプルを出力バッファに入れないようにする
			if (music->decodedSamples + music->mixBufCount > music->totalSamples && music->accurateLength) {
				music->mixBufCount = music->totalSamples - music->decodedSamples;
				if (music->mixBufCount < 0)
					music->mixBufCount = 0;
			}

			// 先頭の要らないサンプルを出力バッファに入れないようにする
			if (music->firstOffset > 0) {
				if (music->mixOutput + music->firstOffset > music->mixBufCount) {
					music->firstOffset -= music->mixBufCount - music->mixOutput;
					music->mixOutput += music->mixBufCount - music->mixOutput; //music->mixOutput = music->mixBufCount;
				} else {
					music->mixOutput += music->firstOffset;
					music->firstOffset -= music->firstOffset; //music->firstOffset = 0;
				}
			}

			int copySamples = music->mixBufCount - music->mixOutput;
			if (copySamples + bufferFill > size)
				copySamples = size - bufferFill;

			if (copySamples > 0) {
				memcpy(&data[bufferFill * 2], &music->mixBuffer[music->mixOutput * 2], copySamples * 2 * sizeof(short));
				music->decodedSamples += copySamples;
				music->mixOutput += copySamples;
				bufferFill += copySamples;
			}
		} else {
			// バグがなければ、ここは実行されない
			DEBUG_PRINTF("bugbug, mixOutput = %d, mixBufCount = %d\n", music->mixOutput, music->mixBufCount);
			assert(0);
			break;
		}
	}
	return bufferFill;
}

static int sflSoundMP3Seek(sflSound *self, int offset, int whence)
{
	sflSoundMP3 *music = (sflSoundMP3*)self;
	long pos = music->decodedSamples;
	switch (whence) {
		case SEEK_SET:
			pos = offset;
			break;
		case SEEK_CUR:
			pos += offset;
			break;
		case SEEK_END:
			// VBRヘッダがない場合は合計サンプル数が分からないのでシークして調べてみる
			if (music->totalSamples == 0) {
				sflSoundMP3SeekSample(music, 0x7fffffff);
				music->totalSamples = music->decodedSamples;
			}
			pos = music->totalSamples + offset;
			break;
		default:
			DEBUG_PRINTF("invalid whence %d\n", whence);
			assert(0);
			return -1;
	}

	if (pos == music->decodedSamples)
		return 0;

	return sflSoundMP3SeekSample(music, pos);
}

static long sflSoundMP3Tell(sflSound *self)
{
	sflSoundMP3 *music = (sflSoundMP3*)self;
	return music->decodedSamples;
}

static int sflSoundMP3SeekSample(sflSoundMP3 *music, long pos)
{
	if (pos < 0)
		pos = 0;

	int unadjustPos = pos + music->encoderDelay + MP3_DECODER_DELAY;
	int wantFrame = unadjustPos / music->samplePerFrame;
	int ignoreFrame = wantFrame - SEEK_IGNORE_FRAMES;
	if (ignoreFrame < 0)
		ignoreFrame = 0;
	
	int firstOffset = unadjustPos - ignoreFrame * music->samplePerFrame;

	if (music->index.fill > 0) {
		int index = ignoreFrame >> music->index.step;
		if (index >= music->index.fill)
			index = music->index.fill - 1;
		DEBUG_PRINTF("ignore=%d, step=%d, fill=%d, index=%d\n", ignoreFrame, music->index.step, music->index.fill, index);
		
		music->frameNum = index << music->index.step;
		int byteOffset = music->index.data[index];
		music->io->lseek32(music->io, byteOffset, SEEK_SET);
		DEBUG_PRINTF("byteOffset = %d\n", byteOffset);
	}

	int spf = music->samplePerFrame;
	int decoded = pos + (music->frameNum - ignoreFrame) * spf;
	int ts = music->totalSamples;
	int tsMinusSpf = ts - spf;
	if (!music->accurateLength)
		tsMinusSpf = INT_MAX;
	
	while(music->frameNum < ignoreFrame) {
		int frameSize = readMpegAudioHeader(music);
		if (frameSize < 0) {
			DEBUG_PRINTF("readMpegAudioHeader failed\n");
			int offset = seekNextFrame(music);
			if (offset < 0) {
				DEBUG_PRINTF("seekNextFrame failed\n");
				return -1;
			}

			continue;
		}

		updateSeekTable(music);
		music->frameNum++;
		
		music->io->lseek32(music->io, frameSize, SEEK_CUR);
		//int streamOffset = music->io->lseek32(music->io, frameSize, SEEK_CUR);
		//DEBUG_PRINTF("frameNum = %d, streamOffset = %d, frameSize = %d\n", music->frameNum, streamOffset, frameSize);

		if (decoded > tsMinusSpf) {
			decoded = ts;
			break;
		}

		decoded += spf;
	}

	music->decodedSamples = decoded;
	music->firstOffset = firstOffset;

	music->mixOutput = 0;
	music->mixBufCount = 0;

	//DEBUG_PRINTF("pos = %d, decoded = %d, mixOutput = %d, mixBufCount = %d\n", pos, decoded, music->mixOutput, music->mixBufCount);
	DEBUG_PRINTF("decoded = %ld, frameNum = %d, offset=%d\n", music->decodedSamples, music->frameNum, music->firstOffset);

	return 0;
}

static int seekNextFrame(sflSoundMP3 *music)
{
	unsigned char buf[1024];
	for(;;) {
		int offset = music->io->lseek32(music->io, 0, SEEK_CUR);
		int size = music->io->read(music->io, buf, sizeof(buf));
		if (size <= 2)
			return -1;
		
		int len = size - 1;
		int i;
		for(i = 0; i < len; ++i) {
			if (buf[i] == 0xff && (buf[i + 1] & 0xe0) == 0xe0) {
				offset += i;
				music->io->lseek32(music->io, offset, SEEK_SET);
				return offset;
			}
		}

		music->io->lseek32(music->io, -1, SEEK_CUR);
	}
}

static int readMpegAudioHeader(sflSoundMP3 *music)
{
	unsigned char buf[4];
	if (music->io->read(music->io, buf, 4) != 4)
		return -1;

	int header = buf[0];
	header = (header << 8) | buf[1];
	header = (header << 8) | buf[2];
	header = (header << 8) | buf[3];
	if ((header & 0xffe00000) != 0xffe00000 || // invalid frame sync
		(header & 0x00600000) == 0x00000000 || // invalid layer
		(header & 0x00180000) == 0x00080000 || // invalid version
		(header & 0x0000f000) == 0x0000f000 || // invalid bitrate
		(header & 0x0000f000) == 0x00000000 || // invalid bitrate
		(header & 0x00000c00) == 0x00000c00) { // invalid samplerate
		DEBUG_PRINTF("invalid header 0x%08x\n", header);
		music->io->lseek32(music->io, -3, SEEK_CUR); //seek back
		return -1;
	}

	music->version = (header & 0x00180000) >> 19;
	music->layer = (header & 0x00060000) >> 17;
	music->bitrate = g_bitrates[music->version != 3][music->layer][(header & 0xf000) >> 12];
	music->samplerate = g_samplerates[music->version][(header & 0xc00) >> 10];
	int padding = (header & 0x200) >> 9;
	music->channelMode = (header & 0xc0) >> 6;
	music->samplePerFrame = g_sample_per_frame[music->version][music->layer];

	music->io->lseek32(music->io, -4, SEEK_CUR); //seek back

	//return (music->samplePerFrame / 8 * music->bitrate * 1000) / music->samplerate + padding;
	return (125 * music->samplePerFrame * music->bitrate) / music->samplerate + padding;
}

static void updateSeekTable(sflSoundMP3 *music)
{
	if (music->frameNum != music->index.fill << music->index.step)
		return;
	
	if (music->index.fill == SEEK_TABLE_SIZE) {
		++music->index.step;
		music->index.fill >>= 1;
		int i;
		for(i = 0; i < music->index.fill; ++i)
			music->index.data[i] = music->index.data[i * 2];
	}

	if (music->frameNum == music->index.fill << music->index.step) {
		music->index.data[music->index.fill] = music->io->lseek32(music->io, 0, SEEK_CUR);
		music->index.fill++;
	}
}

static int sflSoundMP3Decode(sflSoundMP3 *music)
{
	int frameSize = readMpegAudioHeader(music);
	if (frameSize < 0)
		return -1;
	
	updateSeekTable(music);
	music->frameNum++;

	if (music->io->read(music->io, music->dataBuffer, frameSize) != frameSize)
		return -1;

	if (music->getEdram == 0) {
		int res;
		if ((res = sceAudiocodecCheckNeedMem(music->codecBuffer, PSP_CODEC_MP3)) < 0) {
			DEBUG_PRINTF("sceAudiocodecCheckNeedMem failed 0x%08x\n", res);
			return res;
		}
		if ((res = sceAudiocodecGetEDRAM(music->codecBuffer, PSP_CODEC_MP3)) < 0) {
			DEBUG_PRINTF("sceAudiocodecGetEDRAM failed 0x%08x\n", res);
			return res;
		}
		if ((res = sceAudiocodecInit(music->codecBuffer, PSP_CODEC_MP3)) < 0) {
			DEBUG_PRINTF("sceAudiocodecInit failed 0x%08x\n", res);
			sceAudiocodecReleaseEDRAM(music->codecBuffer);
			return res;
		}
		music->getEdram = 1;
	}

	music->codecBuffer[6] = (unsigned long)music->dataBuffer;
	music->codecBuffer[8] = (unsigned long)music->mixBuffer;
	music->codecBuffer[7] = music->codecBuffer[10] = frameSize;
	music->codecBuffer[9] = music->samplePerFrame * 4;

	int ret = sceAudiocodecDecode(music->codecBuffer, PSP_CODEC_MP3);
	if (ret < 0) {
		// ファイルの先頭の無音ぽい部分にシークするとたまに0x807f00fdエラーが発生するが
		// どうせ無音だし問題にはならないだろうからゼロで埋めてそのまま再生する
		memset(music->mixBuffer, 0x00, sizeof(short) * 1152 * 2);
		int pos = music->io->lseek32(music->io, 0, SEEK_CUR) - frameSize;
		DEBUG_PRINTF("sceAudiocodecDecode error 0x%08x %d %d\n", ret, pos, music->frameNum);
		if (ret != 0x807f00fd)
			return ret;
	}

	return 0;
}

static int getID3v2TagSize(sflIoObject *io)
{
	u8 buf[10];
	io->read(io, buf, sizeof(buf));
	io->lseek32(io, -sizeof(buf), SEEK_CUR);
	unsigned long *t = (unsigned long*)buf;
	if ((*t & 0xffffff) == STR_ID3) {
		int size = buf[6];
		size = (size << 7) | buf[7];
		size = (size << 7) | buf[8];
		size = (size << 7) | buf[9];
		size += 10;
		if (buf[5] & 0x10)
			size += 10;
		
		return size;
	}
	return 0;
}

static int getEncoderDelayAndPadding(sflSoundMP3 *music)
{
	// id3タグをスルー
	music->io->lseek32(music->io, music->id3v2tagSize, SEEK_SET);

	int frameSize = readMpegAudioHeader(music);
	if (frameSize < 0)
		return -1;

	if (music->io->read(music->io, music->dataBuffer, frameSize) != frameSize)
		return -1;

	u8 *p;
	if (music->channelMode != 3) { // モノラル以外
		if (music->version == 3) { // mpeg 1
			p = music->dataBuffer + 32 + 4;
		} else {
			p = music->dataBuffer + 17 + 4;
		}
	} else {
		if (music->version == 3) { // mpeg 1
			p = music->dataBuffer + 17 + 4;
		} else {
			p = music->dataBuffer + 9 + 4;
		}
	}

	music->streamSize = music->io->lseek32(music->io, 0, SEEK_END);
	music->streamSize -= music->id3v2tagSize;
	//music->totalSamples = 0x7fffffff;
	// vbrヘッダが存在
	unsigned long t = (p[3] << 24) | (p[2] << 16) | (p[1] << 8) | (p[0]);
	if (t == STR_XING || t == STR_INFO) {
		int offset = 8;
		if (p[7] & 0x01) { //Frames field is present
			// このフレームを含めないMP3フレームの合計数
			int frames = (p[offset] << 24) | (p[offset+1] << 16) | (p[offset+2] << 8) | p[offset+3];
			DEBUG_PRINTF("frames: %d\n", frames);
			DEBUG_PRINTF("samples: %d\n", frames * music->samplePerFrame);
			music->totalSamples = frames * music->samplePerFrame;
			offset += 4;
		}
		if (p[7] & 0x02) { // Bytes field is present
			// このフレームを含めたMP3フレームの合計バイト数
			// 先頭のID3v2タグと末尾のID3v1タグのサイズは含まない
			int bytes = (p[offset] << 24) | (p[offset+1] << 16) | (p[offset+2] << 8) | p[offset+3];
			music->streamSize = bytes - frameSize;
			DEBUG_PRINTF("bytes: %d\n", bytes);
			offset += 4;
		}
		if (p[7] & 0x04) { // TOC field is present
			offset += 100;
		}
		if (p[7] & 0x08) { // quality indicator field is present
			int quality = p[offset+3];
			DEBUG_PRINTF("quality; %d\n", quality);
			offset += 4;
		}

		p += offset + 21;
		music->accurateLength = 1;
		music->encoderDelay = (p[0] << 4) | (p[1] >> 4);
		music->encoderPadding = ((p[1] & 0x0f) << 8) | p[2];
		music->vbrHeaderSize = frameSize;
		music->totalSamples -= music->encoderDelay + music->encoderPadding;
		DEBUG_PRINTF("totalSamples: %ld\n", music->totalSamples);
		DEBUG_PRINTF("ENC_DELAY: %d\n", music->encoderDelay);
		DEBUG_PRINTF("ENC_PADDING: %d\n", music->encoderPadding);
	}

	music->io->lseek32(music->io, music->id3v2tagSize + music->vbrHeaderSize, SEEK_SET);
	
	return 0;
}

static sflSound *sflSoundMP3Clone(sflSound *self)
{
	sflSoundMP3 *sound = (sflSoundMP3*)self;
	sflSoundMP3 *clone = memalign(64, sizeof(sflSoundMP3));
	if (!clone) {
		DEBUG_PRINTF("memalign failed\n");
		return 0;
	}

	*clone = *sound;
	clone->getEdram = 0;
	clone->io = sound->io->clone(sound->io);
	if (!clone->io) {
		DEBUG_PRINTF("sound->io->clone failed\n");
		free(clone);
	}

	return (sflSound*)clone;
}

