/*

 * Copyright (c) 2025, Alliance for Open Media. All rights reserved

 *

 * This source code is subject to the terms of the BSD 3-Clause Clear License

 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear

 * License was not distributed with this source code in the LICENSE file, you

 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  If the

 * Alliance for Open Media Patent License 1.0 was not distributed with this

 * source code in the PATENTS file, you can obtain it at

 * aomedia.org/license/patent-license/.

 */

#include <assert.h>

#include "config/avm_config.h"

#include "config/avm_scale_rtcd.h"

#include "avm/avm_codec.h"

#include "avm_dsp/bitreader_buffer.h"

#include "avm_ports/mem_ops.h"

#include "av2/common/common.h"

#include "av2/common/obu_util.h"

#include "av2/common/timing.h"

#include "av2/decoder/decoder.h"

#include "av2/decoder/decodeframe.h"

#include "av2/decoder/obu.h"

void copy_fgm_from_list(AV2_COMMON *cm, avm_film_grain_t *pars,

                        const struct film_grain_model *fgm) {

  const SequenceHeader *const seq_params = &cm->seq_params;

  pars->bit_depth = seq_params->bit_depth;

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

    pars->fgm_points[c] = fgm->fgm_points[c];

  }

  for (int i = 0; i < 14; i++) {

    for (int j = 0; j < 2; j++) {

      pars->fgm_scaling_points_0[i][j] = fgm->fgm_scaling_points[0][i][j];

      pars->fgm_scaling_points_1[i][j] = fgm->fgm_scaling_points[1][i][j];

      pars->fgm_scaling_points_2[i][j] = fgm->fgm_scaling_points[2][i][j];

    }

  }

  pars->scaling_shift = fgm->scaling_shift;

  pars->ar_coeff_lag = fgm->ar_coeff_lag;

  // 8 bit values

  for (int i = 0; i < 24; i++) pars->ar_coeffs_y[i] = fgm->ar_coeffs_y[i];

  for (int i = 0; i < 25; i++) {

    pars->ar_coeffs_cb[i] = fgm->ar_coeffs_cb[i];

    pars->ar_coeffs_cr[i] = fgm->ar_coeffs_cr[i];

  }

  pars->ar_coeff_shift = fgm->ar_coeff_shift;

  pars->cb_mult = fgm->cb_mult;

  pars->cb_luma_mult = fgm->cb_luma_mult;

  pars->cb_offset = fgm->cb_offset;

  pars->cr_mult = fgm->cr_mult;

  pars->cr_luma_mult = fgm->cr_luma_mult;

  pars->cr_offset = fgm->cr_offset;

  pars->overlap_flag = fgm->overlap_flag;

  pars->clip_to_restricted_range = fgm->clip_to_restricted_range;

  pars->mc_identity = fgm->mc_identity;

  pars->fgm_scale_from_channel0_flag = fgm->fgm_scale_from_channel0_flag;

  pars->grain_scale_shift = fgm->grain_scale_shift;

  pars->block_size = fgm->block_size;

}

