// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// -----------------------------------------------------------------------------
//
// main entry for the encoder
//
// Author: Skal (pascal.massimino@gmail.com)

#include <cassert>

#include "src/common/constants.h"
#include "src/common/header_enc_dec.h"
#include "src/common/vdebug.h"
#include "src/enc/analysis.h"
#include "src/enc/preview/preview_enc.h"
#include "src/enc/tile_enc.h"
#include "src/enc/wp2_enc_i.h"
#include "src/utils/ans_enc.h"
#include "src/utils/utils.h"
#include "src/wp2/base.h"
#include "src/wp2/encode.h"
#include "src/wp2/format_constants.h"

namespace WP2 {

const EncoderConfig EncoderConfig::kDefault;

//------------------------------------------------------------------------------

namespace {

WP2Status WriteChunk(DataView data, Writer* const output) {
  if (data.size == 0) return WP2_STATUS_OK;
  WP2_CHECK_OK(data.size <= kMaxChunkSize, WP2_STATUS_INVALID_PARAMETER);
  uint8_t buf[kMaxVarIntLength];
  const uint32_t buf_size = WriteVarInt(data.size, 1, kMaxChunkSize, buf);
  WP2_CHECK_OK(output->Append(buf, buf_size), WP2_STATUS_BAD_WRITE);
  WP2_CHECK_OK(output->Append(data.bytes, data.size), WP2_STATUS_BAD_WRITE);
  return WP2_STATUS_OK;
}

bool WriteTag(const uint32_t tag, Writer* const output) {
  const uint8_t data[3] = {(uint8_t)(tag >> 0), (uint8_t)(tag >> 8),
                           (uint8_t)(tag >> 16)};
  return output->Append(data, sizeof(data));
}

//------------------------------------------------------------------------------

bool HasPreview(const EncoderConfig& config) {
  return (config.create_preview || config.preview_size > 0);
}

}  // anonymous namespace

//------------------------------------------------------------------------------

WP2Status SetupEncoderInfo(uint32_t width, uint32_t height,
                           const EncoderConfig& config) {
#if !defined(WP2_REDUCE_BINARY_SIZE)
  if (VDMatch(config, "")) {  // Matches any visual debug.
    ArgbBuffer* const debug_output = &config.info->debug_output;
    WP2_CHECK_STATUS(debug_output->Resize(width, height));
    VDDrawUndefinedPattern(debug_output);
  }
#endif  // WP2_REDUCE_BINARY_SIZE

  if (config.info != nullptr) {
#if defined(WP2_BITTRACE)
    config.info->blocks.clear();
#endif
    config.info->selection_info.clear();
  }
  return WP2_STATUS_OK;
}

//------------------------------------------------------------------------------

// Encode header to 'output'.
WP2Status EncodeHeader(const EncoderConfig& config,
                       uint32_t width, uint32_t height, uint32_t rgb_bit_depth,
                       bool has_alpha, bool is_anim, bool loop_forever,
                       Argb38b background_color, RGB12b preview_color,
                       bool has_icc, bool has_trailing_data,
                       Writer* const output) {
  if (!has_alpha) {
    WP2_CHECK_OK(background_color.a == kAlphaMax, WP2_STATUS_INVALID_PARAMETER);
  }
  WP2_CHECK_OK(rgb_bit_depth == 8 || rgb_bit_depth == 10,
               WP2_STATUS_INVALID_PARAMETER);
  WP2_CHECK_OK(CheckPremultiplied(background_color),
               WP2_STATUS_INVALID_PARAMETER);

  uint8_t header[kHeaderMaxSize];
  BitPacker henc(header, sizeof(header), "image_features/");

  henc.PutBits(kSignature, 24, "signature");

  henc.PutBits(width - 1, kImageDimNumBits, "width_m1");
  henc.PutBits(height - 1, kImageDimNumBits, "height_m1");
  henc.PutBits((uint32_t)config.decoding_orientation, 3, "orientation");
  henc.PutBits(has_alpha ? 1 : 0, 1, "has_alpha");
  henc.PutBits(is_anim ? 1 : 0, 1, "is_animation");

  henc.PutBits(ToUInt32(preview_color), 12, "preview_color");

  henc.PutBits(HasPreview(config) ? 1 : 0, 1, "has_preview");
  henc.PutBits(has_icc ? 1 : 0, 1, "has_icc");
  henc.PutBits(has_trailing_data ? 1 : 0, 1, "has_trailing_data");
  henc.PutBits((rgb_bit_depth == 10) ? 1 : 0, 1, "rgb_bit_depth");

  static_assert(TileShape::TILE_SHAPE_AUTO <= (1 << kTileShapeBits),
                "invalid TILE_SHAPE_AUTO value");
  henc.PutBits(FinalTileShape(config), kTileShapeBits, "tile_shape");

  if (config.transfer_function == WP2_TF_ITU_R_BT2020_10BIT) {
    henc.PutBits(/*value=*/1, /*num_bits=*/1, "default_transfer_function");
  } else {
    henc.PutBits(/*value=*/0, /*num_bits=*/1, "default_transfer_function");
    henc.PutBits((uint32_t)config.transfer_function - 1, 4,
                 "transfer_function");
  }

  if (is_anim) {
    henc.PutBits(loop_forever ? 1 : 0, 1, "loop");
    const bool custom_background =
        (background_color.a != kDefaultBackgroundColor.a ||
         background_color.r != kDefaultBackgroundColor.r ||
         background_color.g != kDefaultBackgroundColor.g ||
         background_color.b != kDefaultBackgroundColor.b);
    henc.PutBits(custom_background ? 1 : 0, 1, "background");
    if (custom_background) {
      if (has_alpha) henc.PutBits(background_color.a, 8, "background");
      henc.PutBits(background_color.r, 10, "background");
      henc.PutBits(background_color.g, 10, "background");
      henc.PutBits(background_color.b, 10, "background");
    }
  }

  henc.Pad();  // Unused bits till the next full aligned byte.
  WP2_CHECK_OK(henc.Ok(), WP2_STATUS_BAD_WRITE);

  assert(henc.Used() >= kHeaderMinSize && henc.Used() <= kHeaderMaxSize);
  WP2_CHECK_OK(output->Append(header, henc.Used()), WP2_STATUS_BAD_WRITE);

  return WP2_STATUS_OK;
}

WP2Status EncodePreview(const ArgbBuffer& buffer, const EncoderConfig& config,
                        const ProgressRange& progress, Writer* const output) {
  // TODO(skal): add size limit on preview chunk length?
  if (config.create_preview) {
    MemoryWriter preview;
    WP2_CHECK_STATUS(EncodePreview(buffer,
                                   PreviewConfig(config.quality, config.effort),
                                   progress, &preview));
    WP2_CHECK_STATUS(WriteChunk({preview.mem_, preview.size_}, output));
  } else {
    if (config.preview_size > 0) {
      WP2_CHECK_OK(config.preview != nullptr, WP2_STATUS_INVALID_PARAMETER);
      WP2_CHECK_STATUS(
          WriteChunk({config.preview, config.preview_size}, output));
    }
    WP2_CHECK_STATUS(progress.AdvanceBy(1.));
  }
  return WP2_STATUS_OK;
}

WP2Status EncodeICC(DataView iccp, Writer* const output) {
  if (iccp.size > 0) {
    if (iccp.size == kPredefinedICC1Size &&
        !memcmp(iccp.bytes, kPredefinedICC1, kPredefinedICC1Size)) {
      // kPredefinedICC1 so only write one byte.
      uint8_t buffer[2] = {/*size=*/1, /*type=*/1};
      if (WriteVarInt(1, 1, kMaxChunkSize, buffer) != 1) assert(false);
      WP2_CHECK_OK(output->Append(buffer, 2), WP2_STATUS_BAD_WRITE);
    } else if (iccp.size == kPredefinedICC2Size &&
               !memcmp(iccp.bytes, kPredefinedICC2, kPredefinedICC2Offset) &&
               !memcmp(iccp.bytes + kPredefinedICC2Offset + 1,
                       kPredefinedICC2 + kPredefinedICC2Offset + 1,
                       kPredefinedICC2Size - (kPredefinedICC2Offset + 1)) &&
               iccp.bytes[kPredefinedICC2Offset] <= 0x01) {
      // kPredefinedICC2/3 so only write one byte.
      const uint8_t type = (iccp.bytes[kPredefinedICC2Offset] == 0x00) ? 2 : 3;
      uint8_t buffer[2] = {/*size=*/1, type};
      if (WriteVarInt(1, 1, kMaxChunkSize, buffer) != 1) assert(false);
      WP2_CHECK_OK(output->Append(buffer, 2), WP2_STATUS_BAD_WRITE);
    } else {
      WP2_CHECK_OK(iccp.size > 1, WP2_STATUS_INVALID_PARAMETER);
      WP2_CHECK_OK(iccp.size <= kMaxChunkSize, WP2_STATUS_INVALID_PARAMETER);

      // Custom ICC so write all bytes.
      uint8_t buffer[kMaxVarIntLength];
      uint32_t buffer_size = WriteVarInt(iccp.size, 1, kMaxChunkSize, buffer);
      WP2_CHECK_OK(output->Append(buffer, buffer_size), WP2_STATUS_BAD_WRITE);
      // TODO(skal): use 0-order ANS coding?
      WP2_CHECK_OK(output->Append(iccp.bytes, iccp.size), WP2_STATUS_BAD_WRITE);
    }
  }
  return WP2_STATUS_OK;
}

WP2Status EncodeGLBL(const EncoderConfig& config, const GlobalParams& gparams,
                     bool image_has_alpha, Writer* const output) {
  ANSEnc enc;
  WP2_CHECK_STATUS(gparams.Write(image_has_alpha, &enc));
  WP2_CHECK_STATUS(enc.AssembleToBitstream(/*clear_tokens=*/true));

  uint8_t buf[kMaxVarIntLength];
  const uint32_t buf_size =
      WriteVarInt(enc.GetBitstreamSize(), 1, kMaxChunkSize, buf);
  WP2_CHECK_OK(output->Append(buf, buf_size), WP2_STATUS_BAD_WRITE);
  WP2_CHECK_STATUS(enc.WriteBitstreamTo(*output));
  return WP2_STATUS_OK;
}

WP2Status EncodeMetadata(const Metadata& metadata, Writer* const output) {
  // put a terminating signature to count extra chunks and
  // finish off with trailing metadata
  const bool has_xmp = (metadata.xmp.size > 0);
  const bool has_exif = (metadata.exif.size > 0);
  const uint32_t content_bits = (has_xmp ? 1u : 0u) | (has_exif ? 2u : 0u);
  if (content_bits > 0) {
    WP2_CHECK_OK(WriteTag(kTagMask | (content_bits << 16), output),
                 WP2_STATUS_BAD_WRITE);
    WP2_CHECK_STATUS(
        WriteChunk({metadata.xmp.bytes, metadata.xmp.size}, output));
    WP2_CHECK_STATUS(
        WriteChunk({metadata.exif.bytes, metadata.exif.size}, output));
  }
  return WP2_STATUS_OK;
}

TileShape FinalTileShape(const EncoderConfig& config) {
  if (config.tile_shape == TILE_SHAPE_AUTO) {
    return config.quality > kMaxLossyQuality ? TILE_SHAPE_SQUARE_256
                                             : TILE_SHAPE_SQUARE_512;
  }
  return config.tile_shape;
}

PartitionMethod FinalPartitionMethod(const EncoderConfig& config,
                                     uint32_t tile_width,
                                     uint32_t tile_height) {
  if (config.partition_method == AUTO_PARTITIONING) {
    if (tile_width <= kMinBlockSizePix && tile_height <= kMinBlockSizePix) {
      return ALL_4X4_PARTITIONING;  // Skip any setup.
    }
    if (tile_width < 32 && tile_height < 32) {
      // Small enough to use slow methods.
      if (config.effort == 0) return ALL_16X16_PARTITIONING;
      if (config.effort <= 2) return MULTIPASS_PARTITIONING;
      if (!config.partition_snapping) return MULTIPASS_PARTITIONING;
      if (config.effort <= 5) return AREA_ENCODE_PARTITIONING;
      if (config.effort <= 7) return TILE_ENCODE_PARTITIONING;
      if (tile_width <= 16 && tile_height <= 16) return EXHAUSTIVE_PARTITIONING;
      return TILE_ENCODE_PARTITIONING;
    }
    if (config.effort == 0) return ALL_16X16_PARTITIONING;
    if (config.effort <= 6) return MULTIPASS_PARTITIONING;
    if (!config.partition_snapping) return MULTIPASS_PARTITIONING;
    return AREA_ENCODE_PARTITIONING;
  }
  return config.partition_method;
}

//------------------------------------------------------------------------------
// Still image encoding function

WP2Status Encode(const ArgbBuffer& input, Writer* output,
                 const EncoderConfig& config, PictureHint picture_hint) {
  WP2_CHECK_OK(output != nullptr, WP2_STATUS_NULL_PARAMETER);

  WP2_CHECK_OK(config.IsValid(), WP2_STATUS_INVALID_CONFIGURATION);

  WP2_CHECK_OK(
      input.format() == WP2_Argb_32 || input.format() == WP2_ARGB_32 ||
          (input.format() == WP2_Argb_38 && config.quality == kMaxQuality),
      WP2_STATUS_INVALID_COLORSPACE);
  WP2_CHECK_OK((input.width() > 0) && (input.height() > 0) &&
                   (input.width() <= kImageDimMax) &&
                   (input.height() <= kImageDimMax),
               WP2_STATUS_BAD_DIMENSION);

  (void)picture_hint;
  WP2_CHECK_OK(output != nullptr, WP2_STATUS_NULL_PARAMETER);

  ProgressWatcher progress(config.progress_hook);
  WP2_CHECK_STATUS(progress.Start());

  WP2_CHECK_STATUS(SetupEncoderInfo(input.width(), input.height(), config));

  const uint32_t rgb_bit_depth = (input.format() == WP2_Argb_38) ? 10 : 8;
  const RGB12b preview_color = GetPreviewColor(input);
  const bool has_alpha = input.HasTransparency();
  const bool has_icc = (input.metadata_.iccp.size > 0);
  const bool has_trailing_data =
      (input.metadata_.xmp.size > 0) || (input.metadata_.exif.size > 0);
  WP2_CHECK_STATUS(EncodeHeader(
      config, input.width(), input.height(), rgb_bit_depth, has_alpha,
      /*is_anim=*/false, /*loop_forever=*/true, kDefaultBackgroundColor,
      preview_color, has_icc, has_trailing_data, output));
  WP2_CHECK_STATUS(progress.AdvanceBy(kOnePercent));

  const ProgressRange preview_progress(&progress, kOnePercent);
  WP2_CHECK_STATUS(EncodePreview(input, config, preview_progress, output));

  WP2_CHECK_STATUS(EncodeICC(
      {input.metadata_.iccp.bytes, input.metadata_.iccp.size}, output));
  WP2_CHECK_STATUS(progress.AdvanceBy(kOnePercent));

  const GlobalParams::Type type = DecideGlobalParamsType(config);
  const bool yuv_is_needed =
      (type == GlobalParams::GP_LOSSY || type == GlobalParams::GP_BOTH);

  YUVPlane yuv_input;
  CSPTransform csp_transform;
  if (yuv_is_needed) {
    WP2_CHECK_STATUS(csp_transform.Init(config.csp_type, input));
    WP2_CHECK_STATUS(yuv_input.Import(input, has_alpha, csp_transform,
                                      /*resize_if_needed=*/true,
                                      /*pad=*/kPredWidth));
  }

  const ProgressRange frames_progress(&progress, kProgressFrames);
  WP2_CHECK_STATUS(frames_progress.AdvanceBy(kOnePercent));  // Simulate ANMF
  WP2_CHECK_STATUS(EncodeTiles(input.width(), input.height(), input, yuv_input,
                               csp_transform, config, has_alpha,
                               frames_progress, output));

  WP2_CHECK_STATUS(EncodeMetadata(input.metadata_, output));
  WP2_CHECK_STATUS(progress.Finish());
  return WP2_STATUS_OK;
}

//------------------------------------------------------------------------------

WP2Status Encode(uint32_t width, uint32_t height,
                 const int16_t* c0_buffer, uint32_t c0_step,
                 const int16_t* c1_buffer, uint32_t c1_step,
                 const int16_t* c2_buffer, uint32_t c2_step,
                 const int16_t* a_buffer, uint32_t a_step,
                 const int16_t ccsp_to_rgb_matrix[9],
                 uint32_t ccsp_to_rgb_shift, Writer* output,
                 const EncoderConfig& config, const Metadata& metadata) {
  WP2_CHECK_OK(ccsp_to_rgb_matrix != nullptr, WP2_STATUS_NULL_PARAMETER);
  WP2_CHECK_OK(output != nullptr, WP2_STATUS_NULL_PARAMETER);
  const CSPMtx ccsp_to_rgb(ccsp_to_rgb_matrix, ccsp_to_rgb_shift);

  WP2_CHECK_OK(config.IsValid(), WP2_STATUS_INVALID_CONFIGURATION);

  WP2_CHECK_OK((width > 0 && height > 0) &&
                   (width <= kImageDimMax && height <= kImageDimMax),
               WP2_STATUS_BAD_DIMENSION);
  bool has_alpha = (a_buffer != nullptr);

  WP2_CHECK_OK(c0_buffer != nullptr, WP2_STATUS_NULL_PARAMETER);
  WP2_CHECK_OK(c1_buffer != nullptr, WP2_STATUS_NULL_PARAMETER);
  WP2_CHECK_OK(c2_buffer != nullptr, WP2_STATUS_NULL_PARAMETER);

  WP2_CHECK_OK(c0_step >= width && c1_step >= width && c2_step >= width,
               WP2_STATUS_BAD_DIMENSION);
  if (has_alpha) {
    WP2_CHECK_OK(a_step >= width, WP2_STATUS_BAD_DIMENSION);
  }
  WP2_CHECK_OK(ccsp_to_rgb.shift <= 16, WP2_STATUS_INVALID_PARAMETER);

  ProgressWatcher progress(config.progress_hook);
  WP2_CHECK_STATUS(progress.Start());

  if (has_alpha) {
    has_alpha = false;  // Make sure at least one pixel is transparent.
    const int16_t* a_row = a_buffer;
    for (uint32_t y = 0; y < height; ++y) {
      for (uint32_t x = 0; x < width; ++x) {
        WP2_CHECK_OK(a_row[x] >= 0 && (uint32_t)a_row[x] <= kAlphaMax,
                     WP2_STATUS_INVALID_PARAMETER);
        if ((uint32_t)a_row[x] < kAlphaMax) has_alpha = true;
      }
      a_row += a_step;
    }
  }

  WP2_CHECK_STATUS(SetupEncoderInfo(width, height, config));

  const GlobalParams::Type type = DecideGlobalParamsType(config);
  const bool yuv_is_needed =
      (type == GlobalParams::GP_LOSSY || type == GlobalParams::GP_BOTH);
  const bool rgb_is_needed =
      (type == GlobalParams::GP_LOSSLESS || type == GlobalParams::GP_BOTH ||
       type == GlobalParams::GP_AV1 ||
       (yuv_is_needed &&
        config.csp_type == Csp::kCustom) ||  // for CSPTransform::Optimize()
       config.create_preview);               // for EncodePreview()
  // TODO(yguyon): Some of these cases could also be done directly in YUV space
  //               instead of needing RGB conversion.

  ArgbBuffer rgb_input(WP2_Argb_32);
  YUVPlane yuv_input;
  CSPTransform csp_transform;
  if (rgb_is_needed) {
    WP2_CHECK_STATUS(rgb_input.Resize(width, height));
    WP2_CHECK_STATUS(CSPTransform::CustomToArgb(
        width, height, c0_buffer, c0_step, c1_buffer, c1_step, c2_buffer,
        c2_step, a_buffer, a_step, ccsp_to_rgb, &rgb_input));
  }
  if (yuv_is_needed) {
    WP2_CHECK_STATUS(csp_transform.Init(config.csp_type, rgb_input));
    WP2_CHECK_STATUS(
        yuv_input.Resize(width, height, /*pad=*/kPredWidth, has_alpha));

    WP2_CHECK_STATUS(csp_transform.CustomToYuv(
        width, height, c0_buffer, c0_step, c1_buffer, c1_step, c2_buffer,
        c2_step, ccsp_to_rgb, yuv_input.Y.Row(0), yuv_input.Y.Step(),
        yuv_input.U.Row(0), yuv_input.U.Step(), yuv_input.V.Row(0),
        yuv_input.V.Step()));
    if (has_alpha) {
      Plane16 non_padded_alpha;
      WP2_CHECK_STATUS(
          non_padded_alpha.SetView(yuv_input.A, {0, 0, width, height}));
      non_padded_alpha.From(a_buffer, a_step);
    }

    WP2_CHECK_STATUS(yuv_input.FillPad(width, height));
  }

  const uint32_t rgb_bit_depth = (rgb_input.format() == WP2_Argb_38) ? 10 : 8;
  const RGB12b preview_color =
      yuv_input.IsEmpty() ? GetPreviewColor(rgb_input)
                          : GetPreviewColor(yuv_input, csp_transform);
  const bool has_icc = (metadata.iccp.size > 0);
  const bool has_trailing_data =
      (metadata.xmp.size > 0 || metadata.exif.size > 0);
  WP2_CHECK_STATUS(EncodeHeader(config, width, height, rgb_bit_depth, has_alpha,
                                /*is_anim=*/false, /*loop_forever=*/true,
                                kDefaultBackgroundColor, preview_color, has_icc,
                                has_trailing_data, output));
  WP2_CHECK_STATUS(progress.AdvanceBy(kOnePercent));

  const ProgressRange preview_progress(&progress, kOnePercent);
  WP2_CHECK_STATUS(EncodePreview(rgb_input, config, preview_progress, output));

  WP2_CHECK_STATUS(
      EncodeICC({metadata.iccp.bytes, metadata.iccp.size}, output));
  WP2_CHECK_STATUS(progress.AdvanceBy(kOnePercent));

  const ProgressRange frames_progress(&progress, kProgressFrames);
  WP2_CHECK_STATUS(frames_progress.AdvanceBy(kOnePercent));  // Simulate ANMF
  WP2_CHECK_STATUS(EncodeTiles(width, height, rgb_input, yuv_input,
                               csp_transform, config, has_alpha,
                               frames_progress, output));

  WP2_CHECK_STATUS(EncodeMetadata(metadata, output));
  WP2_CHECK_STATUS(progress.Finish());
  return WP2_STATUS_OK;
}

//------------------------------------------------------------------------------

}  // namespace WP2