// 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.
// -----------------------------------------------------------------------------
//
// WP2 lossy decoding of the alpha plane.
//
// Author: Maryla (maryla@google.com)

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstring>

#include "src/common/global_params.h"
#include "src/common/lossy/block.h"
#include "src/common/symbols.h"
#include "src/dec/symbols_dec.h"
#include "src/dec/tile_dec.h"
#include "src/dec/wp2_dec_i.h"
#include "src/utils/ans.h"
#include "src/utils/ans_utils.h"
#include "src/utils/plane.h"
#include "src/utils/utils.h"
#include "src/wp2/base.h"
#include "src/wp2/decode.h"
#include "src/wp2/format_constants.h"

namespace WP2 {

AlphaReader::AlphaReader(ANSDec* const dec, const Tile& tile)
    : dec_(dec), lossless_decoder_() {
  alpha_tile_.chunk_size_is_known = tile.chunk_size_is_known;
  alpha_tile_.chunk_size = tile.chunk_size;
  alpha_tile_.input = tile.input;
  alpha_tile_.rect = tile.rect;
  alpha_tile_.num_decoded_rows = 0;
  // Alpha values cannot be premultiplied by themselves.
  WP2_ASSERT_STATUS(alpha_tile_.rgb_output.SetFormat(WP2_ARGB_32));
}

WP2Status AlphaReader::Allocate() {
  return alpha_tile_.rgb_output.Resize(alpha_tile_.rect.width,
                                       alpha_tile_.rect.height);
}

WP2Status AlphaReader::ReadHeader(const GlobalParams& gparams) {
  gparams_ = &gparams;
  ANSDebugPrefix prefix(dec_, "Alpha");
  alpha_mode_ = gparams_->maybe_use_lossy_alpha_
                    ? (AlphaMode)dec_->ReadRValue(kAlphaModeNum, "alpha_mode")
                    : kAlphaModeLossless;
  if (alpha_mode_ == kAlphaModeLossy) {
    WP2_CHECK_STATUS(
        mode_predictor_.Init(alpha_tile_.rect.width, alpha_tile_.rect.height));
  }
  return WP2_STATUS_OK;
}

void AlphaReader::GetBlockHeader(SymbolReader* const sr,
                                 CodedBlockBase* const cb) {
  if (alpha_mode_ == kAlphaModeLossy) {
    cb->alpha_mode_ = (BlockAlphaMode)mode_predictor_.Read(*cb, sr);
  } else {
    cb->alpha_mode_ = kBlockAlphaLossless;
  }
}

WP2Status AlphaReader::GetBlock(CodedBlockBase* const cb) {
  ANSDebugPrefix prefix(dec_, "Alpha");
  if (alpha_mode_ == kAlphaModeLossless) {
    WP2_CHECK_STATUS(ReadLossless(cb));
  }
  return WP2_STATUS_OK;
}

void AlphaReader::Reconstruct(CodedBlockBase* const cb) const {
  switch (cb->alpha_mode_) {
    case kBlockAlphaLossless:
      ReconstructLossless(*cb, &cb->out_.A);
      break;
    case kBlockAlphaLossy:
      cb->Reconstruct(kAChannel, /*tf_i=*/0);
      break;
    case kBlockAlphaFullTransp:
      cb->out_.A.Fill(0);
      break;
    case kBlockAlphaFullOpaque:
      cb->out_.A.Fill(kAlphaMax);
      break;
  }
}

void AlphaReader::ReconstructLossless(const CodedBlockBase& cb,
                                      Plane16* const plane) const {
  // Copy from buffer to plane, possibly replicating the right/bottom samples.
  const uint32_t block_width = cb.w_pix();
  const uint32_t block_height = cb.h_pix();
  const uint32_t max_x =
      std::min(block_width, alpha_tile_.rect.width - cb.x_pix());
  const uint32_t max_y =
      std::min(block_height, alpha_tile_.rect.height - cb.y_pix());
  assert(max_x > 0 && max_y > 0);
  for (uint32_t y = 0; y < max_y; ++y) {
    const uint8_t* const alpha_row =
        alpha_tile_.rgb_output.GetRow8(cb.y_pix() + y) + 4 * cb.x_pix() + 1;
    for (uint32_t x = 0; x < max_x; ++x) plane->At(x, y) = alpha_row[4 * x];
    // replicate right samples
    for (uint32_t x = max_x; x < block_width; ++x) {
      plane->At(x, y) = alpha_row[4 * (max_x - 1)];
    }
  }
  // replicate bottom samples
  for (uint32_t y = max_y; y < block_height; ++y) {
    memcpy(plane->Row(y), plane->Row(max_y - 1), block_width * sizeof(int16_t));
  }
}

WP2Status AlphaReader::ReadLossless(CodedBlockBase* const cb) {
  const uint32_t max_line =
      std::min(cb->y_pix() + cb->h_pix(), alpha_tile_.rect.height);

  if (next_line_ < max_line) {
    // Read as many new lines as we need for this block.
    const uint32_t lines_to_read = max_line - next_line_;
    const uint32_t num_used_bytes_before = dec_->GetNumUsedBytes();

    if (next_line_ == 0) {
      // Initialize the lossless decoder and read the header.
      lossless_decoder_.Init(DecoderConfig(), *gparams_,
                             /*image_is_premultiplied=*/false, dec_,
                             &alpha_tile_);
      WP2_CHECK_STATUS(lossless_decoder_.DecodeHeader());
    }

    WP2_CHECK_STATUS(lossless_decoder_.DecodeLines(lines_to_read));

    const uint32_t min_num_used_bytes = ANSDec::GetMinNumUsedBytesDiff(
        num_used_bytes_before, dec_->GetNumUsedBytes());
    cb->alpha_lossless_bytes_ = min_num_used_bytes;

    next_line_ = max_line;
    assert(next_line_ > 0);  // so that we won't init the decoder again
  } else {
    cb->alpha_lossless_bytes_ = 0;
  }

  return WP2_STATUS_OK;
}

}  // namespace WP2