// Copyright 2014 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// This file contains an implementation of an H264 Annex-B video stream parser.

#ifndef MEDIA_PARSERS_H264_PARSER_H_
#define MEDIA_PARSERS_H264_PARSER_H_

#include <stddef.h>
#include <stdint.h>
#include <sys/types.h>

#include <array>
#include <map>
#include <memory>
#include <optional>
#include <utility>
#include <variant>
#include <vector>

#include "base/containers/span.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/raw_span.h"
#include "media/base/media_export.h"
#include "media/base/ranges.h"
#include "media/base/video_codecs.h"
#include "media/base/video_color_space.h"
#include "media/base/video_types.h"
#include "media/parsers/h264_bit_reader.h"

namespace gfx {
class Rect;
class Size;
struct HdrMetadataSmpteSt2086;
struct HdrMetadataCta861_3;
}  // namespace gfx

namespace media {

struct SubsampleEntry;

// For explanations of each struct and its members, see H.264 specification
// at http://www.itu.int/rec/T-REC-H.264.
struct MEDIA_EXPORT H264NALU {
  H264NALU();

  enum Type {
    kUnspecified = 0,
    kNonIDRSlice = 1,
    kSliceDataA = 2,
    kSliceDataB = 3,
    kSliceDataC = 4,
    kIDRSlice = 5,
    kSEIMessage = 6,
    kSPS = 7,
    kPPS = 8,
    kAUD = 9,
    kEOSeq = 10,
    kEOStream = 11,
    kFiller = 12,
    kSPSExt = 13,
    kPrefix = 14,
    kSubsetSPS = 15,
    kDPS = 16,
    kReserved17 = 17,
    kReserved18 = 18,
    kCodedSliceAux = 19,
    kCodedSliceExtension = 20,
  };

  // After (without) start code; we don't own the underlying memory
  // and a shallow copy should be made when copying this struct.
  base::raw_span<const uint8_t, DanglingUntriaged> data;

  int nal_ref_idc = 0;
  int nal_unit_type = 0;
};

enum {
  kH264ScalingList4x4Length = 16,
  kH264ScalingList8x8Length = 64,
};

struct MEDIA_EXPORT H264SPS {
  H264SPS();
  H264SPS(const H264SPS&);
  H264SPS& operator=(const H264SPS&);
  H264SPS(H264SPS&&) noexcept;
  H264SPS& operator=(H264SPS&&) noexcept;

  bool operator==(const H264SPS&) const;

  enum H264ProfileIDC {
    kProfileIDCBaseline = 66,
    kProfileIDCConstrainedBaseline = kProfileIDCBaseline,
    kProfileIDCMain = 77,
    kProfileIDScalableBaseline = 83,
    kProfileIDScalableHigh = 86,
    kProfileIDCHigh = 100,
    kProfileIDHigh10 = 110,
    kProfileIDSMultiviewHigh = 118,
    kProfileIDHigh422 = 122,
    kProfileIDStereoHigh = 128,
    kProfileIDHigh444Predictive = 244,
  };

  enum H264LevelIDC : uint8_t {
    kLevelIDC1p0 = 10,
    kLevelIDC1B = 9,
    kLevelIDC1p1 = 11,
    kLevelIDC1p2 = 12,
    kLevelIDC1p3 = 13,
    kLevelIDC2p0 = 20,
    kLevelIDC2p1 = 21,
    kLevelIDC2p2 = 22,
    kLevelIDC3p0 = 30,
    kLevelIDC3p1 = 31,
    kLevelIDC3p2 = 32,
    kLevelIDC4p0 = 40,
    kLevelIDC4p1 = 41,
    kLevelIDC4p2 = 42,
    kLevelIDC5p0 = 50,
    kLevelIDC5p1 = 51,
    kLevelIDC5p2 = 52,
    kLevelIDC6p0 = 60,
    kLevelIDC6p1 = 61,
    kLevelIDC6p2 = 62,
  };

  enum AspectRatioIdc {
    kExtendedSar = 255,
  };

