// (C) 2016 - 2020 Tino Reichardt

#include "StdAfx.h"
#include "ZstdDecoder.h"

namespace NCompress {
namespace NZSTD {

CDecoder::CDecoder():
  _ctx(NULL),
  _srcBuf(NULL),
  _dstBuf(NULL),
  _srcBufSize(ZSTD_DStreamInSize()),
  _dstBufSize(ZSTD_DStreamOutSize()),
  _processedIn(0),
  _processedOut(0),
  _readSrcPos(0),
  _readSrcLen(0)
{
  _props.clear();
}

CDecoder::~CDecoder()
{
  if (_ctx) {
    ZSTD_freeDCtx(_ctx);
    MyFree(_srcBuf);
    MyFree(_dstBuf);
  }
}

Z7_COM7F_IMF(CDecoder::SetDecoderProperties2(const Byte * prop, UInt32 size))
{
  switch (size) {
  case 1:
    _props._flags = *prop;
    // flags currently unused, todo: maybe some handling is necessary if set to not 0
    // if (_props._flags != 0) ...
    return S_OK;
  // for backwards compatibility only:
  // size 3 or 5 was accepted in previous versions, so allow it here too
  // (since the old props members _ver/_level seemed to be unused anyway)
  case 3:
  case 5:
    return S_OK;
  default:
    return E_NOTIMPL;
  }
}

HRESULT CDecoder::SetOutStreamSizeResume(const UInt64 * /*outSize*/)
{
  _processedOut = 0;
  return S_OK;
}

Z7_COM7F_IMF(CDecoder::SetOutStreamSize(const UInt64 * outSize))
{
  _processedIn = 0;
  _readSrcPos = 0;
  _readSrcLen = 0;
  RINOK(SetOutStreamSizeResume(outSize))
  return S_OK;
}

HRESULT CDecoder::CodeSpec(ISequentialInStream * inStream,
  ISequentialOutStream * outStream, ICompressProgressInfo * progress)
{
  size_t srcBufLen, result;
  ZSTD_inBuffer zIn;
  ZSTD_outBuffer zOut;

  /* 1) create context */
  if (!_ctx) {
    _ctx = ZSTD_createDCtx();
    if (!_ctx)
      return E_OUTOFMEMORY;

    _srcBuf = MyAlloc(_srcBufSize);
    if (!_srcBuf)
      return E_OUTOFMEMORY;

    _dstBuf = MyAlloc(_dstBufSize);
    if (!_dstBuf)
      return E_OUTOFMEMORY;

    result = ZSTD_DCtx_setParameter(_ctx, ZSTD_d_windowLogMax, ZSTD_WINDOWLOG_MAX);
    if (ZSTD_isError(result))
      return E_OUTOFMEMORY;
  } else {
    result = ZSTD_DCtx_reset(_ctx, ZSTD_reset_session_only);
    if (ZSTD_isError(result))
      return E_FAIL;
  }

  zOut.dst = _dstBuf;
  srcBufLen = _srcBufSize;

  /* read first input block */
  RINOK(ReadStream(inStream, _srcBuf, &srcBufLen))
  _processedIn += srcBufLen;

  zIn.src = _srcBuf;
  zIn.size = srcBufLen;
  zIn.pos = 0;

  /* Main decompression Loop */
  for (;;) {
    for (;;) {
      /* decompress loop */
      zOut.size = _dstBufSize;
      zOut.pos = 0;

      result = ZSTD_decompressStream(_ctx, &zOut, &zIn);
      if (ZSTD_isError(result)) {
        switch (ZSTD_getErrorCode(result)) {
          /* @Igor: would be nice, if we have an API to store the errmsg */
          case ZSTD_error_memory_allocation:
            return E_OUTOFMEMORY;
          case ZSTD_error_frameParameter_unsupported:
          case ZSTD_error_parameter_unsupported:
          case ZSTD_error_version_unsupported:
            return E_NOTIMPL;
          case ZSTD_error_frameParameter_windowTooLarge:
          case ZSTD_error_parameter_outOfBound:
            return E_INVALIDARG;
          default:
            return E_FAIL;
        }
      }

      /* write decompressed result */
      if (zOut.pos) {
        RINOK(WriteStream(outStream, _dstBuf, zOut.pos))
        _processedOut += zOut.pos;
        if (progress)
        {
          RINOK(progress->SetRatioInfo(&_processedIn, &_processedOut))
        }
      }

      /* finished with buffer */
      if (zIn.pos == zIn.size)
        break;

      /* end of frame */
      if (result == 0) {
        result = ZSTD_DCtx_reset(_ctx, ZSTD_reset_session_only);
        if (ZSTD_isError(result))
          return E_FAIL;

        /* end of frame, but more data */
        if (zIn.pos < zIn.size)
          continue;

        /* read next input, or eof */
        break;
      }
    } /* for() decompress */

    /* read next input */
    srcBufLen = _srcBufSize;
    RINOK(ReadStream(inStream, _srcBuf, &srcBufLen))
    _processedIn += srcBufLen;

    /* finished */
    if (srcBufLen == 0)
      return S_OK;

    zIn.size = srcBufLen;
    zIn.pos = 0;
  }
}

Z7_COM7F_IMF(CDecoder::Code(ISequentialInStream * inStream, ISequentialOutStream * outStream,
  const UInt64 * /*inSize */, const UInt64 *outSize, ICompressProgressInfo * progress))
{
  SetOutStreamSize(outSize);
  return CodeSpec(inStream, outStream, progress);
}

#ifndef Z7_NO_READ_FROM_CODER
Z7_COM7F_IMF(CDecoder::SetInStream(ISequentialInStream * inStream))
{
  _inStream = inStream;
  return S_OK;
}

Z7_COM7F_IMF(CDecoder::ReleaseInStream())
{
  _inStream.Release();
  return S_OK;
}

Z7_COM7F_IMF(CDecoder::Read(void *data, UInt32 size, UInt32 *processedSize))
{
  if (processedSize)
    *processedSize = 0;

  /* 1) create/reset context */
  if (!_ctx) {
    _ctx = ZSTD_createDCtx();
    if (!_ctx)
      return E_OUTOFMEMORY;
    _srcBuf = MyAlloc(_srcBufSize);
    if (!_srcBuf)
      return E_OUTOFMEMORY;
    _dstBuf = MyAlloc(_dstBufSize);
    if (!_dstBuf)
      return E_OUTOFMEMORY;
    ZSTD_DCtx_setParameter(_ctx, ZSTD_d_windowLogMax, ZSTD_WINDOWLOG_MAX);
    _readSrcPos = 0;
    _readSrcLen = 0;
  }

  UInt32 totalOut = 0;
  Byte *outPtr = (Byte *)data;

  while (totalOut < size) {
    /* refill input buffer if needed */
    if (_readSrcPos >= _readSrcLen) {
      _readSrcLen = _srcBufSize;
      RINOK(ReadStream(_inStream, _srcBuf, &_readSrcLen))
      _processedIn += _readSrcLen;
      _readSrcPos = 0;
      if (_readSrcLen == 0)
        break; /* EOF */
    }

    ZSTD_inBuffer zIn;
    zIn.src = (const Byte *)_srcBuf + _readSrcPos;
    zIn.size = _readSrcLen - _readSrcPos;
    zIn.pos = 0;

    ZSTD_outBuffer zOut;
    zOut.dst = outPtr + totalOut;
    zOut.size = size - totalOut;
    zOut.pos = 0;

    size_t result = ZSTD_decompressStream(_ctx, &zOut, &zIn);
    if (ZSTD_isError(result))
      return E_FAIL;

    _readSrcPos += zIn.pos;
    totalOut += (UInt32)zOut.pos;
    _processedOut += zOut.pos;

    if (result == 0) {
      /* frame complete */
      ZSTD_DCtx_reset(_ctx, ZSTD_reset_session_only);
    }
  }

  if (processedSize)
    *processedSize = totalOut;
  return S_OK;
}
#endif

Z7_COM7F_IMF(CDecoder::SetNumberOfThreads(UInt32 /* numThreads */))
{
  return S_OK;
}

HRESULT CDecoder::CodeResume(ISequentialOutStream * outStream, const UInt64 * outSize, ICompressProgressInfo * progress)
{
  RINOK(SetOutStreamSizeResume(outSize))
  return CodeSpec(_inStream, outStream, progress);
}

}}
