/*

 * H.265 video codec.

 * Copyright (c) 2013-2014 struktur AG, Dirk Farin <farin@struktur.de>

 *

 * This file is part of libde265.

 *

 * libde265 is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libde265 is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libde265.  If not, see <http://www.gnu.org/licenses/>.

 */

#ifndef DE265_DECCTX_H

#define DE265_DECCTX_H

#include "libde265/vps.h"

#include "libde265/sps.h"

#include "libde265/pps.h"

#include "libde265/nal.h"

#include "libde265/slice.h"

#include "libde265/image.h"

#include "libde265/motion.h"

#include "libde265/de265.h"

#include "libde265/dpb.h"

#include "libde265/sei.h"

#include "libde265/threads.h"

#include "libde265/acceleration.h"

#include "libde265/nal-parser.h"

#include <array>

#include <memory>

#include <mutex>

constexpr int DE265_MAX_VPS_SETS = 16;   // this is the maximum as defined in the standard

constexpr int DE265_MAX_SPS_SETS = 16;   // this is the maximum as defined in the standard

constexpr int DE265_MAX_PPS_SETS = 64;   // this is the maximum as defined in the standard

constexpr int MAX_WARNINGS = 20;

class slice_segment_header;

class image_unit;

class slice_unit;

class decoder_context;

class thread_context

{

public:

  thread_context();

  uint32_t CtbAddrInRS;

  uint32_t CtbAddrInTS;

  uint16_t CtbX, CtbY;

  // motion vectors

  PBMotionCoding motion;

  // prediction

  // enum IntraPredMode IntraPredModeC[4]; // chroma intra-prediction mode for current CB

  int ResScaleVal;

  // residual data

  uint8_t cu_transquant_bypass_flag;

  uint8_t transform_skip_flag[3];

  uint8_t explicit_rdpcm_flag;

  uint8_t explicit_rdpcm_dir;

  // we need 16 bytes of extra memory (8*int16) to shift the base for the

  // alignment required for SSE code !

  int16_t _coeffBuf[(32*32)+8];

  int16_t *coeffBuf; // the base pointer for into _coeffBuf, aligned to 16 bytes

  int16_t coeffList[3][32*32];

  int16_t coeffPos[3][32*32];

  int16_t nCoeff[3];

  int32_t residual_luma[32*32]; // only used when cross-comp-prediction is enabled

  // quantization

  int IsCuQpDeltaCoded = 0;

  int CuQpDelta = 0;

  int IsCuChromaQpOffsetCoded = 0;

  int CuQpOffsetCb = 0, CuQpOffsetCr = 0;

  int currentQPY;

  int currentQG_x, currentQG_y;

  int lastQPYinPreviousQG;

  int qPYPrime, qPCbPrime, qPCrPrime;

  CABAC_decoder cabac_decoder;

  context_model_table ctx_model;

  uint8_t StatCoeff[4];

  decoder_context* decctx = nullptr;

  struct de265_image *img = nullptr;

  slice_segment_header* shdr = nullptr;

  image_unit* imgunit = nullptr;

  slice_unit* sliceunit = nullptr;

  thread_task* task; // executing thread_task or nullptr if not multi-threaded

  thread_context(const thread_context&) = delete;

  thread_context& operator=(const thread_context&) = delete;

};

class error_queue

{

 public:

  void add_warning(de265_error warning, bool once);

  de265_error get_warning();

 private:

  std::mutex m_mutex;

  std::vector<de265_error> warnings;

  std::vector<de265_error> warnings_shown; // warnings that have already occurred

};

class slice_unit

{

public:

  slice_unit(decoder_context* decctx);

  ~slice_unit();

  std::unique_ptr<NAL_unit> nal;   // we are the owner

  slice_segment_header* shdr;  // not the owner (de265_image is owner)

  bitreader reader;

  image_unit* imgunit;

  bool flush_reorder_buffer;

  // decoding status

  enum SliceDecodingProgress { Unprocessed,

                               InProgress,

                               Decoded

  } state;

  de265_progress_lock finished_threads;

  int nThreads;

  int first_decoded_CTB_RS; // TODO

  int last_decoded_CTB_RS;  // TODO

  void allocate_thread_contexts(int n);

  thread_context* get_thread_context(int n) {

    assert(n < nThreadContexts);

    return &thread_contexts[n];

  }

  int num_thread_contexts() const { return nThreadContexts; }

private:

  thread_context* thread_contexts; /* NOTE: cannot use std::vector, because thread_context has

                                      no copy constructor. */

  int nThreadContexts;

public:

  decoder_context* ctx;