  enum {
    // Constants for HRD parameters (spec ch. E.2.2).
    kBitRateScaleConstantTerm = 6,  // Equation E-37.
    kCPBSizeScaleConstantTerm = 4,  // Equation E-38.
    kDefaultInitialCPBRemovalDelayLength = 24,
    kDefaultDPBOutputDelayLength = 24,
    kDefaultTimeOffsetLength = 24,
  };

  int profile_idc = 0;
  bool constraint_set0_flag = false;
  bool constraint_set1_flag = false;
  bool constraint_set2_flag = false;
  bool constraint_set3_flag = false;
  bool constraint_set4_flag = false;
  bool constraint_set5_flag = false;
  int level_idc = 0;
  int seq_parameter_set_id = 0;

  int chroma_format_idc = 0;
  bool separate_colour_plane_flag = false;
  int bit_depth_luma_minus8 = 0;
  int bit_depth_chroma_minus8 = 0;
  bool qpprime_y_zero_transform_bypass_flag = false;

  bool seq_scaling_matrix_present_flag = false;
  std::array<std::array<uint8_t, kH264ScalingList4x4Length>, 6>
      scaling_list4x4 = {};
  std::array<std::array<uint8_t, kH264ScalingList8x8Length>, 6>
      scaling_list8x8 = {};

  int log2_max_frame_num_minus4 = 0;
  int pic_order_cnt_type = 0;
  int log2_max_pic_order_cnt_lsb_minus4 = 0;
  bool delta_pic_order_always_zero_flag = false;
  int offset_for_non_ref_pic = 0;
  int offset_for_top_to_bottom_field = 0;
  int num_ref_frames_in_pic_order_cnt_cycle = 0;
  int expected_delta_per_pic_order_cnt_cycle = 0;  // calculated
  std::array<int, 255> offset_for_ref_frame = {};
  int max_num_ref_frames = 0;
  bool gaps_in_frame_num_value_allowed_flag = false;
  int pic_width_in_mbs_minus1 = 0;
  int pic_height_in_map_units_minus1 = 0;
  bool frame_mbs_only_flag = false;
  bool mb_adaptive_frame_field_flag = false;
  bool direct_8x8_inference_flag = false;
  bool frame_cropping_flag = false;
  int frame_crop_left_offset = 0;
  int frame_crop_right_offset = 0;
  int frame_crop_top_offset = 0;
  int frame_crop_bottom_offset = 0;

  bool vui_parameters_present_flag = false;
  int sar_width = 0;   // Set to 0 when not specified.
  int sar_height = 0;  // Set to 0 when not specified.
  bool bitstream_restriction_flag = false;
  int max_num_reorder_frames = 0;
  int max_dec_frame_buffering = 0;
  bool timing_info_present_flag = false;
  int num_units_in_tick = 0;
  int time_scale = 0;
  bool fixed_frame_rate_flag = false;

  bool video_signal_type_present_flag = false;
  int video_format = 0;
  bool video_full_range_flag = false;
  bool colour_description_present_flag = false;
  int colour_primaries = 0;
  int transfer_characteristics = 0;
  int matrix_coefficients = 0;

  // TODO(posciak): actually parse these instead of ParseAndIgnoreHRDParameters.
  bool nal_hrd_parameters_present_flag = false;
  int cpb_cnt_minus1 = 0;
  int bit_rate_scale = 0;
  int cpb_size_scale = 0;
  std::array<int, 32> bit_rate_value_minus1 = {};
  std::array<int, 32> cpb_size_value_minus1 = {};
  std::array<bool, 32> cbr_flag = {};
  int initial_cpb_removal_delay_length_minus_1 = 0;
  int cpb_removal_delay_length_minus1 = 0;
  int dpb_output_delay_length_minus1 = 0;
  int time_offset_length = 0;

  bool low_delay_hrd_flag = false;

  int chroma_array_type = 0;

  // Get corresponding SPS |level_idc| and |constraint_set3_flag| value from
  // requested |profile| and |level| (see Spec A.3.1).
  static void GetLevelConfigFromProfileLevel(VideoCodecProfile profile,
                                             uint8_t level,
                                             int* level_idc,
                                             bool* constraint_set3_flag);

