/**
 * Copyright (c) 2026 Dominic Masters
 *
 * This software is released under the MIT License.
 * https://opensource.org/licenses/MIT
 */

#include "audiostreammp3decoder.h"
#include "audio/audiostream.h"
#include "assert/assert.h"
#include "util/memory.h"
#include <pspmp3.h>
#include <psputility.h>

// Per the PSP SDK's mp3 sample and pspmp3.h's own documentation - these
// are hard minimums sceMp3ReserveMp3Handle() requires, not tunable.
#define AUDIO_MP3_PSP_STREAM_BUF_SIZE (16 * 1024)
#define AUDIO_MP3_PSP_PCM_BUF_SIZE (16 * (1152 / 2))

// sceMp3 manages its own internal ring buffer of compressed bytes rather
// than being handed a whole file/buffer up front (unlike, say,
// zip_source_buffer_create elsewhere in this codebase) - it tells the
// caller exactly where in the source stream it wants more data from
// (sceMp3GetInfoToAddStreamData()'s srcpos) and how much room is free to
// receive it into, and the caller is expected to fetch exactly that and
// call sceMp3NotifyAddStreamData(). This decoder only ever feeds it
// forward, sequentially, from stream->mp3.file (never anything sceMp3
// itself wouldn't naturally request next), so despite the API being built
// for arbitrary positioning, this driver never needs to seek that file
// out of sequence - see audioStreamMp3DecoderFillStream()'s own assertion.

// End-of-stream sentinel returned by sceMp3Decode() once nothing more can
// be decoded - not documented in pspmp3.h, but confirmed by the PSP SDK's
// own mp3 sample (main.c), which explicitly excludes it from its generic
// "sceMp3Decode failed" error path.
#define AUDIO_MP3_PSP_DECODE_END_OF_STREAM 0x80671402

// Set once audioStreamMp3DecoderInit() has loaded sceMp3's firmware
// modules and initialized its resource pool - lazily, on the first MP3
// stream ever created, rather than unconditionally at engine startup, so
// a game that never plays an MP3 never pays for it. Guards
// audioStreamMp3DecoderGlobalDispose() so it only tears down what was
// actually set up.
static bool_t AUDIO_MP3_PSP_RESOURCES_READY = false;

errorret_t audioStreamMp3DecoderFillStream(audiostream_t *stream) {
  audiostreammp3decoder_t *decoder = &stream->mp3.decoder;

  SceUChar8 *dst = NULL;
  SceInt32 toWrite = 0;
  SceInt32 srcPos = 0;
  SceInt32 status = sceMp3GetInfoToAddStreamData(
    decoder->handle, &dst, &toWrite, &srcPos
  );
  if(status < 0) {
    errorThrow("sceMp3GetInfoToAddStreamData failed: 0x%08X", status);
  }

  if(toWrite <= 0) {
    errorOk();
  }

  assertTrue(
    (size_t) srcPos == decoder->streamPosition,
    "sceMp3 requested stream data out of the sequence this decoder feeds it in."
  );

  errorChain(assetFileRead(&stream->mp3.file, dst, (size_t) toWrite));
  const size_t bytesRead = (size_t) stream->mp3.file.lastRead;

  status = sceMp3NotifyAddStreamData(decoder->handle, (SceInt32) bytesRead);
  if(status < 0) {
    errorThrow("sceMp3NotifyAddStreamData failed: 0x%08X", status);
  }

  decoder->streamPosition += bytesRead;
  if(bytesRead < (size_t) toWrite) decoder->endOfFile = true;

  errorOk();
}

