/*
 * Copyright (c) 2023, 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 www.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
 * www.aomedia.org/license/patent.
 */
#include "iamf/cli/tests/cli_test_utils.h"

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <ios>
#include <list>
#include <memory>
#include <numbers>
#include <numeric>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "iamf/cli/audio_element_with_data.h"
#include "iamf/cli/audio_frame_with_data.h"
#include "iamf/cli/demixing_module.h"
#include "iamf/cli/obu_processor.h"
#include "iamf/cli/obu_with_data_generator.h"
#include "iamf/cli/parameter_block_with_data.h"
#include "iamf/cli/proto/codec_config.pb.h"
#include "iamf/cli/proto/mix_presentation.pb.h"
#include "iamf/cli/proto/user_metadata.pb.h"
#include "iamf/cli/proto_conversion/proto_to_obu/audio_element_generator.h"
#include "iamf/cli/proto_conversion/proto_to_obu/codec_config_generator.h"
#include "iamf/cli/proto_conversion/proto_to_obu/mix_presentation_generator.h"
#include "iamf/cli/renderer/audio_element_renderer_base.h"
#include "iamf/cli/user_metadata_builder/audio_element_metadata_builder.h"
#include "iamf/cli/user_metadata_builder/iamf_input_layout.h"
#include "iamf/cli/wav_reader.h"
#include "iamf/common/leb_generator.h"
#include "iamf/common/read_bit_buffer.h"
#include "iamf/common/utils/macros.h"
#include "iamf/common/write_bit_buffer.h"
#include "iamf/obu/audio_element.h"
#include "iamf/obu/codec_config.h"
#include "iamf/obu/decoder_config/aac_decoder_config.h"
#include "iamf/obu/decoder_config/flac_decoder_config.h"
#include "iamf/obu/decoder_config/lpcm_decoder_config.h"
#include "iamf/obu/decoder_config/opus_decoder_config.h"
#include "iamf/obu/demixing_info_parameter_data.h"
#include "iamf/obu/demixing_param_definition.h"
#include "iamf/obu/ia_sequence_header.h"
#include "iamf/obu/mix_presentation.h"
#include "iamf/obu/obu_base.h"
#include "iamf/obu/obu_header.h"
#include "iamf/obu/param_definitions.h"
#include "iamf/obu/types.h"
#include "src/google/protobuf/io/zero_copy_stream_impl.h"
#include "src/google/protobuf/repeated_ptr_field.h"
#include "src/google/protobuf/text_format.h"

