// 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/passes_state.h"

#include <jxl/memory_manager.h>

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

#include "lib/jxl/chroma_from_luma.h"

#include "lib/jxl/coeff_order.h"

#include "lib/jxl/frame_dimensions.h"

namespace jxl {

Status InitializePassesSharedState(const FrameHeader& frame_header,

                                   PassesSharedState* JXL_RESTRICT shared,

                                   bool encoder) {

  JXL_ENSURE(frame_header.nonserialized_metadata != nullptr);

  shared->metadata = frame_header.nonserialized_metadata;

  shared->frame_dim = frame_header.ToFrameDimensions();

  shared->image_features.patches.SetShared(&shared->reference_frames);

  const FrameDimensions& frame_dim = shared->frame_dim;

  JxlMemoryManager* memory_manager = shared->memory_manager;

  JXL_ASSIGN_OR_RETURN(

      shared->ac_strategy,

      AcStrategyImage::Create(memory_manager, frame_dim.xsize_blocks,

                              frame_dim.ysize_blocks));

  JXL_ASSIGN_OR_RETURN(shared->raw_quant_field,

                       ImageI::Create(memory_manager, frame_dim.xsize_blocks,

                                      frame_dim.ysize_blocks));

  JXL_ASSIGN_OR_RETURN(shared->epf_sharpness,

                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,

                                      frame_dim.ysize_blocks));

  JXL_ASSIGN_OR_RETURN(

      shared->cmap, ColorCorrelationMap::Create(memory_manager, frame_dim.xsize,

                                                frame_dim.ysize));

  // In the decoder, we allocate coeff orders afterwards, when we know how many

  // we will actually need.

  shared->coeff_order_size = kCoeffOrderMaxSize;

  if (encoder &&

      shared->coeff_orders.size() <

          frame_header.passes.num_passes * kCoeffOrderMaxSize &&

      frame_header.encoding == FrameEncoding::kVarDCT) {

    shared->coeff_orders.resize(frame_header.passes.num_passes *

                                kCoeffOrderMaxSize);

  }

  JXL_ASSIGN_OR_RETURN(shared->quant_dc,

                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,

                                      frame_dim.ysize_blocks));

  bool use_dc_frame = ((frame_header.flags & FrameHeader::kUseDcFrame) != 0u);

  if (!encoder && use_dc_frame) {

    if (frame_header.dc_level == 4) {

      return JXL_FAILURE("Invalid DC level for kUseDcFrame: %u",

                         frame_header.dc_level);

    }

    shared->dc_storage = Image3F();

    shared->dc = &shared->dc_frames[frame_header.dc_level];

    if (shared->dc->xsize() == 0) {

      return JXL_FAILURE(

          "kUseDcFrame specified for dc_level %u, but no frame was decoded "

          "with level %u",

          frame_header.dc_level, frame_header.dc_level + 1);

    }

    ZeroFillImage(&shared->quant_dc);

  } else {

    JXL_ASSIGN_OR_RETURN(shared->dc_storage,

                         Image3F::Create(memory_manager, frame_dim.xsize_blocks,

                                         frame_dim.ysize_blocks));

    shared->dc = &shared->dc_storage;

  }

  return true;

}

}  // namespace jxl