/*
   HARK Library
   Copyright (C) 2021 Honda Research Institute Japan Inc.
 */

#include "hark-lib/node/AudioStreamFromWave.h"

#include "hark-lib/matrix.h"
#include <istream>
#include <fstream>

namespace hark
{
namespace node
{
    AudioStreamFromWave::AudioStreamFromWave():
        NodeDef("AudioStreamFromWave")
    {
        *this
            << Terminal::in("INPUT", "An audio input stream (IStream).")
            << Terminal::in("LENGTH", "The frame length of each channel (in samples) [default: 512].").def(512)
            << Terminal::in("ADVANCE", "The shift length beween adjacent frames (in samples)[default: 160].").def(160)
            << Terminal::in("USE_WAIT", "If true, real recording is simulated [default: false].").def(false)
            >> Terminal::out("AUDIO", "Windowed sound data. A row index is channel, and a column index is time.")
            >> Terminal::out("NOT_EOF", "True unless End-of-File.");
    }

    void
    AudioStreamFromWave::reset()
    {
        this->_length = 0;
        this->_advance = 0;
        this->_buffer.clear();
        this->_fbuffer.clear();
    }

    void
    AudioStreamFromWave::initialize(std::istream &streamInput, InputDataBundle &input)
    {
        this->read_wave_header(streamInput, this->_channels, this->_bytes_per_sample, this->_sampling_rate);

        this->_length = input["LENGTH"].last().as<int>();
        this->_advance = input["ADVANCE"].last().as<int>();
        this->_buffer.clear();
        this->_fbuffer.resize(this->_channels * this->_length);
    }