namespace iamf_tools {

namespace {

void SetParamDefinitionCommonFields(DecodedUleb128 parameter_id,
                                    DecodedUleb128 parameter_rate,
                                    DecodedUleb128 duration,
                                    ParamDefinition& param_definition) {
  param_definition.parameter_id_ = parameter_id;
  param_definition.parameter_rate_ = parameter_rate;
  param_definition.param_definition_mode_ = 0;
  param_definition.reserved_ = 0;
  param_definition.duration_ = duration;
  param_definition.constant_subblock_duration_ = duration;
}

template <typename ParamDefinitionType>
void AddParamDefinition(DecodedUleb128 parameter_id,
                        DecodedUleb128 parameter_rate, DecodedUleb128 duration,
                        AudioElementObu& audio_element_obu,
                        ParamDefinitionType& param_definition) {
  SetParamDefinitionCommonFields(parameter_id, parameter_rate, duration,
                                 param_definition);

  // Add to the Audio Element OBU.
  audio_element_obu.audio_element_params_.emplace_back(
      AudioElementParam{param_definition});
}

}  // namespace

using ::absl_testing::IsOk;

absl::Status CollectObusFromIaSequence(
    ReadBitBuffer& read_bit_buffer, IASequenceHeaderObu& ia_sequence_header,
    absl::flat_hash_map<DecodedUleb128, CodecConfigObu>& codec_config_obus,
    absl::flat_hash_map<DecodedUleb128, AudioElementWithData>& audio_elements,
    std::list<MixPresentationObu>& mix_presentations,
    std::list<AudioFrameWithData>& audio_frames,
    std::list<ParameterBlockWithData>& parameter_blocks) {
  bool insufficient_data = false;
  auto obu_processor = ObuProcessor::Create(
      /*is_exhaustive_and_exact=*/false, &read_bit_buffer, insufficient_data);
  EXPECT_FALSE(insufficient_data);

  bool continue_processing = true;
  int temporal_unit_count = 0;
  LOG(INFO) << "Starting Temporal Unit OBU processing";
  while (continue_processing) {
    std::optional<ObuProcessor::OutputTemporalUnit> output_temporal_unit;
    RETURN_IF_NOT_OK(obu_processor->ProcessTemporalUnit(
        /*eos_is_end_of_sequence=*/true, output_temporal_unit,
        continue_processing));
    if (output_temporal_unit.has_value()) {
      audio_frames.splice(audio_frames.end(),
                          output_temporal_unit->output_audio_frames);
      parameter_blocks.splice(parameter_blocks.end(),
                              output_temporal_unit->output_parameter_blocks);
      temporal_unit_count++;
    }
  }
  LOG(INFO) << "Processed " << temporal_unit_count << " Temporal Unit OBUs";

  // Move the processed data to the output.
  ia_sequence_header = obu_processor->ia_sequence_header_;
  codec_config_obus.swap(obu_processor->codec_config_obus_);
  audio_elements.swap(obu_processor->audio_elements_);
  mix_presentations.swap(obu_processor->mix_presentations_);
  return absl::OkStatus();
}

void AddLpcmCodecConfig(
    DecodedUleb128 codec_config_id, uint32_t num_samples_per_frame,
    uint8_t sample_size, uint32_t sample_rate,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  // Initialize the Codec Config OBU.
  ASSERT_EQ(codec_config_obus.find(codec_config_id), codec_config_obus.end());

  auto obu = CodecConfigObu::Create(
      ObuHeader(), codec_config_id,
      {.codec_id = CodecConfig::kCodecIdLpcm,
       .num_samples_per_frame = num_samples_per_frame,
       .decoder_config = LpcmDecoderConfig{
           .sample_format_flags_bitmask_ = LpcmDecoderConfig::kLpcmLittleEndian,
           .sample_size_ = sample_size,
           .sample_rate_ = sample_rate}});
  ASSERT_THAT(obu, IsOk());

  codec_config_obus.emplace(codec_config_id, *std::move(obu));
}

void AddLpcmCodecConfigWithIdAndSampleRate(
    uint32_t codec_config_id, uint32_t sample_rate,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  // Many tests either don't care about the details. Or assumed these "default"
  // values.
  constexpr uint32_t kNumSamplesPerFrame = 8;
  constexpr uint8_t kSampleSize = 16;
  AddLpcmCodecConfig(codec_config_id, kNumSamplesPerFrame, kSampleSize,
                     sample_rate, codec_config_obus);
}

void AddOpusCodecConfig(
    uint32_t codec_config_id, uint32_t num_samples_per_frame,
    uint32_t sample_rate,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  // Initialize the Codec Config OBU.
  ASSERT_EQ(codec_config_obus.find(codec_config_id), codec_config_obus.end());

  auto obu = CodecConfigObu::Create(
      ObuHeader(), codec_config_id,
      {.codec_id = CodecConfig::kCodecIdOpus,
       .num_samples_per_frame = num_samples_per_frame,
       .decoder_config = OpusDecoderConfig{.version_ = 1,
                                           .pre_skip_ = 312,
                                           .input_sample_rate_ = sample_rate}});
  ASSERT_THAT(obu, IsOk());

  codec_config_obus.emplace(codec_config_id, *std::move(obu));
}

void AddOpusCodecConfigWithId(
    uint32_t codec_config_id,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  const uint32_t kNumSamplesPerFrame = 8;
  const uint32_t kSampleRate = 48000;

  AddOpusCodecConfig(codec_config_id, kNumSamplesPerFrame, kSampleRate,
                     codec_config_obus);
}

void AddFlacCodecConfig(
    uint32_t codec_config_id, uint32_t num_samples_per_frame,
    uint32_t sample_rate, uint8_t sample_size,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  // Initialize the Codec Config OBU.
  ASSERT_EQ(codec_config_obus.find(codec_config_id), codec_config_obus.end());

  auto obu = CodecConfigObu::Create(
      ObuHeader(), codec_config_id,
      {.codec_id = CodecConfig::kCodecIdFlac,
       .num_samples_per_frame = num_samples_per_frame,
       .decoder_config = FlacDecoderConfig(
           {{{.header = {.last_metadata_block_flag = true,
                         .block_type = FlacMetaBlockHeader::kFlacStreamInfo,
                         .metadata_data_block_length = 34},
              .payload = FlacMetaBlockStreamInfo{
                  .minimum_block_size =
                      static_cast<uint16_t>(num_samples_per_frame),
                  .maximum_block_size =
                      static_cast<uint16_t>(num_samples_per_frame),
                  .sample_rate = sample_rate,
                  .bits_per_sample = static_cast<uint8_t>(sample_size - 1),
                  .total_samples_in_stream = 0}}}})});
  ASSERT_THAT(obu, IsOk());

  codec_config_obus.emplace(codec_config_id, *std::move(obu));
}

void AddFlacCodecConfigWithId(
    uint32_t codec_config_id,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  const uint32_t kNumSamplesPerFrame = 16;
  const uint32_t kSampleRate = 48000;
  const uint8_t kSampleSize = 16;
  AddFlacCodecConfig(codec_config_id, kNumSamplesPerFrame, kSampleRate,
                     kSampleSize, codec_config_obus);
}

void AddAacCodecConfig(
    uint32_t codec_config_id, uint32_t num_samples_per_frame,
    uint32_t sample_rate,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  // Initialize the Codec Config OBU.
  ASSERT_EQ(codec_config_obus.find(codec_config_id), codec_config_obus.end());

  auto obu = CodecConfigObu::Create(
      ObuHeader(), codec_config_id,
      {.codec_id = CodecConfig::kCodecIdAacLc,
       .num_samples_per_frame = num_samples_per_frame,
       .decoder_config = AacDecoderConfig{
           .decoder_specific_info_ = {
               .audio_specific_config = {
                   .sample_frequency_index_ =
                       AudioSpecificConfig::SampleFrequencyIndex::kEscapeValue,
                   .sampling_frequency_ = sample_rate}}}});
  ASSERT_THAT(obu, IsOk());

  codec_config_obus.emplace(codec_config_id, *std::move(obu));
}

void AddAacCodecConfigWithId(
    uint32_t codec_config_id,
    absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus) {
  constexpr uint32_t kNumSamplesPerFrame = 1024;
  constexpr uint32_t kSampleRate = 48000;
  AddAacCodecConfig(codec_config_id, kNumSamplesPerFrame, kSampleRate,
                    codec_config_obus);
}

void AddAmbisonicsMonoAudioElementWithSubstreamIds(
    DecodedUleb128 audio_element_id, uint32_t codec_config_id,
    absl::Span<const DecodedUleb128> substream_ids,
    const absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus,
    absl::flat_hash_map<DecodedUleb128, AudioElementWithData>& audio_elements) {
  // Check the `codec_config_id` is known and this is a new
  // `audio_element_id`.
  auto codec_config_iter = codec_config_obus.find(codec_config_id);
  ASSERT_NE(codec_config_iter, codec_config_obus.end());
  ASSERT_EQ(audio_elements.find(audio_element_id), audio_elements.end());

  // Initialize the Audio Element OBU without any parameters.
  AudioElementObu obu = AudioElementObu(
      ObuHeader(), audio_element_id, AudioElementObu::kAudioElementSceneBased,
      0, codec_config_id);
  obu.InitializeParams(0);
  obu.InitializeAudioSubstreams(substream_ids.size());
  obu.audio_substream_ids_.assign(substream_ids.begin(), substream_ids.end());

  // Initialize to n-th order ambisonics. Choose the lowest order that can fit
  // all `substream_ids`. This may result in mixed-order ambisonics.
  uint8_t next_valid_output_channel_count;
  ASSERT_THAT(AmbisonicsConfig::GetNextValidOutputChannelCount(
                  substream_ids.size(), next_valid_output_channel_count),
              IsOk());
  EXPECT_THAT(obu.InitializeAmbisonicsMono(next_valid_output_channel_count,
                                           substream_ids.size()),
              IsOk());

  auto& channel_mapping =
      std::get<AmbisonicsMonoConfig>(
          std::get<AmbisonicsConfig>(obu.config_).ambisonics_config)
          .channel_mapping;
  // Map the first n channels from [0, n] in input order. Leave the rest of
  // the channels as unmapped.
  std::fill(channel_mapping.begin(), channel_mapping.end(),
            AmbisonicsMonoConfig::kInactiveAmbisonicsChannelNumber);
  std::iota(channel_mapping.begin(),
            channel_mapping.begin() + substream_ids.size(), 0);

  AudioElementWithData audio_element = {
      .obu = std::move(obu), .codec_config = &codec_config_iter->second};
  ASSERT_THAT(ObuWithDataGenerator::FinalizeAmbisonicsConfig(
                  audio_element.obu, audio_element.substream_id_to_labels),
              IsOk());

  audio_elements.emplace(audio_element_id, std::move(audio_element));
}

// Adds a scalable Audio Element OBU based on the input arguments.
void AddScalableAudioElementWithSubstreamIds(
    IamfInputLayout input_layout, DecodedUleb128 audio_element_id,
    uint32_t codec_config_id, absl::Span<const DecodedUleb128> substream_ids,
    const absl::flat_hash_map<uint32_t, CodecConfigObu>& codec_config_obus,
    absl::flat_hash_map<DecodedUleb128, AudioElementWithData>& audio_elements) {
  google::protobuf::RepeatedPtrField<
      iamf_tools_cli_proto::AudioElementObuMetadata>
      audio_element_metadatas;
  AudioElementMetadataBuilder builder;

  auto& new_audio_element_metadata = *audio_element_metadatas.Add();
  ASSERT_THAT(builder.PopulateAudioElementMetadata(
                  audio_element_id, codec_config_id, input_layout,
                  new_audio_element_metadata),
              IsOk());
  // Check that this is a scalable Audio Element, and override the substream
  // IDs.
  ASSERT_TRUE(new_audio_element_metadata.has_scalable_channel_layout_config());
  ASSERT_EQ(new_audio_element_metadata.audio_substream_ids().size(),
            substream_ids.size());
  for (int i = 0; i < substream_ids.size(); ++i) {
    new_audio_element_metadata.mutable_audio_substream_ids()->Set(
        i, substream_ids[i]);
  }

  // Generate the Audio Element OBU.
  AudioElementGenerator generator(audio_element_metadatas);
  ASSERT_THAT(generator.Generate(codec_config_obus, audio_elements), IsOk());
}

void AddMixPresentationObuWithAudioElementIds(
    DecodedUleb128 mix_presentation_id,
    const std::vector<DecodedUleb128>& audio_element_ids,
    DecodedUleb128 common_parameter_id, DecodedUleb128 common_parameter_rate,
    std::list<MixPresentationObu>& mix_presentations) {
  // Configure one of the simplest mix presentation. Mix presentations REQUIRE
  // at least one sub-mix and a stereo layout.
  AddMixPresentationObuWithConfigurableLayouts(
      mix_presentation_id, audio_element_ids, common_parameter_id,
      common_parameter_rate,
      {LoudspeakersSsConventionLayout::kSoundSystemA_0_2_0}, mix_presentations);
}

void AddMixPresentationObuWithConfigurableLayouts(
    DecodedUleb128 mix_presentation_id,
    const std::vector<DecodedUleb128>& audio_element_ids,
    DecodedUleb128 common_parameter_id, DecodedUleb128 common_parameter_rate,
    const std::vector<LoudspeakersSsConventionLayout::SoundSystem>&
        sound_system_layouts,
    std::list<MixPresentationObu>& mix_presentations) {
  MixGainParamDefinition common_mix_gain_param_definition;
  common_mix_gain_param_definition.parameter_id_ = common_parameter_id;
  common_mix_gain_param_definition.parameter_rate_ = common_parameter_rate;
  common_mix_gain_param_definition.param_definition_mode_ = true;
  common_mix_gain_param_definition.default_mix_gain_ = 0;
  std::vector<MixPresentationLayout> layouts;
  for (const auto& sound_system : sound_system_layouts) {
    layouts.push_back(
        {.loudness_layout = {.layout_type =
                                 Layout::kLayoutTypeLoudspeakersSsConvention,
                             .specific_layout =
                                 LoudspeakersSsConventionLayout{
                                     .sound_system = sound_system,
                                     .reserved = 0}},
         .loudness = {
             .info_type = 0, .integrated_loudness = 0, .digital_peak = 0}});
  }

  std::vector<MixPresentationSubMix> sub_mixes = {
      {.output_mix_gain = common_mix_gain_param_definition,
       .layouts = layouts}};
  for (const auto& audio_element_id : audio_element_ids) {
    sub_mixes[0].audio_elements.push_back({
        .audio_element_id = audio_element_id,
        .localized_element_annotations = {},
        .rendering_config =
            {.headphones_rendering_mode =
                 RenderingConfig::kHeadphonesRenderingModeStereo,
             .reserved = 0,
             .rendering_config_extension_size = 0,
             .rendering_config_extension_bytes = {}},
        .element_mix_gain = common_mix_gain_param_definition,
    });
  }

  mix_presentations.push_back(
      MixPresentationObu(ObuHeader(), mix_presentation_id,
                         /*count_label=*/0, {}, {}, sub_mixes));
}

void AddParamDefinitionWithMode0AndOneSubblock(
    DecodedUleb128 parameter_id, DecodedUleb128 parameter_rate,
    DecodedUleb128 duration,
    absl::flat_hash_map<DecodedUleb128, MixGainParamDefinition>&
        param_definitions) {
  MixGainParamDefinition param_definition;
  SetParamDefinitionCommonFields(parameter_id, parameter_rate, duration,
                                 param_definition);
  param_definitions.emplace(parameter_id, param_definition);
}

void AddDemixingParamDefinition(DecodedUleb128 parameter_id,
                                DecodedUleb128 parameter_rate,
                                DecodedUleb128 duration,
                                AudioElementObu& audio_element_obu) {
  DemixingParamDefinition param_definition;

  // Specific fields of demixing param definitions.
  param_definition.default_demixing_info_parameter_data_.dmixp_mode =
      DemixingInfoParameterData::kDMixPMode1;
  param_definition.default_demixing_info_parameter_data_.reserved = 0;
  param_definition.default_demixing_info_parameter_data_.default_w = 10;
  param_definition.default_demixing_info_parameter_data_
      .reserved_for_future_use = 0;

  AddParamDefinition(parameter_id, parameter_rate, duration, audio_element_obu,
                     param_definition);
}

void AddReconGainParamDefinition(DecodedUleb128 parameter_id,
                                 DecodedUleb128 parameter_rate,
                                 DecodedUleb128 duration,
                                 AudioElementObu& audio_element_obu) {
  ReconGainParamDefinition param_definition(
      audio_element_obu.GetAudioElementId());

  AddParamDefinition(parameter_id, parameter_rate, duration, audio_element_obu,
                     param_definition);
}

WavReader CreateWavReaderExpectOk(const std::string& filename,
                                  int num_samples_per_frame) {
  auto wav_reader = WavReader::CreateFromFile(filename, num_samples_per_frame);
  EXPECT_THAT(wav_reader, IsOk());
  return std::move(*wav_reader);
}

void RenderAndFlushExpectOk(
    const LabeledFrame& labeled_frame, AudioElementRendererBase* renderer,
    std::vector<std::vector<InternalSampleType>>& output_samples) {
  ASSERT_NE(renderer, nullptr);
  EXPECT_THAT(renderer->RenderLabeledFrame(labeled_frame), IsOk());
  EXPECT_THAT(renderer->Finalize(), IsOk());
  EXPECT_TRUE(renderer->IsFinalized());
  renderer->Flush(output_samples);
}

std::string GetAndCleanupOutputFileName(absl::string_view suffix) {
  const testing::TestInfo* const test_info =
      testing::UnitTest::GetInstance()->current_test_info();
  std::string filename = absl::StrCat(test_info->name(), "-",
                                      test_info->test_suite_name(), suffix);

  // It is possible that the test suite name contains the '/' character.
  // Replace it with '-' to form a legal file name.
  absl::StrReplaceAll({{"/", "-"}}, &filename);
  const std::filesystem::path test_specific_filename =
      std::filesystem::path(::testing::TempDir()) / filename;

  std::filesystem::remove(test_specific_filename);
  return test_specific_filename.string();
}

std::string GetAndCreateOutputDirectory(absl::string_view suffix) {
  const std::string output_directory = GetAndCleanupOutputFileName(suffix);
  std::error_code error_code;
  EXPECT_TRUE(
      std::filesystem::create_directories(output_directory, error_code));
  return output_directory;
}

std::vector<uint8_t> SerializeObusExpectOk(
    const std::list<const ObuBase*>& obus, const LebGenerator& leb_generator) {
  using ::absl_testing::IsOk;
  WriteBitBuffer serialized_obus(0, leb_generator);
  for (const auto* obu : obus) {
    EXPECT_NE(obu, nullptr);
    EXPECT_THAT(obu->ValidateAndWriteObu(serialized_obus), IsOk());
  }

  return serialized_obus.bit_buffer();
}

void ParseUserMetadataAssertSuccess(
    const std::string& textproto_filename,
    iamf_tools_cli_proto::UserMetadata& user_metadata) {
  ASSERT_TRUE(std::filesystem::exists(textproto_filename));
  std::ifstream user_metadata_file(textproto_filename, std::ios::in);
  google::protobuf::io::IstreamInputStream input_stream(&user_metadata_file);
  ASSERT_TRUE(
      google::protobuf::TextFormat::Parse(&input_stream, &user_metadata));
}

double GetLogSpectralDistance(
    const absl::Span<const InternalSampleType>& first_log_spectrum,
    const absl::Span<const InternalSampleType>& second_log_spectrum) {
  const int num_samples = first_log_spectrum.size();
  if (num_samples != second_log_spectrum.size()) {
    LOG(ERROR) << "Spectrum sizes are not equal.";
    return false;
  }
  double log_spectral_distance = 0.0;
  for (int i = 0; i < num_samples; ++i) {
    log_spectral_distance += (first_log_spectrum[i] - second_log_spectrum[i]) *
                             (first_log_spectrum[i] - second_log_spectrum[i]);
  }
  return (10 * std::sqrt(log_spectral_distance / num_samples));
}

uint32_t GetSampleRateForCodecConfigMetadata(
    const iamf_tools_cli_proto::UserMetadata& user_metadata,
    DecodedUleb128 codec_config_id) {
  auto codec_config_generator =
      CodecConfigGenerator(user_metadata.codec_config_metadata());
  absl::flat_hash_map<uint32_t, CodecConfigObu> codec_config_obus;
  EXPECT_THAT(codec_config_generator.Generate(codec_config_obus), IsOk());
  if (codec_config_obus.contains(codec_config_id)) {
    return codec_config_obus.at(codec_config_id).GetOutputSampleRate();
  }
  return 0;
}

uint32_t GetSampleRateForAudioElementMetadata(
    const iamf_tools_cli_proto::UserMetadata& user_metadata,
    DecodedUleb128 audio_element_id) {
  for (const auto& audio_element : user_metadata.audio_element_metadata()) {
    if (audio_element.audio_element_id() == audio_element_id) {
      return GetSampleRateForCodecConfigMetadata(
          user_metadata, audio_element.codec_config_id());
    }
  }
  return 0;
}

std::vector<DecodeSpecification> GetDecodeSpecifications(
    const iamf_tools_cli_proto::UserMetadata& user_metadata) {
  std::vector<DecodeSpecification> decode_specifications;
  for (const auto& mix_presentation :
       user_metadata.mix_presentation_metadata()) {
    for (int i = 0; i < mix_presentation.sub_mixes_size(); ++i) {
      for (int j = 0; j < mix_presentation.sub_mixes(i).layouts_size(); ++j) {
        DecodeSpecification decode_specification;
        decode_specification.mix_presentation_id =
            mix_presentation.mix_presentation_id();
        decode_specification.sub_mix_index = i;
        if (mix_presentation.sub_mixes(i)
                .layouts(j)
                .loudness_layout()
                .has_ss_layout()) {
          auto sound_system_status = MixPresentationGenerator::CopySoundSystem(
              mix_presentation.sub_mixes(i)
                  .layouts(j)
                  .loudness_layout()
                  .ss_layout()
                  .sound_system(),
              decode_specification.sound_system);
          if (!sound_system_status.ok()) {
            LOG(ERROR) << "Failed to copy sound system: "
                       << sound_system_status;
            continue;
          }
        }
        decode_specification.sample_rate = GetSampleRateForAudioElementMetadata(
            user_metadata,
            mix_presentation.sub_mixes(i).audio_elements(0).audio_element_id());
        decode_specification.layout_index = j;
        decode_specifications.push_back(decode_specification);
      }
    }
  }
  return decode_specifications;
}

std::vector<InternalSampleType> Int32ToInternalSampleType(
    absl::Span<const int32_t> samples) {
  std::vector<InternalSampleType> result(samples.size());
  Int32ToInternalSampleType(samples, absl::MakeSpan(result));
  return result;
}

std::vector<std::vector<InternalSampleType>> Int32ToInternalSampleType2D(
    const std::vector<std::vector<int32_t>>& samples) {
  const auto num_channels = samples.size();
  std::vector<std::vector<InternalSampleType>> internal_samples(num_channels);
  for (int c = 0; c < num_channels; c++) {
    internal_samples[c] = Int32ToInternalSampleType(samples[c]);
  }

  return internal_samples;
}

std::vector<InternalSampleType> GenerateSineWav(uint64_t start_tick,
                                                uint32_t num_samples,
                                                uint32_t sample_rate_hz,
                                                double frequency_hz,
                                                double amplitude) {
  std::vector<InternalSampleType> samples(num_samples, 0.0);
  constexpr double kPi = std::numbers::pi_v<InternalSampleType>;
  const double time_base = 1.0 / sample_rate_hz;

  for (int frame_tick = 0; frame_tick < num_samples; ++frame_tick) {
    const double t = start_tick + frame_tick;
    samples[frame_tick] =
        amplitude * sin(2.0 * kPi * frequency_hz * t * time_base);
  }
  return samples;
}

void AccumulateZeroCrossings(
    absl::Span<const absl::Span<const InternalSampleType>> channel_time_samples,
    std::vector<ZeroCrossingState>& zero_crossing_states,
    std::vector<int>& zero_crossing_counts) {
  using enum ZeroCrossingState;
  const auto num_channels = channel_time_samples.size();
  // Seed the data structures, or check they contain the right number of
  // channels.
  if (zero_crossing_counts.empty()) {
    zero_crossing_counts.resize(num_channels, 0);
  } else {
    ASSERT_EQ(num_channels, zero_crossing_counts.size());
  }
  if (zero_crossing_states.empty()) {
    zero_crossing_states.resize(num_channels, ZeroCrossingState::kUnknown);
  } else {
    ASSERT_EQ(num_channels, zero_crossing_states.size());
  }

  // Zero crossing threshold determined empirically for -18 dB sine waves to
  // skip encoding artifacts (e.g. a small ringing artifact < -40 dB after
  // the sine wave stopped.)  Note that -18 dB would correspond to dividing
  // by 8, while dividing by 100 is -40 dB.
  constexpr InternalSampleType kThreshold = 0.01;

  for (int c = 0; c < num_channels; c++) {
    const auto& channel = channel_time_samples[c];
    for (int t = 0; t < channel.size(); ++t) {
      ZeroCrossingState next_state = (channel[t] > kThreshold)    ? kPositive
                                     : (channel[t] < -kThreshold) ? kNegative
                                                                  : kUnknown;
      if (next_state == kUnknown) {
        // Don't do anything if it's not clearly positive or negative.
        continue;
      } else if (zero_crossing_states[c] != next_state) {
        // If we clearly flipped states, count it as a zero crossing.
        zero_crossing_counts[c]++;
        zero_crossing_states[c] = next_state;
      }
    }
  }
}

absl::Status ReadFileToBytes(const std::filesystem::path& file_path,
                             std::vector<uint8_t>& buffer) {
  if (!std::filesystem::exists(file_path)) {
    return absl::NotFoundError("File not found.");
  }
  std::ifstream ifs(file_path, std::ios::binary | std::ios::in);

  // Increase the size of the buffer. Write to the original end (before
  // resizing).
  const auto file_size = std::filesystem::file_size(file_path);
  const auto original_buffer_size = buffer.size();
  buffer.resize(original_buffer_size + file_size);
  ifs.read(reinterpret_cast<char*>(buffer.data() + original_buffer_size),
           file_size);
  return absl::OkStatus();
}

absl::Status EverySecondTickResampler::PushFrameDerived(
    absl::Span<const absl::Span<const InternalSampleType>>
        channel_time_samples) {
  for (int c = 0; c < num_channels_; c++) {
    // `SampleProcessorBase` should ensure this.
    EXPECT_TRUE(output_channel_time_samples_[c].empty());
    for (size_t i = 0; i < channel_time_samples[c].size(); ++i) {
      if (i % 2 == 1) {
        output_channel_time_samples_[c].push_back(channel_time_samples[c][i]);
      }
    }
  }
  return absl::OkStatus();
}

absl::Status EverySecondTickResampler::FlushDerived() {
  // `SampleProcessorBase` should ensure this.
  for (int c = 0; c < num_channels_; c++) {
    // `SampleProcessorBase` should ensure this.
    EXPECT_TRUE(output_channel_time_samples_[c].empty());
  }
  return absl::OkStatus();
}

absl::Status OneFrameDelayer::PushFrameDerived(
    absl::Span<const absl::Span<const InternalSampleType>>
        channel_time_samples) {
  // Swap the delayed samples with the output samples from the base class.
  std::swap(delayed_samples_, output_channel_time_samples_);

  num_delayed_ticks_ = 0;
  for (int c = 0; c < num_channels_; c++) {
    delayed_samples_[c].resize(channel_time_samples[c].size());

    // Cache the new samples to delay.
    std::copy(channel_time_samples[c].begin(), channel_time_samples[c].end(),
              delayed_samples_[c].begin());

    if (num_delayed_ticks_ == 0) {
      num_delayed_ticks_ = channel_time_samples[c].size();
    }
    EXPECT_EQ(num_delayed_ticks_, channel_time_samples[c].size());
  }

  return absl::OkStatus();
}

absl::Status OneFrameDelayer::FlushDerived() {
  // Pushing in an empty frame will cause the delayed frame to be available.
  auto empty_frame =
      std::vector<absl::Span<const InternalSampleType>>(num_channels_);
  return PushFrameDerived(absl::MakeConstSpan(empty_frame));
}

}  // namespace iamf_tools