// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

#include "lib/jxl/dec_cache.h"

#include <jxl/memory_manager.h>

#include <algorithm>

#include <cstddef>

#include <cstdint>

#include <utility>

#include "lib/jxl/ac_strategy.h"

#include "lib/jxl/base/bits.h"

#include "lib/jxl/base/data_parallel.h"

#include "lib/jxl/base/status.h"

#include "lib/jxl/blending.h"

#include "lib/jxl/cms/color_encoding_cms.h"

#include "lib/jxl/coeff_order.h"

#include "lib/jxl/color_encoding_internal.h"

#include "lib/jxl/common.h"  // JXL_HIGH_PRECISION

#include "lib/jxl/frame_dimensions.h"

#include "lib/jxl/frame_header.h"

#include "lib/jxl/image.h"

#include "lib/jxl/image_bundle.h"

#include "lib/jxl/image_metadata.h"

#include "lib/jxl/loop_filter.h"

#include "lib/jxl/memory_manager_internal.h"

#include "lib/jxl/render_pipeline/render_pipeline.h"

#include "lib/jxl/render_pipeline/stage_blending.h"

#include "lib/jxl/render_pipeline/stage_chroma_upsampling.h"

#include "lib/jxl/render_pipeline/stage_cms.h"

#include "lib/jxl/render_pipeline/stage_epf.h"

#include "lib/jxl/render_pipeline/stage_from_linear.h"

#include "lib/jxl/render_pipeline/stage_gaborish.h"

#include "lib/jxl/render_pipeline/stage_noise.h"

#include "lib/jxl/render_pipeline/stage_patches.h"

#include "lib/jxl/render_pipeline/stage_splines.h"

#include "lib/jxl/render_pipeline/stage_spot.h"

#include "lib/jxl/render_pipeline/stage_to_linear.h"

#include "lib/jxl/render_pipeline/stage_tone_mapping.h"

#include "lib/jxl/render_pipeline/stage_upsampling.h"

#include "lib/jxl/render_pipeline/stage_write.h"

#include "lib/jxl/render_pipeline/stage_xyb.h"

#include "lib/jxl/render_pipeline/stage_ycbcr.h"

namespace jxl {

Status GroupDecCache::InitOnce(JxlMemoryManager* memory_manager,

                               size_t num_passes, size_t used_acs) {

  for (size_t i = 0; i < num_passes; i++) {

    if (num_nzeroes[i].xsize() == 0) {

      // Allocate enough for a whole group - partial groups on the

      // right/bottom border just use a subset. The valid size is passed via

      // Rect.

      JXL_ASSIGN_OR_RETURN(num_nzeroes[i],

                           Image3I::Create(memory_manager, kGroupDimInBlocks,

                                           kGroupDimInBlocks));

    }

  }

  size_t max_block_area = 0;

  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {

    AcStrategy acs = AcStrategy::FromRawStrategy(o);

    if ((used_acs & (1 << o)) == 0) continue;

    size_t area =

        acs.covered_blocks_x() * acs.covered_blocks_y() * kDCTBlockSize;

    max_block_area = std::max(area, max_block_area);

  }

  if (max_block_area > max_block_area_) {

    max_block_area_ = max_block_area;

    // We need 3x float blocks for dequantized coefficients and 1x for scratch

    // space for transforms.

    JXL_ASSIGN_OR_RETURN(

        float_memory_,

        AlignedMemory::Create(memory_manager,

                              max_block_area_ * 7 * sizeof(float)));

    // We need 3x int32 or int16 blocks for quantized coefficients.

    JXL_ASSIGN_OR_RETURN(

        int32_memory_,

        AlignedMemory::Create(memory_manager,

                              max_block_area_ * 3 * sizeof(int32_t)));

    JXL_ASSIGN_OR_RETURN(

        int16_memory_,

        AlignedMemory::Create(memory_manager,

                              max_block_area_ * 3 * sizeof(int16_t)));

  }

  dec_group_block = float_memory_.address<float>();

  scratch_space = dec_group_block + max_block_area_ * 3;

  dec_group_qblock = int32_memory_.address<int32_t>();

  dec_group_qblock16 = int16_memory_.address<int16_t>();

  return true;

}

// Initialize the decoder state after all of DC is decoded.

Status PassesDecoderState::InitForAC(size_t num_passes, ThreadPool* pool) {

  shared_storage.coeff_order_size = 0;

  for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {

    if (((1 << o) & used_acs) == 0) continue;

    uint8_t ord = kStrategyOrder[o];

    shared_storage.coeff_order_size =

        std::max(kCoeffOrderOffset[3 * (ord + 1)] * kDCTBlockSize,

                 shared_storage.coeff_order_size);

  }

  size_t sz = num_passes * shared_storage.coeff_order_size;

  if (sz > shared_storage.coeff_orders.size()) {

    shared_storage.coeff_orders.resize(sz);

  }

  return true;

}

Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,