errorret_t audioStreamMp3DecoderInit(audiostream_t *stream) {
  assertNotNull(stream, "Stream cannot be NULL.");
  assertTrue(stream->type == AUDIO_STREAM_TYPE_MP3, "Stream is not MP3.");

  if(!AUDIO_MP3_PSP_RESOURCES_READY) {
    int status = sceUtilityLoadModule(PSP_MODULE_AV_AVCODEC);
    if(status < 0) {
      errorThrow("Failed to load PSP_MODULE_AV_AVCODEC: 0x%08X", status);
    }
    status = sceUtilityLoadModule(PSP_MODULE_AV_MP3);
    if(status < 0) {
      errorThrow("Failed to load PSP_MODULE_AV_MP3: 0x%08X", status);
    }
    status = sceMp3InitResource();
    if(status < 0) {
      errorThrow("sceMp3InitResource failed: 0x%08X", status);
    }
    AUDIO_MP3_PSP_RESOURCES_READY = true;
  }

  const assetmp3file_t *mp3 = &stream->asset->data.mp3;
  audiostreammp3decoder_t *decoder = &stream->mp3.decoder;

  decoder->mp3Buf = memoryAlign(64, AUDIO_MP3_PSP_STREAM_BUF_SIZE);
  decoder->pcmBuf = memoryAlign(64, AUDIO_MP3_PSP_PCM_BUF_SIZE);
  decoder->streamPosition = mp3->dataOffset;
  decoder->endOfFile = false;

  SceMp3InitArg initArg;
  memoryZero(&initArg, sizeof(initArg));
  initArg.mp3StreamStart = (SceOff) mp3->dataOffset;
  initArg.mp3StreamEnd = (SceOff) (mp3->dataOffset + mp3->dataSize);
  initArg.mp3Buf = decoder->mp3Buf;
  initArg.mp3BufSize = AUDIO_MP3_PSP_STREAM_BUF_SIZE;
  initArg.pcmBuf = decoder->pcmBuf;
  initArg.pcmBufSize = AUDIO_MP3_PSP_PCM_BUF_SIZE;

  const SceInt32 handle = sceMp3ReserveMp3Handle(&initArg);
  if(handle < 0) {
    memoryFree(decoder->mp3Buf);
    memoryFree(decoder->pcmBuf);
    errorThrow("sceMp3ReserveMp3Handle failed: 0x%08X", handle);
  }
  decoder->handle = handle;

  // sceMp3Init() needs some stream data already fed before it can
  // determine the stream's format - see the PSP SDK's own mp3 sample.
  errorChain(audioStreamMp3DecoderFillStream(stream));

  const SceInt32 initStatus = sceMp3Init(decoder->handle);
  if(initStatus < 0) {
    sceMp3ReleaseMp3Handle(decoder->handle);
    memoryFree(decoder->mp3Buf);
    memoryFree(decoder->pcmBuf);
    errorThrow("sceMp3Init failed: 0x%08X", initStatus);
  }

  errorOk();
}

errorret_t audioStreamMp3DecoderDispose(audiostream_t *stream) {
  assertNotNull(stream, "Stream cannot be NULL.");
  assertTrue(stream->type == AUDIO_STREAM_TYPE_MP3, "Stream is not MP3.");

  audiostreammp3decoder_t *decoder = &stream->mp3.decoder;

  sceMp3ReleaseMp3Handle(decoder->handle);
  memoryFree(decoder->mp3Buf);
  decoder->mp3Buf = NULL;
  memoryFree(decoder->pcmBuf);
  decoder->pcmBuf = NULL;

  errorOk();
}

errorret_t audioStreamMp3DecoderRewind(audiostream_t *stream) {
  assertNotNull(stream, "Stream cannot be NULL.");
  assertTrue(stream->type == AUDIO_STREAM_TYPE_MP3, "Stream is not MP3.");

  const assetmp3file_t *mp3 = &stream->asset->data.mp3;

  errorChain(assetFileRewind(&stream->mp3.file));
  errorChain(assetFileRead(&stream->mp3.file, NULL, mp3->dataOffset));

  // Deliberately not sceMp3ResetPlayPosition(): the SDK doesn't document
  // precisely enough whether it fully clears sceMp3's own internal
  // compressed-data ring buffer or just resets a playback cursor, leaving
  // room for it to disagree with the fresh stream position we just fed it
  // from - a mismatch that would only be audible as an occasional click
  // right at the loop point, depending on how much of sceMp3's internal
  // buffer happened to still hold stale data at reset time. Fully tearing
  // down and recreating the decoder instead guarantees a clean state by
  // construction, via the exact same path already proven correct for the
  // stream's very first Init - one extra handle release/reserve cycle per
  // loop wrap is negligible next to how rarely loops actually happen.
  errorChain(audioStreamMp3DecoderDispose(stream));
  errorChain(audioStreamMp3DecoderInit(stream));

  errorOk();
}