  slice_unit(const slice_unit&) = delete;

  slice_unit& operator=(const slice_unit&) = delete;

};

class image_unit

{

public:

  image_unit();

  ~image_unit();

  de265_image* img = nullptr;

  de265_image  sao_output; // if SAO is used, this is allocated and used as SAO output buffer

  std::vector<slice_unit*> slice_units;

  std::vector<sei_message> suffix_SEIs;

  slice_unit* get_next_unprocessed_slice_segment() const {

    for (size_t i=0;i<slice_units.size();i++) {

      if (slice_units[i]->state == slice_unit::Unprocessed) {

        return slice_units[i];

      }

    }

    return nullptr;

  }

  slice_unit* get_prev_slice_segment(slice_unit* s) const {

    for (size_t i=1; i<slice_units.size(); i++) {

      if (slice_units[i]==s) {

        return slice_units[i-1];

      }

    }

    return nullptr;

  }

  slice_unit* get_next_slice_segment(slice_unit* s) const {

    for (size_t i=0; i<slice_units.size()-1; i++) {

      if (slice_units[i]==s) {

        return slice_units[i+1];

      }

    }

    return nullptr;

  }

  void dump_slices() const {

    for (size_t i=0; i<slice_units.size(); i++) {

      printf("[%zu] = %p\n",i,static_cast<void*>(slice_units[i]));

    }

  }

  bool all_slice_segments_processed() const {

    if (slice_units.size()==0) return true;

    if (slice_units.back()->state != slice_unit::Unprocessed) return true;

    return false;

  }

  bool is_first_slice_segment(const slice_unit* s) const {

    if (slice_units.size()==0) return false;

    return (slice_units[0] == s);

  }

  enum { Invalid, // headers not read yet

         Unknown, // SPS/PPS available

         Reference, // will be used as reference

         Leaf       // not a reference picture

  } role = Invalid;

  enum { Unprocessed,

         InProgress,

         Decoded,

         Dropped         // will not be decoded

  } state = Unprocessed;

  std::vector<thread_task*> tasks; // we are the owner

  /* Saved context models for WPP.

     There is one saved model for the initialization of each CTB row.

     The array is unused for non-WPP streams.

     Threading: context_model_table is a reference-counted handle (model

     pointer plus refcount pointer) that is not thread-safe, so each slot is

     touched by exactly one producer and one consumer and never concurrently.

     The row task of row N stores into ctx_models[N] after decoding CTB x=1

     and only then signals CTB_PROGRESS_PREFILTER for that CTB; the row task

     of row N+1 waits for that progress before it copies and releases the

     slot. The acquire/release ordering of de265_progress_lock makes the

     store visible to the consumer, so no lock is needed. This relies on the

     wait never being skipped because of stale progress: slice segments must

     arrive in increasing address order (slice_segment_order_is_valid()) and

     decode_slice_unit_WPP() resets the progress of the rows it schedules

     (GHSA-xp3h-6f5r-8cxp). */

  std::vector<context_model_table> ctx_models;  // TODO: move this into image ?

  /* Saved StatCoeff[] (persistent_rice_adaptation state) parallel to ctx_models.

     Per HEVC RExt, this state must be carried across WPP CTB rows together

     with the CABAC context. */

  std::vector<std::array<uint8_t, 4>> StatCoeff_models;

};

class base_context : public error_queue

{

 public:

  base_context();

  virtual ~base_context() { }

  // --- accelerated DSP functions ---

  void set_acceleration_functions(enum de265_acceleration);

  struct acceleration_functions acceleration; // CPU optimized functions

  //virtual /* */ de265_image* get_image(uint16_t dpb_index)       { return dpb.get_image(dpb_index); }

  virtual const de265_image* get_image(uint16_t frame_id) const = 0;

  virtual bool has_image(uint16_t frame_id) const = 0;

};

class decoder_context : public base_context {

 public:

  decoder_context();

  ~decoder_context();

  de265_error start_thread_pool(int nThreads);

  void        stop_thread_pool();

  void reset();

  bool has_sps(int id) const { return sps[id] != nullptr; }

  bool has_pps(int id) const { return pps[id] != nullptr; }

  std::shared_ptr<const seq_parameter_set> get_shared_sps(int id) { return sps[id]; }

  std::shared_ptr<const pic_parameter_set> get_shared_pps(int id) { return pps[id]; }

  /* */ seq_parameter_set* get_sps(int id)       { return sps[id].get(); }

  const seq_parameter_set* get_sps(int id) const { return sps[id].get(); }

  /* */ pic_parameter_set* get_pps(int id)       { return pps[id].get(); }