                                           const ImageMetadata* metadata,

                                           ImageBundle* decoded,

                                           PipelineOptions options) {

  JxlMemoryManager* memory_manager = this->memory_manager();

  size_t num_c = 3 + frame_header.nonserialized_metadata->m.num_extra_channels;

  size_t num_tmp_c = options.render_noise ? 3 : 0;

  if (frame_header.CanBeReferenced()) {

    // Necessary so that SetInputSizes() can allocate output buffers as needed.

    frame_storage_for_referencing = ImageBundle(memory_manager, metadata);

  }

  RenderPipeline::Builder builder(memory_manager, num_c + num_tmp_c);

  if (options.use_slow_render_pipeline) {

    builder.UseSimpleImplementation();

  }

  if (!frame_header.chroma_subsampling.Is444()) {

    for (size_t c = 0; c < 3; c++) {

      if (frame_header.chroma_subsampling.HShift(c) != 0) {

        JXL_RETURN_IF_ERROR(

            builder.AddStage(GetChromaUpsamplingStage(c, /*horizontal=*/true)));

      }

      if (frame_header.chroma_subsampling.VShift(c) != 0) {

        JXL_RETURN_IF_ERROR(builder.AddStage(

            GetChromaUpsamplingStage(c, /*horizontal=*/false)));

      }

    }

  }

  if (frame_header.loop_filter.gab) {

    JXL_RETURN_IF_ERROR(

        builder.AddStage(GetGaborishStage(frame_header.loop_filter)));

  }

  {

    const LoopFilter& lf = frame_header.loop_filter;

    if (lf.epf_iters >= 3) {

      JXL_RETURN_IF_ERROR(

          builder.AddStage(GetEPFStage(lf, sigma, EpfStage::Zero)));

    }

    if (lf.epf_iters >= 1) {

      JXL_RETURN_IF_ERROR(

          builder.AddStage(GetEPFStage(lf, sigma, EpfStage::One)));

    }

    if (lf.epf_iters >= 2) {

      JXL_RETURN_IF_ERROR(

          builder.AddStage(GetEPFStage(lf, sigma, EpfStage::Two)));

    }

  }

  bool late_ec_upsample = frame_header.upsampling != 1;

  for (auto ecups : frame_header.extra_channel_upsampling) {

    if (ecups != frame_header.upsampling) {

      // If patches are applied, either frame_header.upsampling == 1 or

      // late_ec_upsample is true.

      late_ec_upsample = false;

    }

  }

  if (!late_ec_upsample) {

    for (size_t ec = 0; ec < frame_header.extra_channel_upsampling.size();

         ec++) {

      if (frame_header.extra_channel_upsampling[ec] != 1) {

        JXL_RETURN_IF_ERROR(builder.AddStage(GetUpsamplingStage(

            memory_manager, frame_header.nonserialized_metadata->transform_data,

            3 + ec,

            CeilLog2Nonzero(frame_header.extra_channel_upsampling[ec]))));

      }

    }

  }

  if ((frame_header.flags & FrameHeader::kPatches) != 0) {

    JXL_RETURN_IF_ERROR(builder.AddStage(GetPatchesStage(

        &shared->image_features.patches,

        &frame_header.nonserialized_metadata->m.extra_channel_info)));

  }

  if ((frame_header.flags & FrameHeader::kSplines) != 0) {

    JXL_RETURN_IF_ERROR(

        builder.AddStage(GetSplineStage(&shared->image_features.splines)));

  }

  if (frame_header.upsampling != 1) {

    size_t nb_channels =

        3 +

        (late_ec_upsample ? frame_header.extra_channel_upsampling.size() : 0);

    for (size_t c = 0; c < nb_channels; c++) {

      JXL_RETURN_IF_ERROR(builder.AddStage(GetUpsamplingStage(

          memory_manager, frame_header.nonserialized_metadata->transform_data,

          c, CeilLog2Nonzero(frame_header.upsampling))));

    }

  }

  // Starting from this line all the stages considered to have zero xextra.

  // Upsampling does not have xextra as well (even if it happens before

  // splines/patches for EC).

  if (options.render_noise) {

    JXL_RETURN_IF_ERROR(builder.AddStage(GetConvolveNoiseStage(num_c)));

    JXL_RETURN_IF_ERROR(builder.AddStage(GetAddNoiseStage(

        shared->image_features.noise_params, shared->cmap.base(), num_c)));

  }

  if (frame_header.dc_level != 0) {

    JXL_RETURN_IF_ERROR(builder.AddStage(GetWriteToImage3FStage(

        memory_manager, &shared_storage.dc_frames[frame_header.dc_level - 1])));

  }

  if (frame_header.CanBeReferenced() &&

      frame_header.save_before_color_transform) {

    JXL_RETURN_IF_ERROR(builder.AddStage(GetWriteToImageBundleStage(

        &frame_storage_for_referencing, output_encoding_info)));

  }

  bool has_alpha = false;

  size_t alpha_c = 0;

  for (size_t i = 0; i < metadata->extra_channel_info.size(); i++) {

    if (metadata->extra_channel_info[i].type == ExtraChannel::kAlpha) {

      has_alpha = true;

      alpha_c = 3 + i;

      break;

    }

  }

  if (fast_xyb_srgb8_conversion) {

#if !JXL_HIGH_PRECISION

    JXL_ENSURE(!NeedsBlending(frame_header));

    JXL_ENSURE(!frame_header.CanBeReferenced() ||

               frame_header.save_before_color_transform);

    JXL_ENSURE(!options.render_spotcolors ||

               !metadata->Find(ExtraChannel::kSpotColor));

    bool is_rgba = (main_output.format.num_channels == 4);

    uint8_t* rgb_output = reinterpret_cast<uint8_t*>(main_output.buffer);

    JXL_RETURN_IF_ERROR(builder.AddStage(

        GetFastXYBTosRGB8Stage(rgb_output, main_output.stride, width, height,

                               is_rgba, has_alpha, alpha_c)));

#endif

  } else {

    bool linear = false;

    if (frame_header.color_transform == ColorTransform::kYCbCr) {

      JXL_RETURN_IF_ERROR(builder.AddStage(GetYCbCrStage()));

    } else if (frame_header.color_transform == ColorTransform::kXYB) {

      JXL_RETURN_IF_ERROR(builder.AddStage(GetXYBStage(output_encoding_info)));

      if (output_encoding_info.color_encoding.GetColorSpace() !=

          ColorSpace::kXYB) {

        linear = true;

      }

    }  // Nothing to do for kNone.

    if (options.coalescing && NeedsBlending(frame_header)) {

      if (linear) {

        JXL_RETURN_IF_ERROR(

            builder.AddStage(GetFromLinearStage(output_encoding_info)));

        linear = false;

      }

      JXL_RETURN_IF_ERROR(builder.AddStage(GetBlendingStage(

          frame_header, this, output_encoding_info.color_encoding)));

    }

    if (options.coalescing && frame_header.CanBeReferenced() &&

        !frame_header.save_before_color_transform) {

      if (linear) {

        JXL_RETURN_IF_ERROR(

            builder.AddStage(GetFromLinearStage(output_encoding_info)));

        linear = false;

      }

      JXL_RETURN_IF_ERROR(builder.AddStage(GetWriteToImageBundleStage(

          &frame_storage_for_referencing, output_encoding_info)));

    }

    if (options.render_spotcolors &&

        frame_header.nonserialized_metadata->m.Find(ExtraChannel::kSpotColor)) {

      for (size_t i = 0; i < metadata->extra_channel_info.size(); i++) {

        // Don't use Find() because there may be multiple spot color channels.

        const ExtraChannelInfo& eci = metadata->extra_channel_info[i];

        if (eci.type == ExtraChannel::kSpotColor) {

          JXL_RETURN_IF_ERROR(

              builder.AddStage(GetSpotColorStage(i, eci.spot_color)));

        }

      }

    }

    auto tone_mapping_stage = GetToneMappingStage(output_encoding_info);

    if (tone_mapping_stage) {

      if (!linear) {

        auto to_linear_stage = GetToLinearStage(output_encoding_info);

        if (!to_linear_stage) {

          if (!output_encoding_info.cms_set) {

            return JXL_FAILURE("Cannot tonemap this colorspace without a CMS");

          }

          auto cms_stage = GetCmsStage(output_encoding_info);

          if (cms_stage) {

            JXL_RETURN_IF_ERROR(builder.AddStage(std::move(cms_stage)));

          }

        } else {

          JXL_RETURN_IF_ERROR(builder.AddStage(std::move(to_linear_stage)));

        }

        linear = true;

      }

      JXL_RETURN_IF_ERROR(builder.AddStage(std::move(tone_mapping_stage)));

    }

    if (linear) {

      const size_t channels_src =

          (output_encoding_info.orig_color_encoding.IsCMYK()

               ? 4

               : output_encoding_info.orig_color_encoding.Channels());

      const size_t channels_dst =

          output_encoding_info.color_encoding.Channels();

      bool mixing_color_and_grey = (channels_dst != channels_src);

      if ((output_encoding_info.color_encoding_is_original) ||

          (!output_encoding_info.cms_set) || mixing_color_and_grey) {

        // in those cases we only need a linear stage in other cases we attempt

        // to obtain a cms stage: the cases are

        // - output_encoding_info.color_encoding_is_original: no cms stage

        // needed because it would be a no-op

        // - !output_encoding_info.cms_set: can't use the cms, so no point in

        // trying to add a cms stage

        // - mixing_color_and_grey: cms stage can't handle that

        // TODO(firsching): remove "mixing_color_and_grey" condition after

        // adding support for greyscale to cms stage.

        JXL_RETURN_IF_ERROR(

            builder.AddStage(GetFromLinearStage(output_encoding_info)));

      } else {

        if (!output_encoding_info.linear_color_encoding.CreateICC()) {

          return JXL_FAILURE("Failed to create ICC");

        }

        auto cms_stage = GetCmsStage(output_encoding_info);

        if (cms_stage) {

          JXL_RETURN_IF_ERROR(builder.AddStage(std::move(cms_stage)));

        }

      }

      linear = false;

    } else {

      auto cms_stage = GetCmsStage(output_encoding_info, false);

      if (cms_stage) {

        JXL_RETURN_IF_ERROR(builder.AddStage(std::move(cms_stage)));

      }

    }

    (void)linear;

    if (main_output.callback.IsPresent() || main_output.buffer) {

      JXL_RETURN_IF_ERROR(builder.AddStage(GetWriteToOutputStage(

          main_output, width, height, has_alpha, unpremul_alpha, alpha_c,

          undo_orientation, extra_output, memory_manager)));

    } else {

      JXL_RETURN_IF_ERROR(builder.AddStage(

          GetWriteToImageBundleStage(decoded, output_encoding_info)));

    }

  }

  JXL_ASSIGN_OR_RETURN(render_pipeline,

                       std::move(builder).Finalize(shared->frame_dim));

  return render_pipeline->IsInitialized();

}

}  // namespace jxl