  // Helpers to compute frequently-used values. These methods return
  // std::nullopt if they encounter integer overflow. They do not verify that
  // the results are in-spec for the given profile or level.
  std::optional<gfx::Size> GetCodedSize() const;
  std::optional<gfx::Rect> GetVisibleRect() const;
  VideoColorSpace GetColorSpace() const;
  VideoChromaSampling GetChromaSampling() const;

  // Helper to compute indicated level from parsed SPS data. The value of
  // indicated level would be included in H264LevelIDC enum representing the
  // level as in name.
  uint8_t GetIndicatedLevel() const;
  // Helper to check if indicated level is lower than or equal to
  // |target_level|.
  bool CheckIndicatedLevelWithinTarget(uint8_t target_level) const;
};

struct MEDIA_EXPORT H264PPS {
  H264PPS();
  H264PPS(const H264PPS&);
  H264PPS& operator=(const H264PPS&);
  H264PPS(H264PPS&&) noexcept;
  H264PPS& operator=(H264PPS&&) noexcept;

  bool operator==(const H264PPS&) const;

  int pic_parameter_set_id = 0;
  int seq_parameter_set_id = 0;
  bool entropy_coding_mode_flag = false;
  bool bottom_field_pic_order_in_frame_present_flag = false;
  int num_slice_groups_minus1 = 0;
  // TODO(posciak): Slice groups not implemented, could be added at some point.
  int num_ref_idx_l0_default_active_minus1 = 0;
  int num_ref_idx_l1_default_active_minus1 = 0;
  bool weighted_pred_flag = false;
  int weighted_bipred_idc = 0;
  int pic_init_qp_minus26 = 0;
  int pic_init_qs_minus26 = 0;
  int chroma_qp_index_offset = 0;
  bool deblocking_filter_control_present_flag = false;
  bool constrained_intra_pred_flag = false;
  bool redundant_pic_cnt_present_flag = false;
  bool transform_8x8_mode_flag = false;

  bool pic_scaling_matrix_present_flag = false;
  std::array<std::array<uint8_t, kH264ScalingList4x4Length>, 6>
      scaling_list4x4 = {};
  std::array<std::array<uint8_t, kH264ScalingList8x8Length>, 6>
      scaling_list8x8 = {};

  int second_chroma_qp_index_offset = 0;
};

struct MEDIA_EXPORT H264ModificationOfPicNum {
  int modification_of_pic_nums_idc = 0;
  union {
    int abs_diff_pic_num_minus1 = 0;
    int long_term_pic_num;
  };
};

struct MEDIA_EXPORT H264WeightingFactors {
  bool luma_weight_flag = false;
  bool chroma_weight_flag = false;
  std::array<int, 32> luma_weight = {};
  std::array<int, 32> luma_offset = {};
  std::array<std::array<int, 2>, 32> chroma_weight = {};
  std::array<std::array<int, 2>, 32> chroma_offset = {};
};

struct MEDIA_EXPORT H264DecRefPicMarking {
  int memory_mgmnt_control_operation = 0;
  int difference_of_pic_nums_minus1 = 0;
  int long_term_pic_num = 0;
  int long_term_frame_idx = 0;
  int max_long_term_frame_idx_plus1 = 0;
};

struct MEDIA_EXPORT H264SliceHeader {
  H264SliceHeader();
  H264SliceHeader(const H264SliceHeader&);
  H264SliceHeader& operator=(const H264SliceHeader&);
  H264SliceHeader(H264SliceHeader&&);
  H264SliceHeader& operator=(H264SliceHeader&&);

  enum { kRefListSize = 32, kRefListModSize = kRefListSize };

  enum Type {
    kPSlice = 0,
    kBSlice = 1,
    kISlice = 2,
    kSPSlice = 3,
    kSISlice = 4,
  };

  bool IsPSlice() const;
  bool IsBSlice() const;
  bool IsISlice() const;
  bool IsSPSlice() const;
  bool IsSISlice() const;