  const pic_parameter_set* get_pps(int id) const { return pps[id].get(); }

  /*

  const slice_segment_header* get_SliceHeader_atCtb(int ctb) {

    return img->slices[img->get_SliceHeaderIndex_atIndex(ctb)];

  }

  */

  uint8_t get_nal_unit_type() const { return nal_unit_type; }

  bool    get_RapPicFlag() const { return RapPicFlag; }

  de265_error decode_NAL(std::unique_ptr<NAL_unit> nal);

  de265_error decode(int* more);

  de265_error decode_some(bool* did_work);

  de265_error decode_slice_unit_sequential(image_unit* imgunit, slice_unit* sliceunit);

  de265_error decode_slice_unit_parallel(image_unit* imgunit, slice_unit* sliceunit);

  de265_error decode_slice_unit_WPP(image_unit* imgunit, slice_unit* sliceunit);

  de265_error decode_slice_unit_tiles(image_unit* imgunit, slice_unit* sliceunit);

  void process_nal_hdr(nal_header*);

  bool process_slice_segment_header(slice_segment_header*,

                                    de265_error*, de265_PTS pts,

                                    nal_header* nal_hdr, void* user_data);

  //void push_current_picture_to_output_queue();

  de265_error push_picture_to_output_queue(image_unit*);

  // --- parameters ---

  bool param_sei_check_hash = false;

  bool param_conceal_stream_errors = true;

  bool param_suppress_faulty_pictures = false;

  int  param_sps_headers_fd = -1;

  int  param_vps_headers_fd = -1;

  int  param_pps_headers_fd = -1;

  int  param_slice_headers_fd = -1;

  bool param_disable_deblocking = false;

  bool param_disable_sao = false;

  //bool param_disable_mc_residual_idct;  // not implemented yet

  //bool param_disable_intra_residual_idct;  // not implemented yet

  de265_security_limits param_security_limits = {

    1,                // version

    8192 * 8192,      // max_image_size_pixels

    16 * 1024 * 1024, // max_NAL_size_bytes

    256               // max_SEI_messages

  };

  void set_image_allocation_functions(de265_image_allocation* allocfunc, void* userdata);

  de265_image_allocation param_image_allocation_functions; // initialized in constructor

  void*                  param_image_allocation_userdata = nullptr;

  // --- input stream data ---

  NAL_Parser nal_parser;

  int get_num_worker_threads() const { return num_worker_threads; }

  /* */ de265_image* get_image(uint16_t dpb_index)       { return dpb.get_image(dpb_index); }

  const de265_image* get_image(uint16_t dpb_index) const override { return dpb.get_image(dpb_index); }

  bool has_image(uint16_t dpb_index) const override { return dpb_index<dpb.size(); }

  de265_image* get_next_picture_in_output_queue() { return dpb.get_next_picture_in_output_queue(); }

  int          num_pictures_in_output_queue() const { return dpb.num_pictures_in_output_queue(); }

  void         pop_next_picture_in_output_queue() { dpb.pop_next_picture_in_output_queue(); }

 private:

  de265_error read_vps_NAL(bitreader&);

  de265_error read_sps_NAL(bitreader&);

  de265_error read_pps_NAL(bitreader&);

  de265_error read_sei_NAL(bitreader& reader, bool suffix);

  de265_error read_eos_NAL(bitreader& reader);

  de265_error read_slice_NAL(bitreader&, std::unique_ptr<NAL_unit> nal, nal_header& nal_hdr);

 private:

  // --- internal data ---

  std::shared_ptr<video_parameter_set>  vps[ DE265_MAX_VPS_SETS ];

  std::shared_ptr<seq_parameter_set>    sps[ DE265_MAX_SPS_SETS ];

  std::shared_ptr<pic_parameter_set>    pps[ DE265_MAX_PPS_SETS ];

  std::shared_ptr<video_parameter_set>  current_vps;

  std::shared_ptr<seq_parameter_set>    current_sps;

  std::shared_ptr<pic_parameter_set>    current_pps;

 public:

  thread_pool thread_pool_;

 private:

  int num_worker_threads = 0;

 public:

  // --- frame dropping ---

  void set_limit_TID(int tid);

  int  get_highest_TID() const;

  int  get_current_TID() const { return current_HighestTid; }

  int  change_framerate(int more_vs_less); // 1: more, -1: less

  void set_framerate_ratio(int percent);

 private:

  // input parameters

  int limit_HighestTid = 6;    // never switch to a layer above this one

  int framerate_ratio = 100;

  // current control parameters

  int goal_HighestTid = 6;     // this is the layer we want to decode at