static void read_film_grain_model(struct film_grain_model *fgm, int chroma_idc,

                                  struct avm_read_bit_buffer *rb,

                                  struct avm_internal_error_info *error_info) {

  int monochrome = chroma_idc == CHROMA_FORMAT_400;

  int subsampling_x = chroma_idc == CHROMA_FORMAT_444 ? 0 : 1;

  int subsampling_y =

      (chroma_idc == CHROMA_FORMAT_422 || chroma_idc == CHROMA_FORMAT_444) ? 0

                                                                           : 1;

  // Scaling functions parameters

  int fgmNumChannels = monochrome ? 1 : 3;

  if (fgmNumChannels > 1) {

    fgm->fgm_scale_from_channel0_flag = avm_rb_read_bit(rb);

  } else {

    fgm->fgm_scale_from_channel0_flag = 0;

  }

  int fgmNumScalingChannels =

      fgm->fgm_scale_from_channel0_flag ? 1 : fgmNumChannels;

  for (int c = 0; c < fgmNumScalingChannels; c++) {

    fgm->fgm_points[c] = avm_rb_read_literal(rb, 4);  // max 14

    if (fgm->fgm_points[c] > 14) {

      avm_internal_error(error_info, AVM_CODEC_UNSUP_BITSTREAM,

                         "Number of points for film grain %s scaling "

                         "function exceeds the maximum value.",

                         c == 0   ? "y"

                         : c == 1 ? "cb"

                                  : "cr");

    }

    if (fgm->fgm_points[c]) {

      int point_value_increment_bits_minus1 = avm_rb_read_literal(rb, 3);

      int point_scaling_bits_minus5 = avm_rb_read_literal(rb, 2);

      int bitsIncr = point_value_increment_bits_minus1 + 1;

      int bitsScal = point_scaling_bits_minus5 + 5;

      for (int i = 0; i < fgm->fgm_points[c]; i++) {

        int fgm_value_increment = avm_rb_read_literal(rb, bitsIncr);

        if (i == 0)

          fgm->fgm_scaling_points[c][i][0] = fgm_value_increment;

        else

          fgm->fgm_scaling_points[c][i][0] =

              fgm->fgm_scaling_points[c][i - 1][0] + fgm_value_increment;

        if (i && fgm->fgm_scaling_points[c][i - 1][0] >=

                     fgm->fgm_scaling_points[c][i][0]) {

          avm_internal_error(error_info, AVM_CODEC_UNSUP_BITSTREAM,

                             "First coordinate of the %s scaling function "

                             "points shall be increasing.",

                             c == 0   ? "y"

                             : c == 1 ? "cb"

                                      : "cr");

        }

        fgm->fgm_scaling_points[c][i][1] = avm_rb_read_literal(rb, bitsScal);

      }

    }

  }

  // Initialize unsignaled values to zero

  for (int c = fgmNumScalingChannels; c < 3; c++) fgm->fgm_points[c] = 0;

  if ((subsampling_x == 1) && (subsampling_y == 1) &&

      (((fgm->fgm_points[1] == 0) && (fgm->fgm_points[2] != 0)) ||

       ((fgm->fgm_points[1] != 0) && (fgm->fgm_points[2] == 0)))) {

    avm_internal_error(error_info, AVM_CODEC_UNSUP_BITSTREAM,

                       "In YCbCr 4:2:0, film grain shall be applied "

                       "to both chroma components or neither.");

  }

  fgm->scaling_shift = avm_rb_read_literal(rb, 2) + 8;  // 8 + value

  // AR coefficients

  // Only sent if the corresponsing scaling function has

  // more than 0 points

  fgm->ar_coeff_lag = avm_rb_read_literal(rb, 2);

  int num_pos_luma = 2 * fgm->ar_coeff_lag * (fgm->ar_coeff_lag + 1);

  int num_pos_chroma = num_pos_luma;

  if (fgm->fgm_points[0]) {

    ++num_pos_chroma;

    int bits_per_ar_coeff_y_minus5 = avm_rb_read_literal(rb, 2);

    int BitsArY = bits_per_ar_coeff_y_minus5 + 5;

    int midPointY = 1 << (BitsArY - 1);

    for (int i = 0; i < num_pos_luma; i++)

      fgm->ar_coeffs_y[i] = avm_rb_read_literal(rb, BitsArY) - midPointY;

  }

  if (fgm->fgm_points[1] || fgm->fgm_scale_from_channel0_flag) {

    int bits_per_ar_coeff_cb_minus5 = avm_rb_read_literal(rb, 2);

    int BitsArCb = bits_per_ar_coeff_cb_minus5 + 5;

    int midPointCb = 1 << (BitsArCb - 1);

    for (int i = 0; i < num_pos_chroma; i++)

      fgm->ar_coeffs_cb[i] = avm_rb_read_literal(rb, BitsArCb) - midPointCb;

  }

  if (fgm->fgm_points[2] || fgm->fgm_scale_from_channel0_flag) {

    int bits_per_ar_coeff_cr_minus5 = avm_rb_read_literal(rb, 2);

    int BitsArCr = bits_per_ar_coeff_cr_minus5 + 5;

    int midPointCr = 1 << (BitsArCr - 1);

    for (int i = 0; i < num_pos_chroma; i++)

      fgm->ar_coeffs_cr[i] = avm_rb_read_literal(rb, BitsArCr) - midPointCr;

  }

  fgm->ar_coeff_shift = avm_rb_read_literal(rb, 2) + 6;  // 6 + value

  fgm->grain_scale_shift = avm_rb_read_literal(rb, 2);

  if (fgm->fgm_points[1] > 0) {

    fgm->cb_mult = avm_rb_read_literal(rb, 8);

    fgm->cb_luma_mult = avm_rb_read_literal(rb, 8);

    fgm->cb_offset = avm_rb_read_literal(rb, 9);

  }

  if (fgm->fgm_points[2] > 0) {

    fgm->cr_mult = avm_rb_read_literal(rb, 8);

    fgm->cr_luma_mult = avm_rb_read_literal(rb, 8);

    fgm->cr_offset = avm_rb_read_literal(rb, 9);

  }

  fgm->overlap_flag = avm_rb_read_bit(rb);

  fgm->clip_to_restricted_range = avm_rb_read_bit(rb);

  if (fgm->clip_to_restricted_range)

    fgm->mc_identity = avm_rb_read_bit(rb);

  else

    fgm->mc_identity = 0;

  fgm->block_size = avm_rb_read_bit(rb);

}

