#ifndef __WAVWRITER_H__
#define __WAVWRITER_H__

#include <cstring>
#include <fstream>
#include <memory>
#include <string>

template <typename SampleType>
class WavWriter {
    struct Header {
        char riffHeader[4];     // "RIFF"
        uint32_t wavSize;       // File size - 8
        char waveHeader[4];     // "WAVE"
        char fmtHeader[4];      // "fmt "
        uint32_t fmtChunkSize;  // Size of the fmt chunk
        uint16_t audioFormat;   // Audio format (1 for PCM)
        uint16_t numChannels;   // Number of channels
        uint32_t sampleRate;    // Sampling frequency
        uint32_t byteRate;      // (SampleRate * NumChannels * BitsPerSample) / 8
        uint16_t blockAlign;    // (NumChannels * BitsPerSample) / 8
        uint16_t bitsPerSample; // Bits per sample
        char dataHeader[4];     // "data"
        uint32_t dataSize;      // Size of the data section
    };

public:
    WavWriter(const std::string &filename, uint16_t channels, uint32_t sampleRate) {
        m_ofs = std::make_unique<std::ofstream>(filename, std::ofstream::binary);
        std::memcpy(m_header.riffHeader, "RIFF", 4);
        m_header.wavSize = 0;
        std::memcpy(m_header.waveHeader, "WAVE", 4);
        std::memcpy(m_header.fmtHeader, "fmt ", 4);
        m_header.fmtChunkSize = 16;
        m_header.audioFormat = 1;
        m_header.numChannels = channels;
        m_header.sampleRate = sampleRate;
        m_header.byteRate = sampleRate * channels * sizeof(SampleType);
        m_header.blockAlign = channels * sizeof(SampleType);
        m_header.bitsPerSample = sizeof(SampleType) * 8;
        std::memcpy(m_header.dataHeader, "data", 4);
        m_header.dataSize = 0;

        m_ofs->write(reinterpret_cast<const char *>(&m_header), sizeof(m_header));
    }
    ~WavWriter() {
        close();
    }

    void write(const char *data, size_t size) {
        if (m_ofs) {
            m_ofs->write(data, size);
            m_header.dataSize += size;
        }
    }

    void close() {
        if (m_ofs) {
            m_header.wavSize = 36 + m_header.dataSize;
            m_ofs->seekp(0, std::ios::beg);
            m_ofs->write(reinterpret_cast<const char *>(&m_header), sizeof(m_header));
            m_ofs.reset();
        }
    }

private:
    std::unique_ptr<std::ofstream> m_ofs;
    Header m_header;
};
#endif // __WAVWRITER_H__