  int layer_framerate_ratio = 100; // ratio of frames to keep in the current layer

  int current_HighestTid = 6;  // the layer which we are currently decoding

  struct {

    int8_t tid;

    int8_t ratio;

  } framedrop_tab[100+1];

  int framedrop_tid_index[6+1];

  void compute_framedrop_table();

  void calc_tid_and_framerate_ratio();

 private:

  // --- decoded picture buffer ---

  decoded_picture_buffer dpb;

  int current_image_poc_lsb = -1;

  bool first_decoded_picture = true;

  bool NoRaslOutputFlag = false;

  bool HandleCraAsBlaFlag = false;

  bool FirstAfterEndOfSequenceNAL = false;

  int  PicOrderCntMsb = 0;

  int prevPicOrderCntLsb = 0;  // at precTid0Pic

  int prevPicOrderCntMsb = 0;  // at precTid0Pic

  de265_image* img = nullptr;

 public:

  const slice_segment_header* previous_slice_header = nullptr; /* Remember the last slice for a successive

                  dependent slice. */

  // --- motion compensation ---

 public:

  int PocLsbLt[MAX_NUM_REF_PICS]{};

  int UsedByCurrPicLt[MAX_NUM_REF_PICS]{};

  uint32_t DeltaPocMsbCycleLt[MAX_NUM_REF_PICS]{};

 private:

  int CurrDeltaPocMsbPresentFlag[MAX_NUM_REF_PICS]{};

  int FollDeltaPocMsbPresentFlag[MAX_NUM_REF_PICS]{};

  // The number of entries in the lists below.

  int NumPocStCurrBefore = 0;

  int NumPocStCurrAfter = 0;

  int NumPocStFoll = 0;

  int NumPocLtCurr = 0;

  int NumPocLtFoll = 0;

  // These lists contain absolute POC values.

  int PocStCurrBefore[MAX_NUM_REF_PICS]{}; // used for reference in current picture, smaller POC

  int PocStCurrAfter[MAX_NUM_REF_PICS]{};  // used for reference in current picture, larger POC

  int PocStFoll[MAX_NUM_REF_PICS]{}; // not used for reference in current picture, but in future picture

  int PocLtCurr[MAX_NUM_REF_PICS]{}; // used in current picture

  int PocLtFoll[MAX_NUM_REF_PICS]{}; // used in some future picture

  // These lists contain indices into the DPB.

  int RefPicSetStCurrBefore[MAX_NUM_REF_PICS]{};

  int RefPicSetStCurrAfter[MAX_NUM_REF_PICS]{};

  int RefPicSetStFoll[MAX_NUM_REF_PICS]{};

  int RefPicSetLtCurr[MAX_NUM_REF_PICS]{};

  int RefPicSetLtFoll[MAX_NUM_REF_PICS]{};

  // --- parameters derived from parameter sets ---

  // NAL

  uint8_t nal_unit_type = 0;

  bool IdrPicFlag = false;

  bool RapPicFlag = false;

  // --- image unit queue ---

  std::vector<image_unit*> image_units;

  bool flush_reorder_buffer_at_this_frame = false;

 private:

  void init_thread_context(thread_context* tctx);

  void add_task_decode_CTB_row(thread_context* tctx, bool firstSliceSubstream, uint16_t ctbRow);

  void add_task_decode_slice_segment(thread_context* tctx, bool firstSliceSubstream,

                                     uint16_t ctbX, uint16_t ctbY);

  /* Check that the slice segment 'shdr' may be appended to 'imgunit', i.e. that

     its first CTB follows the previous slice segment of the picture in tile-scan

     order (H.265 7.4.2.4.5). Adds a warning and marks the image as faulty when

     the slice segment has to be dropped. */

  bool slice_segment_order_is_valid(image_unit* imgunit, const slice_segment_header* shdr);

  void mark_whole_slice_as_processed(image_unit* imgunit,

                                     slice_unit* sliceunit,

                                     int progress);

  void process_picture_order_count(slice_segment_header* hdr);

  /*

  If there is no space for a new image, returns the negative value of an de265_error.

  I.e. you can check for error by return_value<0, which is error (-return_value);

   */

  int generate_unavailable_reference_picture(const seq_parameter_set* sps,

                                             int POC, bool longTerm);

  de265_error process_reference_picture_set(slice_segment_header* hdr);

  bool construct_reference_picture_lists(slice_segment_header* hdr);

  void remove_images_from_dpb(const std::vector<int>& removeImageList);

  void run_postprocessing_filters_sequential(struct de265_image* img);

  void run_postprocessing_filters_parallel(image_unit* img);

};

#endif