    void
    AudioStreamFromWave::read_wave_header(std::istream &file, int &channels, int &bytes_per_sample, int &sampling_rate)
    {
        char buf[5] = {'\0'};
        unsigned size;
        //short wFormatID;
        short wChannels;
        short wBits;
        unsigned dwSamplingRate;

        // std::cout << "read_wave_header: " << (bool)file << std::endl;

        file.read(buf, 4);
        if(!file){
            std::stringstream ss;
            ss << "Failed to read file marker: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
            throw RuntimeError(ss.str());
        }else if(std::string("RIFF") != buf){
            throw RuntimeError("Not a RIFF file");
        }
        file.read((char *)&size, 4);
        if(!file){
            std::stringstream ss;
            ss << "Failed to read file size: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
            throw RuntimeError(ss.str());
        }
        
        file.read(buf, 4);
        if(!file){
            std::stringstream ss;
            ss << "Failed to read wave marker: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
            throw RuntimeError(ss.str());
        }else if(std::string("WAVE") != buf){
            throw RuntimeError("Not a WAVE file");
        }

        while(!file.eof()){
            file.read(buf, 4);
            if(!file){
                std::stringstream ss;
                ss << "Failed to read chunk marker: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                throw RuntimeError(ss.str());
            }

            if(std::string("fact") == buf || std::string("LIST") == buf){
                file.read((char *)&size, 4);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to read chunk size: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }
                file.seekg(size, std::ios::cur);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to skip chunk: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

            }else if(std::string("fmt ") == buf){
                file.read((char *)&size, 4);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to read format chunk size: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

                file.seekg(2, std::ios::cur);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to skip extra format values: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

                file.read((char *)&wChannels, sizeof(short));
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to read channel count: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }
                // num_channels = wChannels;
                file.read((char *)&dwSamplingRate, sizeof(unsigned));
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to read sampling rate: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

                // length_ms = (length * 1000) / dwSamplingRate;
                // advance_ms = (advance * 1000) / dwSamplingRate;
                // overlapped_usec = (length - advance) * 1000000 / dwSamplingRate;

                file.seekg(6, std::ios::cur);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to skip extra format values: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

                file.read((char *)&wBits, sizeof(short));
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to read bits per sample: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

                int nbits = wBits;
                int nbytes = nbits / 8;
                file.seekg(size - 16, std::ios::cur);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to skip extra format values: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }

            }else if(std::string("data") == buf){
                file.seekg(4, std::ios::cur);
                if(!file){
                    std::stringstream ss;
                    ss << "Failed to skip data size: eof=" << file.eof() << ", fail=" << file.fail() << ", bad=" << file.bad() << ", read " << file.gcount() << " bytes.";
                    throw RuntimeError(ss.str());
                }
                break;

            }else{
                throw RuntimeError("Unexpected WAVE chunk header");
            }
        }

        channels = wChannels;
        bytes_per_sample = wBits / 8;
        sampling_rate = dwSamplingRate;
    }

    bool
    AudioStreamFromWave::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        data streamInputData = input["INPUT"].last();
        std::istream &streamInput = *streamInputData.cast<std::istream *>();
        if(count == 0){
            try{
                this->initialize(streamInput, input);
            }catch(std::exception &ex){
                std::stringstream ss;
                ss << "AudioStreamFromWave: Failed to initialize: " << ex.what();
                throw RuntimeError(ss.str());
            }
        }

        data audioOutputData = data::create<hark::Matrix<float>>(this->_channels, this->_length);
        auto &audioOutput = audioOutputData.cast<hark::Matrix<float>>();
        bool eofOutput = false;

        const int bytes_per_frame = this->_channels * this->_bytes_per_sample;
        auto &buffer = this->_buffer;
        auto &fbuffer = this->_fbuffer;

        if(buffer.empty()){
            buffer.resize(this->_length * bytes_per_frame);
            streamInput.read(buffer.data(), this->_length * bytes_per_frame);
            auto end = this->convert_bytes_to_float(buffer.begin(), buffer.end(), fbuffer.begin(), this->_bytes_per_sample);
            // std::cout << "convert: " << (end - fbuffer.begin()) << " elements." << std::endl;
        }else{
            memmove(buffer.data(), buffer.data() + (this->_advance * bytes_per_frame), (this->_length - this->_advance) * bytes_per_frame);
            streamInput.read(buffer.data() + (this->_length - this->_advance) * bytes_per_frame, this->_advance * bytes_per_frame);
            auto end = this->convert_bytes_to_float(buffer.begin(), buffer.end(), fbuffer.begin(), this->_bytes_per_sample);
        }
        // audioOutput = Eigen::Map<hark::Matrix<float>>(fbuffer.data(), this->_channels, this->_length);
        // fbuffer[0] = 1000 + count;
        audioOutput = hark::Matrix<float>::from_col_major_image(fbuffer.data(), this->_channels, this->_length);

        if(streamInput.eof()){
            eofOutput = true;
        }else{
            // std::cout << "AudioStreamFromWave: count=" << count << ", output=" << audioOutput << std::endl;
            output["AUDIO"].push(audioOutputData);
        }
        output["NOT_EOF"].push(data(!eofOutput));

        return true;
    }

    std::vector<float>::iterator
    AudioStreamFromWave::convert_bytes_to_float(std::vector<std::istream::char_type>::iterator start,
                                                std::vector<std::istream::char_type>::iterator end,
                                                std::vector<float>::iterator into,
                                                size_t bytes_per_sample)
    {
        std::vector<std::istream::char_type>::iterator sit;
        std::vector<float>::iterator dit;

        switch(bytes_per_sample){
        case 3:
            for(sit = start, dit = into; sit != end; ){
                std::istream::char_type ch2 = *sit++;
                std::istream::char_type ch1 = *sit++;
                std::istream::char_type ch0 = *sit++;
                *dit++ = (float)(((int)ch0 << 16) + ((int)(unsigned char)ch1 << 8) + (int)(unsigned char)ch2);
            }
            break;

        case 2:
            for(sit = start, dit = into; sit != end; ){
                std::istream::char_type ch1 = *sit++;
                std::istream::char_type ch0 = *sit++;
                *dit = (float)(((int)ch0 << 8) + (int)(unsigned char)ch1);
                // std::cout << *dit << " = " << (int)ch0 << ", " << (int)(unsigned char)ch1 << std::endl;
                dit++;
            }
            break;

        }
        return dit;
    }

}
}

/* end of file */