  bool idr_pic_flag = false;  // from NAL header
  int nal_ref_idc = 0;        // from NAL header
  raw_ptr<const uint8_t, AllowPtrArithmetic | DanglingUntriaged>
      nalu_data;             // from NAL header
  off_t nalu_size = 0;       // from NAL header
  off_t header_bit_size = 0;  // calculated

  int first_mb_in_slice = 0;
  int slice_type = 0;
  int pic_parameter_set_id = 0;
  int colour_plane_id = 0;  // TODO(posciak): use this!  http://crbug.com/139878
  int frame_num = 0;
  bool field_pic_flag = false;
  bool bottom_field_flag = false;
  int idr_pic_id = 0;
  int pic_order_cnt_lsb = 0;
  int delta_pic_order_cnt_bottom = 0;
  int delta_pic_order_cnt0 = 0;
  int delta_pic_order_cnt1 = 0;
  int redundant_pic_cnt = 0;
  bool direct_spatial_mv_pred_flag = false;

  bool num_ref_idx_active_override_flag = false;
  int num_ref_idx_l0_active_minus1 = 0;
  int num_ref_idx_l1_active_minus1 = 0;
  bool ref_pic_list_modification_flag_l0 = false;
  bool ref_pic_list_modification_flag_l1 = false;
  std::array<H264ModificationOfPicNum, kRefListModSize>
      ref_list_l0_modifications;
  std::array<H264ModificationOfPicNum, kRefListModSize>
      ref_list_l1_modifications;

  int luma_log2_weight_denom = 0;
  int chroma_log2_weight_denom = 0;

  bool luma_weight_l0_flag = false;
  bool chroma_weight_l0_flag = false;
  H264WeightingFactors pred_weight_table_l0;

  bool luma_weight_l1_flag = false;
  bool chroma_weight_l1_flag = false;
  H264WeightingFactors pred_weight_table_l1;

  bool no_output_of_prior_pics_flag = false;
  bool long_term_reference_flag = false;

  bool adaptive_ref_pic_marking_mode_flag = false;
  std::array<H264DecRefPicMarking, kRefListSize> ref_pic_marking;

  int cabac_init_idc = 0;
  int slice_qp_delta = 0;
  bool sp_for_switch_flag = false;
  int slice_qs_delta = 0;
  int disable_deblocking_filter_idc = 0;
  int slice_alpha_c0_offset_div2 = 0;
  int slice_beta_offset_div2 = 0;

  // Calculated.
  // Size in bits of dec_ref_pic_marking() syntax element.
  size_t dec_ref_pic_marking_bit_size = 0;
  size_t pic_order_cnt_bit_size = 0;

  // This is when we are using full sample encryption and only the portions
  // needed for DPB management are filled in, the rest will already be known
  // by the accelerator and we will not need to specify it.
  bool full_sample_encryption = false;
  // This is used by some accelerators to handle decoding after slice header
  // parsing.
  uint32_t full_sample_index = 0;
};

struct MEDIA_EXPORT H264SEIRecoveryPoint {
  int recovery_frame_cnt = 0;
  bool exact_match_flag = false;
  bool broken_link_flag = false;
  int changing_slice_group_idc = 0;
};

struct MEDIA_EXPORT H264SEIMasteringDisplayInfo {
  enum {
    kNumDisplayPrimaries = 3,
    kDisplayPrimaryComponents = 2,
  };

  std::array<std::array<uint16_t, kDisplayPrimaryComponents>,
             kNumDisplayPrimaries>
      display_primaries = {};
  std::array<uint16_t, 2> white_points = {};
  uint32_t max_luminance = 0;
  uint32_t min_luminance = 0;

  gfx::HdrMetadataSmpteSt2086 ToGfx() const;
};

struct MEDIA_EXPORT H264SEIContentLightLevelInfo {
  uint16_t max_content_light_level = 0;
  uint16_t max_picture_average_light_level = 0;

