// 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.

#ifndef LIB_JXL_DEC_MODULAR_H_

#define LIB_JXL_DEC_MODULAR_H_

#include <jxl/memory_manager.h>

#include <cstddef>

#include <cstdint>

#include <string>

#include <vector>

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

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

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

#include "lib/jxl/dec_bit_reader.h"

#include "lib/jxl/dec_cache.h"

#include "lib/jxl/frame_dimensions.h"

#include "lib/jxl/frame_header.h"

#include "lib/jxl/modular/encoding/dec_ma.h"

#include "lib/jxl/modular/encoding/encoding.h"

#include "lib/jxl/modular/modular_image.h"

#include "lib/jxl/quant_weights.h"

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

namespace jxl {

struct ModularStreamId {

  enum class Kind {

    GlobalData,

    VarDCTDC,

    ModularDC,

    ACMetadata,

    QuantTable,

    ModularAC

  };

  Kind kind;

  size_t quant_table_id;

  size_t group_id;  // DC or AC group id.

  size_t pass_id;   // Only for kModularAC.

  size_t ID(const FrameDimensions& frame_dim) const {

    size_t id = 0;

    switch (kind) {

      case Kind::GlobalData:

        id = 0;

        break;

      case Kind::VarDCTDC:

        id = 1 + group_id;

        break;

      case Kind::ModularDC:

        id = 1 + frame_dim.num_dc_groups + group_id;

        break;

      case Kind::ACMetadata:

        id = 1 + 2 * frame_dim.num_dc_groups + group_id;

        break;

      case Kind::QuantTable:

        id = 1 + 3 * frame_dim.num_dc_groups + quant_table_id;

        break;

      case Kind::ModularAC:

        id = 1 + 3 * frame_dim.num_dc_groups + kNumQuantTables +

             frame_dim.num_groups * pass_id + group_id;

        break;

    };

    return id;

  }

  static ModularStreamId Global() {

    return ModularStreamId{Kind::GlobalData, 0, 0, 0};

  }

  static ModularStreamId VarDCTDC(size_t group_id) {

    return ModularStreamId{Kind::VarDCTDC, 0, group_id, 0};

  }

  static ModularStreamId ModularDC(size_t group_id) {

    return ModularStreamId{Kind::ModularDC, 0, group_id, 0};

  }

  static ModularStreamId ACMetadata(size_t group_id) {

    return ModularStreamId{Kind::ACMetadata, 0, group_id, 0};

  }

  static StatusOr<ModularStreamId> QuantTable(size_t quant_table_id) {

    JXL_ENSURE(quant_table_id < kNumQuantTables);

    return ModularStreamId{Kind::QuantTable, quant_table_id, 0, 0};

  }

  static ModularStreamId ModularAC(size_t group_id, size_t pass_id) {

    return ModularStreamId{Kind::ModularAC, 0, group_id, pass_id};

  }

  static size_t Num(const FrameDimensions& frame_dim, size_t passes) {

    return ModularAC(0, passes).ID(frame_dim);

  }

  std::string DebugString() const;

};

class ModularFrameDecoder {

 public:

  explicit ModularFrameDecoder(JxlMemoryManager* memory_manager)

      : memory_manager_(memory_manager), full_image(memory_manager) {}

  void Init(const FrameDimensions& new_frame_dim) { frame_dim = new_frame_dim; }

  Status DecodeGlobalInfo(BitReader* reader, const FrameHeader& frame_header,

                          bool allow_truncated_group);

  Status DecodeGroup(const FrameHeader& frame_header, const Rect& rect,

                     BitReader* reader, int minShift, int maxShift,

                     const ModularStreamId& stream, bool zerofill,

                     PassesDecoderState* dec_state,

                     RenderPipelineInput* render_pipeline_input,

                     bool allow_truncated, bool* should_run_pipeline = nullptr);

  // Decodes a VarDCT DC group (`group_id`) from the given `reader`.

  Status DecodeVarDCTDC(const FrameHeader& frame_header, size_t group_id,

                        BitReader* reader, PassesDecoderState* dec_state);

  // Decodes a VarDCT AC Metadata group (`group_id`) from the given `reader`.

  Status DecodeAcMetadata(const FrameHeader& frame_header, size_t group_id,

                          BitReader* reader, PassesDecoderState* dec_state);

  // Decodes a RAW quant table from `br` into the given `encoding`, of size

  // `required_size_x x required_size_y`. If `modular_frame_decoder` is passed,

  // its global tree is used, otherwise no global tree is used.

  static Status DecodeQuantTable(JxlMemoryManager* memory_manager,

                                 size_t required_size_x, size_t required_size_y,

                                 BitReader* br, QuantEncoding* encoding,

                                 size_t idx,

                                 ModularFrameDecoder* modular_frame_decoder);

  // if inplace is true, this can only be called once

  // if it is false, it can be called multiple times (e.g. for progressive

  // steps)

  Status FinalizeDecoding(const FrameHeader& frame_header,

                          PassesDecoderState* dec_state, jxl::ThreadPool* pool,

                          bool inplace);

  bool have_dc() const { return have_something; }

  void MaybeDropFullImage();

  bool UsesFullImage() const { return use_full_image; }

  JxlMemoryManager* memory_manager() const { return memory_manager_; }

 private:

  Status ModularImageToDecodedRect(const FrameHeader& frame_header, Image& gi,

                                   PassesDecoderState* dec_state,

                                   jxl::ThreadPool* pool,

                                   RenderPipelineInput& render_pipeline_input,

                                   Rect modular_rect) const;

  JxlMemoryManager* memory_manager_;

  Image full_image;

  std::vector<Transform> global_transform;

  FrameDimensions frame_dim;

  bool do_color;

  bool have_something;

  bool use_full_image = true;

  bool all_same_shift;

  Tree tree;

  ANSCode code;

  std::vector<uint8_t> context_map;

  GroupHeader global_header;

};

}  // namespace jxl

#endif  // LIB_JXL_DEC_MODULAR_H_