/*
* Copyright (c) 2017-2022, NVIDIA Corporation.  All rights reserved.
*
* NVIDIA Corporation and its licensors retain all intellectual property
* and proprietary rights in and to this software, related documentation
* and any modifications thereto.  Any use, reproduction, disclosure or
* distribution of this software and related documentation without an express
* license agreement from NVIDIA Corporation is strictly prohibited.
*/

#pragma once

#include <stdint.h>
#include <string.h>
#include <assert.h>

#include <memory>
#include <string>
#include <vector>

#include "wave.hpp"

#define MAKEFOURCC(a, b, c, d) ( (uint32_t) (((d) << 24) | ((c) << 16) | ((b) << 8) | (a)) )

struct RiffHeader {
  uint32_t chunkId;
  uint32_t chunkSize;
  uint32_t fileTag;
};

struct RiffChunk {
  uint32_t chunkId;
  uint32_t chunkSize;
};

struct WaveFileInfo {
  waveFormat_ext wfx;
  uint8_t* audioData;
  uint32_t audioDataSize;
};

enum WaveFileFlags {
  READ_WAVEFILE = 0,
  WRITE_WAVEFILE = 1
};

class CWaveFileRead {
 public:
  explicit CWaveFileRead(std::string wavFile);
  uint32_t GetSampleRate() const { return m_WaveFormatEx.nSamplesPerSec; }
  uint32_t GetRawPCMDataSizeInBytes() const { return m_WaveDataSize; }
  uint32_t GetNumSamples() const { return m_nNumSamples; }
  uint32_t GetNumAlignedSamples() const { return m_NumAlignedSamples; }
  uint8_t* GetRawPCMData() { return m_WaveData.get(); }
  const float *GetFloatPCMData();
  const float *GetFloatPCMDataAligned(int alignSamples);
  waveFormat_ext& GetWaveFormat() { return m_WaveFormatEx; }
  int GetBitsPerSample();
  bool isValid() const { return validFile; }
  std::vector<float>* GetFloatVector();

 private:
  const RiffChunk* FindChunk(const uint8_t* data, size_t sizeBytes, uint32_t fourcc);
  int readPCM(const char* szFileName);

 private:
  std::string m_wavFile;
  uint32_t m_nNumSamples;
  bool validFile;
  std::unique_ptr<uint8_t[]> m_WaveData;
  std::vector<float> m_floatWaveData;
  uint32_t m_WaveDataSize;
  std::unique_ptr<float[]> m_floatWaveDataAligned;
  waveFormat_ext m_WaveFormatEx;
  uint32_t m_NumAlignedSamples;
};

class CWaveFileWrite {
 public:
  CWaveFileWrite(std::string wavFile, uint32_t samplesPerSec, uint32_t numChannels,
                 uint16_t bitsPerSample, bool isFloat);
  ~CWaveFileWrite();
  bool initFile();
  // can be called 'n' times.
  bool writeChunk(const void *data, uint32_t len);
  bool commitFile();
  uint32_t getWrittenCount() { return m_cumulativeCount; }
  std::string getFileName() { return m_wavFile; }
 private:
  bool m_validState = false;
  std::string m_wavFile;
  FILE *m_fp = nullptr;
  uint32_t m_cumulativeCount = 0;
  waveFormat_ext wfx;
  bool m_commitDone = false;
};

bool ReadWavFile(const std::string& filename, uint32_t expected_sample_rate,
                 std::vector<float>** data, unsigned* original_num_samples,
                 std::vector<int>* file_end_offset, int align_samples = -1);