  gfx::HdrMetadataCta861_3 ToGfx() const;
};

using H264SEIMessage = std::variant<std::monostate,
                                    H264SEIRecoveryPoint,
                                    H264SEIMasteringDisplayInfo,
                                    H264SEIContentLightLevelInfo>;

struct MEDIA_EXPORT H264SEI {
  H264SEI();
  ~H264SEI();

  std::vector<H264SEIMessage> msgs;
};

// Class to parse an Annex-B H.264 stream,
// as specified in chapters 7 and Annex B of the H.264 spec.
class MEDIA_EXPORT H264Parser {
 public:
  enum Result {
    kOk,
    kInvalidStream,      // error in stream
    kUnsupportedStream,  // stream not supported by the parser
    kEOStream,           // end of stream
  };

  // Find offset from start of data to next NALU start code
  // and size of found start code (3 or 4 bytes).
  // If no start code is found, offset is pointing to the first unprocessed byte
  // (i.e. the first byte that was not considered as a possible start of a start
  // code) and |*start_code_size| is set to 0.
  // Preconditions:
  // - |data_size| >= 0
  // Postconditions:
  // - |*offset| is between 0 and |data_size| included.
  //   It is strictly less than |data_size| if |data_size| > 0.
  // - |*start_code_size| is either 0, 3 or 4.
  static bool FindStartCode(base::span<const uint8_t> data,
                            size_t* offset,
                            size_t* start_code_size);

  // Wrapper for FindStartCode() that skips over start codes that
  // may appear inside of |encrypted_ranges_|.
  // Returns true if a start code was found. Otherwise returns false.
  static bool FindStartCodeInClearRanges(base::span<const uint8_t> data,
                                         const Ranges<const uint8_t*>& ranges,
                                         size_t* offset,
                                         size_t* start_code_size);

  // DEPRECATED: Use the above `base::span` variant to avoid unsafe buffer
  // usage.
  // TODO(https://crbug.com/40284755): Remove this once the callers are gone.
  static bool FindStartCodeInClearRanges(const uint8_t* data,
                                         off_t data_size,
                                         const Ranges<const uint8_t*>& ranges,
                                         off_t* offset,
                                         off_t* start_code_size);

  static VideoCodecProfile ProfileIDCToVideoCodecProfile(int profile_idc);

  // Parses the input stream and returns all the NALUs through |nalus|. Returns
  // false if the stream is invalid.
  static bool ParseNALUs(const uint8_t* stream,
                         size_t stream_size,
                         std::vector<H264NALU>* nalus);

  H264Parser();

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

  ~H264Parser();

  void Reset();
  // Set current stream pointer to |stream| of |stream_size| in bytes,
  // |stream| owned by caller.
  // |subsamples| contains information about what parts of |stream| are
  // encrypted.
  void SetStream(base::span<const uint8_t> stream);

  // DEPRECATED: Use the above `base::span` variant to avoid unsafe buffer
  // usage.
  // TODO(https://crbug.com/40284755): Remove this once the callers are gone.
  void SetStream(const uint8_t* stream, off_t stream_size);

  void SetEncryptedStream(base::span<const uint8_t> stream,
                          const std::vector<SubsampleEntry>& subsamples);

  // Read the stream to find the next NALU, identify it and return
  // that information in |*nalu|. This advances the stream to the beginning
  // of this NALU, but not past it, so subsequent calls to NALU-specific
  // parsing functions (ParseSPS, etc.)  will parse this NALU.
  // If the caller wishes to skip the current NALU, it can call this function
  // again, instead of any NALU-type specific parse functions below.
  Result AdvanceToNextNALU(H264NALU* nalu);

  // NALU-specific parsing functions.
  // These should be called after AdvanceToNextNALU().

  // SPSes and PPSes are owned by the parser class and the memory for their
  // structures is managed here, not by the caller, as they are reused
  // across NALUs.
  //
  // Parse an SPS/PPS NALU and save their data in the parser, returning id
  // of the parsed structure in |*pps_id|/|*sps_id|.
  // To get a pointer to a given SPS/PPS structure, use GetSPS()/GetPPS(),
  // passing the returned |*sps_id|/|*pps_id| as parameter.
  // TODO(posciak,fischman): consider replacing returning Result from Parse*()
  // methods with a scoped_ptr and adding an AtEOS() function to check for EOS
  // if Parse*() return NULL.
  Result ParseSPS(int* sps_id);
  Result ParsePPS(int* pps_id);