errorret_t audioStreamMp3DecoderDecodeFrame(
  audiostream_t *stream,
  int16_t *out,
  size_t *outFrames
) {
  assertNotNull(stream, "Stream cannot be NULL.");
  assertNotNull(out, "Out buffer cannot be NULL.");
  assertNotNull(outFrames, "outFrames cannot be NULL.");
  assertTrue(stream->type == AUDIO_STREAM_TYPE_MP3, "Stream is not MP3.");

  audiostreammp3decoder_t *decoder = &stream->mp3.decoder;
  const size_t frameSize = stream->channels * sizeof(int16_t);

  for(;;) {
    if(sceMp3CheckStreamDataNeeded(decoder->handle) > 0 && !decoder->endOfFile) {
      errorChain(audioStreamMp3DecoderFillStream(stream));
    }

    SceShort16 *pcm = NULL;
    const SceInt32 bytesDecoded = sceMp3Decode(decoder->handle, &pcm);

    if(bytesDecoded > 0) {
      // `out` (stream->mp3.pending) is sized for
      // AUDIO_MP3_MAX_SAMPLES_PER_FRAME samples - sized specifically to
      // cover sceMp3's pcmBuf provisioning (see this decoder's own
      // AUDIO_MP3_PSP_PCM_BUF_SIZE comment). Asserted rather than trusted:
      // a real hardware overflow here was previously silent, corrupting
      // whatever followed stream->mp3.pending in memory, and only showed
      // up as very intermittent audio clicking.
      assertTrue(
        (size_t) bytesDecoded <= AUDIO_MP3_MAX_SAMPLES_PER_FRAME * sizeof(int16_t),
        "sceMp3Decode returned more PCM than the destination buffer can hold."
      );
      const size_t frames = (size_t) bytesDecoded / frameSize;
      memoryCopy(out, pcm, (size_t) bytesDecoded);
      *outFrames = frames;
      errorOk();
    }

    if(
      bytesDecoded < 0 &&
      (uint32_t) bytesDecoded != AUDIO_MP3_PSP_DECODE_END_OF_STREAM
    ) {
      errorThrow("sceMp3Decode failed: 0x%08X", bytesDecoded);
    }

    // Decoded nothing this call - if sceMp3 doesn't need more data either,
    // it's genuinely out of frames to produce.
    if(sceMp3CheckStreamDataNeeded(decoder->handle) <= 0) {
      *outFrames = 0;
      errorOk();
    }

    if(decoder->endOfFile) {
      // sceMp3 wants more data but the file has none left - one last
      // decode attempt already happened above with whatever was fed;
      // nothing further will ever arrive.
      *outFrames = 0;
      errorOk();
    }
  }
}

errorret_t audioStreamMp3DecoderGlobalDispose(void) {
  if(!AUDIO_MP3_PSP_RESOURCES_READY) {
    errorOk();
  }

  const SceInt32 status = sceMp3TermResource();
  if(status < 0) {
    errorThrow("sceMp3TermResource failed: 0x%08X", status);
  }

  sceUtilityUnloadModule(PSP_MODULE_AV_MP3);
  sceUtilityUnloadModule(PSP_MODULE_AV_AVCODEC);
  AUDIO_MP3_PSP_RESOURCES_READY = false;

  errorOk();
}