#pragma once

#include <cstdint>

#include <array>
#include <fstream>
#include <ostream>
#include <string>

using std::array;
using std::ofstream;
using std::ostream;
using std::string;

class WavWriter {
public:
    enum FormatCode : uint16_t {
        WAVE_FORMAT_PCM = 0x01,
        WAVE_FORMAT_IEEE_FLOAT = 0x03,
        WAVE_FORMAT_ALAW = 0x06,
        WAVE_FORMAT_MULAW = 0x07,
        WAVE_FORMAT_EXTENSIBLE = 0xfffe,
    };

    WavWriter(const string& filename, WavWriter::FormatCode format,
        uint16_t nChannels, uint16_t nBitsPerSample, uint32_t nSamplesPerSec)
        : f(filename, std::ofstream::binary)
        , format(format)
        , nChannels(nChannels)
        , nBitsPerSample(nBitsPerSample)
        , nSamplesPerSec(nSamplesPerSec)
    {
        f.write("RIFF\0\0\0\0WAVE", 12);
        updateHeader();
        f.seekp(0, std::ios_base::end);
    };
    ~WavWriter() { f.close(); }

    template <typename T> void writeSamples(const T data[], size_t n)
    {
        f.write(reinterpret_cast<const char*>(data), n);
        updateHeader();
    };

    void close();

private:
    uint32_t dataSize;
    FormatCode format;
    uint16_t nChannels;
    uint32_t nSamplesPerSec;
    uint16_t nBitsPerSample;
    ofstream f;

    void updateHeader()
    {
        auto curPos = f.tellp();
        dataSize = curPos - std::streamoff(8);
        f.seekp(0);
        // Update main header total size
        f.write("RIFF", 4);
        f.write(reinterpret_cast<const char*>(&dataSize), sizeof(dataSize));
        f.write("WAVE", 4);
        // Write new format and data size
        WavWriter::fmtChunk fmt(this);
        fmt.writeTo(f);
        f.write("data", 4);
        uint32_t ckSize = curPos - f.tellp();
        f.write(reinterpret_cast<const char*>(&ckSize), sizeof(ckSize));

        f.seekp(curPos);
    };

    struct Chunk {
    public:
        char ckID[4];

    protected:
        Chunk(WavWriter* parent, const char* ckID)
            : parent(parent)
            , ckID{ ckID[0], ckID[1], ckID[2], ckID[3] } {};
        virtual void writeTo(ostream& of) = 0;

        WavWriter* parent;
    };

    struct fmtChunk : public Chunk {
        fmtChunk(WavWriter* parent)
            : Chunk(parent, "fmt ")
            {};

        uint32_t cksize;
        uint16_t cbSize = 0;
        uint32_t dwChannelMask;
        char SubFormat[16];

        void writeTo(ostream& of)
        {
            of.write(this->ckID, sizeof(this->ckID));
            cksize = parent->format == WAVE_FORMAT_PCM ? 16 : 18;
            of.write(reinterpret_cast<const char*>(&cksize), sizeof(cksize));
            of.write(reinterpret_cast<const char*>(&parent->format),
                sizeof(parent->format));
            of.write(reinterpret_cast<const char*>(&parent->nChannels),
                sizeof(parent->nChannels));
            of.write(reinterpret_cast<const char*>(&parent->nSamplesPerSec),
                sizeof(parent->nSamplesPerSec));
            uint32_t byteRate = parent->nSamplesPerSec * parent->nChannels
                * parent->nBitsPerSample / 8;
            of.write(
                reinterpret_cast<const char*>(&byteRate), sizeof(byteRate));
            uint16_t blockAlign
                = parent->nChannels * parent->nBitsPerSample / 8;
            of.write(
                reinterpret_cast<const char*>(&blockAlign), sizeof(blockAlign));
            of.write(reinterpret_cast<const char*>(&parent->nBitsPerSample),
                sizeof(parent->nBitsPerSample));
            if (parent->format != WAVE_FORMAT_PCM)
                of.write(
                    reinterpret_cast<const char*>(&cbSize), sizeof(cbSize));
            if (parent->format == WAVE_FORMAT_EXTENSIBLE) {
                of.write(reinterpret_cast<const char*>(&parent->nBitsPerSample),
                    sizeof(parent->nBitsPerSample));
                of.write(reinterpret_cast<const char*>(&dwChannelMask),
                    sizeof(dwChannelMask));
                of.write(reinterpret_cast<const char*>(SubFormat),
                    sizeof(SubFormat));
            }
        }
    };
    struct factChunk : public Chunk {
        factChunk(WavWriter* parent)
            : Chunk(parent, "fact"){};
        uint32_t cksize = 4;
        uint32_t dwSampleLength();
    };
};