/*
 * One-call LZMA decompression with chunked flash output.
 *
 * The caller supplies:
 *   - all compressed bytes through one contiguous memory pointer;
 *   - the five LZMA property bytes separately;
 *   - the expected uncompressed size;
 *   - a callback that writes each decoded chunk to flash.
 *
 * This function owns the CLzmaDec state and temporary output buffer for its
 * entire lifetime. No decoder context is retained after the function returns.
 */

#include "LzmaDec.h"


#define LZMA_STREAM_OUTPUT_CHUNK_SIZE 256


/*
 * The callback must persist all bytes before returning.
 * Return SZ_OK on success, or SZ_ERROR_WRITE on failure.
 */
typedef SRes (*LzmaStreamWriteFunc)(
    void *context, const Byte *data, SizeT size);


/*
 * compressedInput starts immediately after the five property bytes.
 * It must remain valid for the complete duration of this function.
 *
 * srcLen:
 *   In : compressed byte count in compressedInput.
 *   Out: number of compressed bytes consumed.
 *
 * total_file_size is the exact expected uncompressed size.
 *
 * The function returns SZ_OK only after all total_file_size bytes have
 * been passed successfully to the write callback.
 */
SRes LzmaStreamDecode(
    const Byte *compressedInput,
    SizeT *srcLen,
    const Byte *propData,
    unsigned propSize,
    UInt64 total_file_size,
    ELzmaStatus *status,
    LzmaStreamWriteFunc write,
    void *writeContext,
    ISzAllocPtr alloc)
{
  CLzmaDec dec;
  Byte outputBuffer[LZMA_STREAM_OUTPUT_CHUNK_SIZE];
  SRes res;
  SizeT remainingInput;
  SizeT consumedTotal = 0;
  UInt64 writtenTotal = 0;

  if (compressedInput == NULL || srcLen == NULL ||
      propData == NULL || status == NULL ||
      write == NULL || alloc == NULL ||
      total_file_size == 0)
    return SZ_ERROR_PARAM;

  remainingInput = *srcLen;
  *srcLen = 0;
  *status = LZMA_STATUS_NOT_SPECIFIED;

  LzmaDec_Construct(&dec);

  /*
   * Allocate the probability tables and the dictionary derived from
   * propData. The dictionary is separate from the flash output.
   */
  res = LzmaDec_Allocate(&dec, propData, propSize, alloc);
  if (res != SZ_OK)
    return res;

  LzmaDec_Init(&dec);

  while (writtenTotal < total_file_size)
  {
    SizeT outputSize = sizeof(outputBuffer);
    SizeT inputSize = remainingInput;
    UInt64 remainingOutput = total_file_size - writtenTotal;

    /*
     * Use FINISH_ANY because the target raw LZMA stream has no end marker.
     * Completion is determined by total_file_size.
     */
    res = LzmaDec_DecodeToBuf(
        &dec,
        outputBuffer, &outputSize,
        compressedInput, &inputSize,
        LZMA_FINISH_ANY,
        status);

    compressedInput += inputSize;
    remainingInput -= inputSize;
    consumedTotal += inputSize;

    if (res != SZ_OK)
      break;

    if ((UInt64)outputSize > remainingOutput)
    {
      res = SZ_ERROR_DATA;
      break;
    }

    if (outputSize > 0)
    {
      res = write(writeContext, outputBuffer, outputSize);
      if (res != SZ_OK)
        break;

      writtenTotal += outputSize;
    }

    if (writtenTotal == total_file_size)
    {
      if (*status != LZMA_STATUS_FINISHED_WITH_MARK)
        *status = LZMA_STATUS_MAYBE_FINISHED_WITHOUT_MARK;
      break;
    }

    if (*status == LZMA_STATUS_FINISHED_WITH_MARK)
    {
      /*
       * End marker encountered before the expected output size.
       */
      res = SZ_ERROR_DATA;
      break;
    }

    /*
     * When compressed input is exhausted, continue once more with zero
     * input so pending matches in the decoder can be emitted. If no more
     * output is produced, the stream is truncated.
     */
    if (remainingInput == 0 && inputSize == 0 && outputSize == 0)
    {
      res = SZ_ERROR_INPUT_EOF;
      break;
    }

    /*
     * Protect against a decoder state that cannot make progress.
     */
    if (inputSize == 0 && outputSize == 0)
    {
      res = SZ_ERROR_DATA;
      break;
    }
  }

  *srcLen = consumedTotal;
  LzmaDec_Free(&dec, alloc);
  return res;
}


/*
 * Example flash callback. Adapt flash_write() to the MCU driver.
 *
 * outputBuffer is reused after this callback returns, so the callback must
 * not retain its pointer.
 */
#if 0
typedef struct
{
  UInt32 flashAddress;
  UInt32 written;
} FlashContext;

static SRes FlashWrite(
    void *context, const Byte *data, SizeT size)
{
  FlashContext *flash = (FlashContext *)context;

  if (!flash_write(flash->flashAddress + flash->written,
                   data, size))
    return SZ_ERROR_WRITE;

  flash->written += (UInt32)size;
  return SZ_OK;
}
#endif


/*
 * Example:
 *
 *   SizeT compressedSize = compressedEnd - compressedInput;
 *   ELzmaStatus status;
 *
 *   result = LzmaStreamDecode(
 *       compressedInput,       // bytes after the five LZMA properties
 *       &compressedSize,
 *       props,                 // five LZMA property bytes
 *       5,
 *       uncompressedSize,      // exact expected output size
 *       &status,
 *       FlashWrite,
 *       &flashContext,
 *       &g_Alloc);
 */