// acc_fgm_id_bitmap is an in/out parameter. The caller should set

// *acc_fgm_id_bitmap to 0 before the first call to read_fgm_obu(). Each

// read_fgm_obu() call updates *acc_fgm_id_bitmap by bitwise-ORing the

// fgm_bit_map from the FGM OBU with *acc_fgm_id_bitmap.

uint32_t read_fgm_obu(AV2Decoder *pbi, const int obu_tlayer_id,

                      const int obu_mlayer_id, uint32_t *acc_fgm_id_bitmap,

                      struct avm_read_bit_buffer *rb) {

  const uint32_t saved_bit_offset = rb->bit_offset;

  int fgm_bit_map = avm_rb_read_literal(rb, MAX_FGM_NUM);

  if (*acc_fgm_id_bitmap & (uint32_t)fgm_bit_map) {

    avm_internal_error(

        &pbi->common.error, AVM_CODEC_INVALID_PARAM,

        "fgm_bit_map(%d) overlaps the accumulated fgm_bit_map(%d)", fgm_bit_map,

        acc_fgm_id_bitmap);

  } else {

    *acc_fgm_id_bitmap |= fgm_bit_map;

  }

  int fgm_chroma_idc = avm_rb_read_uvlc(rb);

  if (fgm_chroma_idc >= NUM_CHROMA_FORMATS) {

    avm_internal_error(&pbi->common.error, AVM_CODEC_UNSUP_BITSTREAM,

                       "Invalid fgm_chroma_idc [%d].", fgm_chroma_idc);

  }

  for (int j = 0; j < MAX_FGM_NUM; j++) {

    // This process overwrites the position(pbi->fg_list[fgm_id]) if the fgm_id

    // is the same.

    if (fgm_bit_map & (1 << j)) {

      int fgm_id = j;

      memset(&pbi->fgm_list[fgm_id], 0, sizeof(pbi->fgm_list[fgm_id]));

      pbi->fgm_list[fgm_id].fgm_id = fgm_id;

      pbi->fgm_list[fgm_id].fgm_mlayer_id = obu_mlayer_id;

      pbi->fgm_list[fgm_id].fgm_tlayer_id = obu_tlayer_id;

      pbi->fgm_list[fgm_id].fgm_chroma_idc = fgm_chroma_idc;

      read_film_grain_model(&pbi->fgm_list[fgm_id], fgm_chroma_idc, rb,

                            &pbi->common.error);

    }

  }

  if (av2_check_trailing_bits(pbi, rb) != 0) {

    // cm->error.error_code is already set.

    return 0;

  }

  return ((rb->bit_offset - saved_bit_offset + 7) >> 3);

}