// Copyright (c) the JPEG XL Project
//
// 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
//
//      http://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.

#include "lib/jxl/coeff_order.h"

#include <stdint.h>

#include <algorithm>
#include <vector>

#include "lib/jxl/ans_params.h"
#include "lib/jxl/aux_out.h"
#include "lib/jxl/aux_out_fwd.h"
#include "lib/jxl/base/padded_bytes.h"
#include "lib/jxl/base/profiler.h"
#include "lib/jxl/base/span.h"
#include "lib/jxl/coeff_order_fwd.h"
#include "lib/jxl/dec_ans.h"
#include "lib/jxl/dec_bit_reader.h"
#include "lib/jxl/enc_ans.h"
#include "lib/jxl/enc_bit_writer.h"
#include "lib/jxl/entropy_coder.h"
#include "lib/jxl/lehmer_code.h"
#include "lib/jxl/modular/encoding/encoding.h"
#include "lib/jxl/modular/modular_image.h"

namespace jxl {

void SetDefaultOrder(AcStrategy acs, coeff_order_t* JXL_RESTRICT order) {
  PROFILER_FUNC;
  const size_t size =
      kDCTBlockSize * acs.covered_blocks_x() * acs.covered_blocks_y();
  const coeff_order_t* natural_coeff_order = acs.NaturalCoeffOrder();
  for (size_t k = 0; k < size; ++k) {
    order[k] = natural_coeff_order[k];
  }
}

uint32_t CoeffOrderContext(uint32_t val) {
  uint32_t token, nbits, bits;
  HybridUintConfig(0, 0, 0).Encode(val, &token, &nbits, &bits);
  return std::min(token, kPermutationContexts - 1);
}

namespace {
Status ReadPermutation(size_t skip, size_t size, coeff_order_t* order,
                       BitReader* br, ANSSymbolReader* reader,
                       const std::vector<uint8_t>& context_map) {
  std::vector<LehmerT> lehmer(size);
  // temp space needs to be as large as the next power of 2, so doubling the
  // allocated size is enough.
  std::vector<uint32_t> temp(size * 2);
  uint32_t end =
      reader->ReadHybridUint(CoeffOrderContext(size), br, context_map) + skip;
  if (end > size) {
    return JXL_FAILURE("Invalid permutation size");
  }
  uint32_t last = 0;
  for (size_t i = skip; i < end; ++i) {
    lehmer[i] =
        reader->ReadHybridUint(CoeffOrderContext(last), br, context_map);
    last = lehmer[i];
    if (lehmer[i] + i >= size) {
      return JXL_FAILURE("Invalid lehmer code");
    }
  }
  if (order == nullptr) return true;
  DecodeLehmerCode(lehmer.data(), temp.data(), size, order);
  return true;
}

}  // namespace

Status DecodePermutation(size_t skip, size_t size, coeff_order_t* order,
                         BitReader* br) {
  std::vector<uint8_t> context_map;
  ANSCode code;
  JXL_RETURN_IF_ERROR(
      DecodeHistograms(br, kPermutationContexts, &code, &context_map));
  ANSSymbolReader reader(&code, br);
  JXL_RETURN_IF_ERROR(
      ReadPermutation(skip, size, order, br, &reader, context_map));
  if (!reader.CheckANSFinalState()) {
    return JXL_FAILURE("Invalid ANS stream");
  }
  return true;
}

namespace {

Status DecodeCoeffOrder(AcStrategy acs, coeff_order_t* order, BitReader* br,
                        ANSSymbolReader* reader,
                        const std::vector<uint8_t>& context_map) {
  PROFILER_FUNC;
  const size_t llf = acs.covered_blocks_x() * acs.covered_blocks_y();
  const size_t size = kDCTBlockSize * llf;

  JXL_RETURN_IF_ERROR(
      ReadPermutation(llf, size, order, br, reader, context_map));
  if (order == nullptr) return true;
  const coeff_order_t* natural_coeff_order = acs.NaturalCoeffOrder();
  for (size_t k = 0; k < size; ++k) {
    order[k] = natural_coeff_order[order[k]];
  }
  return true;
}

}  // namespace

Status DecodeCoeffOrders(uint16_t used_orders, uint32_t used_acs,
                         coeff_order_t* order, BitReader* br) {
  uint16_t computed = 0;
  std::vector<uint8_t> context_map;
  ANSCode code;
  std::unique_ptr<ANSSymbolReader> reader;
  // Bitstream does not have histograms if no coefficient order is used.
  if (used_orders != 0) {
    JXL_RETURN_IF_ERROR(
        DecodeHistograms(br, kPermutationContexts, &code, &context_map));
    reader = make_unique<ANSSymbolReader>(&code, br);
  }
  uint32_t acs_mask = 0;
  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {
    if ((used_acs & (1 << o)) == 0) continue;
    acs_mask |= 1 << kStrategyOrder[o];
  }
  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {
    uint8_t ord = kStrategyOrder[o];
    if (computed & (1 << ord)) continue;
    computed |= 1 << ord;
    AcStrategy acs = AcStrategy::FromRawStrategy(o);
    bool used = (acs_mask & (1 << ord)) != 0;
    if ((used_orders & (1 << ord)) == 0) {
      // No need to set the default order if no ACS uses this order.
      if (used) {
        for (size_t c = 0; c < 3; c++) {
          SetDefaultOrder(acs, &order[CoeffOrderOffset(ord, c)]);
        }
      }
    } else {
      for (size_t c = 0; c < 3; c++) {
        coeff_order_t* dest = used ? &order[CoeffOrderOffset(ord, c)] : nullptr;
        JXL_RETURN_IF_ERROR(
            DecodeCoeffOrder(acs, dest, br, reader.get(), context_map));
      }
    }
  }
  if (used_orders && !reader->CheckANSFinalState()) {
    return JXL_FAILURE("Invalid ANS stream");
  }
  return true;
}

}  // namespace jxl
