/*

 * HEIF codec.

 * Copyright (c) 2024 Brad Hards <bradh@frogmouth.net>

 *

 * This file is part of libheif.

 *

 * libheif is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libheif is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.

 */

#include "compression.h"

#if HAVE_BROTLI

const size_t BUF_SIZE = (1 << 18);

#include <brotli/decode.h>

#include <brotli/encode.h>

#include <cstring>

#include <cstdio>

#include <string>

#include <vector>

#include "error.h"

#include "security_limits.h"

Result<std::vector<uint8_t>> decompress_brotli(const std::vector<uint8_t> &compressed_input,

                                               const heif_security_limits* limits)

{

    BrotliDecoderResult result = BROTLI_DECODER_RESULT_ERROR;

    std::vector<uint8_t> buffer(BUF_SIZE, 0);

    size_t available_in = compressed_input.size();

    const std::uint8_t *next_in = reinterpret_cast<const std::uint8_t *>(compressed_input.data());

    size_t available_out = buffer.size();

    std::uint8_t *next_output = buffer.data();

    std::unique_ptr<BrotliDecoderState, void(*)(BrotliDecoderState*)> state(BrotliDecoderCreateInstance(0, 0, 0), BrotliDecoderDestroyInstance);

    std::vector<uint8_t> output;

    // Track the growth of `output` against the security limits. Without this, a

    // high-ratio "decompression bomb" would expand a few KB of input into GBs of

    // output, bypassing max_memory_block_size / max_total_memory (GHSA-24wx-9w62-c96w).

    MemoryHandle output_memory_handle;

    while (true)

    {

        result = BrotliDecoderDecompressStream(state.get(), &available_in, &next_in, &available_out, &next_output, 0);

        if (result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT)

        {

            size_t n_new_bytes = static_cast<size_t>(std::distance(buffer.data(), next_output));

            if (Error memErr = output_memory_handle.alloc(n_new_bytes, limits, "brotli decompression output")) {

                return memErr;

            }

            output.insert(output.end(), buffer.data(), buffer.data() + n_new_bytes);

            available_out = buffer.size();

            next_output = buffer.data();

        }

        else if (result == BROTLI_DECODER_RESULT_SUCCESS)

        {

            size_t n_new_bytes = static_cast<size_t>(std::distance(buffer.data(), next_output));

            if (Error memErr = output_memory_handle.alloc(n_new_bytes, limits, "brotli decompression output")) {

                return memErr;

            }

            output.insert(output.end(), buffer.data(), buffer.data() + n_new_bytes);

            break;

        }

        else if (result == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT)

        {

            std::stringstream sstr;

            sstr << "Error performing brotli inflate - insufficient data.\n";

            return Error(heif_error_Invalid_input, heif_suberror_Decompression_invalid_data, sstr.str());

        }

        else if (result == BROTLI_DECODER_RESULT_ERROR)

        {

            const char* errorMessage = BrotliDecoderErrorString(BrotliDecoderGetErrorCode(state.get()));

            std::stringstream sstr;

            sstr << "Error performing brotli inflate - " << errorMessage << "\n";

            return Error(heif_error_Invalid_input, heif_suberror_Decompression_invalid_data, sstr.str());

        }

        else

        {

            const char* errorMessage = BrotliDecoderErrorString(BrotliDecoderGetErrorCode(state.get()));

            std::stringstream sstr;

            sstr << "Unknown error performing brotli inflate - " << errorMessage << "\n";

            return Error(heif_error_Invalid_input, heif_suberror_Decompression_invalid_data, sstr.str());

        }

    }

    return output;

}

std::vector<uint8_t> compress_brotli(const uint8_t* input, size_t size)

{

  std::unique_ptr<BrotliEncoderState, void(*)(BrotliEncoderState*)> state(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr), BrotliEncoderDestroyInstance);

  size_t available_in = size;

  const uint8_t* next_in = input;

  std::vector<uint8_t> tmp(BUF_SIZE);

  size_t available_out = BUF_SIZE;

  uint8_t* next_out = tmp.data();

  std::vector<uint8_t> result;

  for (;;) {

    BROTLI_BOOL success = BrotliEncoderCompressStream(state.get(),

                                                      BROTLI_OPERATION_FINISH,

                                                      &available_in,

                                                      &next_in,

                                                      &available_out,

                                                      &next_out,

                                                      nullptr);

    if (!success) {

      return {};

    }

    if (next_out != tmp.data()) {

      result.insert(result.end(), tmp.data(), tmp.data() + std::distance(tmp.data(), next_out));

      available_out = BUF_SIZE;

      next_out = tmp.data();

    }

    if (BrotliEncoderIsFinished(state.get())) {

      break;

    }

  }

  return result;

}

#endif