  // Parses the SPS ID from the SPSExt, but otherwise does nothing.
  Result ParseSPSExt(int* sps_id);

  // Return a pointer to SPS/PPS with given |sps_id|/|pps_id| or NULL if not
  // present.
  const H264SPS* GetSPS(int sps_id) const;
  const H264PPS* GetPPS(int pps_id) const;

  // Slice headers and SEI messages are not used across NALUs by the parser
  // and can be discarded after current NALU, so the parser does not store
  // them, nor does it manage their memory.
  // The caller has to provide and manage it instead.

  // Parse a slice header, returning it in |*shdr|. |*nalu| must be set to
  // the NALU returned from AdvanceToNextNALU() and corresponding to |*shdr|.
  Result ParseSliceHeader(const H264NALU& nalu, H264SliceHeader* shdr);

  // Parse a SEI, returning it in |*sei|, provided and managed by the caller.
  Result ParseSEI(H264SEI* sei);

  // The return value of this method changes for every successful call to
  // AdvanceToNextNALU().
  // This returns the subsample information for the last NALU that was output
  // from AdvanceToNextNALU().
  std::vector<SubsampleEntry> GetCurrentSubsamples();

 private:
  // Move the stream pointer to the beginning of the next NALU,
  // i.e. pointing at the next start code.
  // Return true if a NALU has been found.
  // If a NALU is found:
  // - its size in bytes is returned in |*nalu_size| and includes
  //   the start code as well as the trailing zero bits.
  // - the size in bytes of the start code is returned in |*start_code_size|.
  bool LocateNALU(size_t* nalu_size, size_t* start_code_size);

  // Parse scaling lists (see spec).
  Result ParseScalingList(base::span<uint8_t> scaling_list, bool* use_default);
  Result ParseSPSScalingLists(H264SPS* sps);
  Result ParsePPSScalingLists(const H264SPS& sps, H264PPS* pps);

  // Parse optional VUI parameters in SPS (see spec).
  Result ParseVUIParameters(H264SPS* sps);
  // Set |hrd_parameters_present| to true only if they are present.
  Result ParseAndIgnoreHRDParameters(bool* hrd_parameters_present);

  // Parse reference picture lists' modifications (see spec).
  Result ParseRefPicListModifications(H264SliceHeader* shdr);
  Result ParseRefPicListModification(int num_ref_idx_active_minus1,
                                     H264ModificationOfPicNum* ref_list_mods);

  // Parse prediction weight table (see spec).
  Result ParsePredWeightTable(const H264SPS& sps, H264SliceHeader* shdr);

  // Parse weighting factors (see spec).
  Result ParseWeightingFactors(int num_ref_idx_active_minus1,
                               int chroma_array_type,
                               int luma_log2_weight_denom,
                               int chroma_log2_weight_denom,
                               H264WeightingFactors* w_facts);

  // Parse decoded reference picture marking information (see spec).
  Result ParseDecRefPicMarking(H264SliceHeader* shdr);

  // Pointer to the current NALU in the stream.
  base::raw_span<const uint8_t, DanglingUntriaged> stream_;

  H264BitReader br_;

  // PPSes and SPSes stored for future reference.
  std::map<int, std::unique_ptr<H264SPS>> active_SPSes_;
  std::map<int, std::unique_ptr<H264PPS>> active_PPSes_;

  // Ranges of encrypted bytes in the buffer passed to
  // SetEncryptedStream().
  Ranges<const uint8_t*> encrypted_ranges_;

  // This contains the range of the previous NALU found in
  // AdvanceToNextNalu(). Holds exactly one range.
  Ranges<const uint8_t*> previous_nalu_range_;
};

}  // namespace media

#endif  // MEDIA_PARSERS_H264_PARSER_H_