#include <iostream>
#include <fstream>
#include <cstring>
#include <array>
#include <list>
#include <vector>
#include <cstdint>
#include <sstream>
#include "huffman.hpp"

int show_usage()
{
    std::cerr << "Usage: sound (compress|decompress) input_file output_file." << std::endl;
    return 1;
}

std::array<char, 44> read_header(std::ifstream& input_file)
{
    std::array<char, 44> wav_header;

    input_file.read(wav_header.data(), wav_header.size());

    return wav_header;
}

std::list<char> read_first_context(std::istreambuf_iterator<char>& input_iterator)
{
    std::list<char> result;

    result.push_back(*input_iterator);
    ++input_iterator;
    result.push_back(*input_iterator);
    ++input_iterator;

    return result;
}

int fc(const std::list<char>& context)
{
    auto c1 = context.front();
    auto c2 = context.back();

    return c1 + (c1 - c2);
}

int archive(const char *input_name, const char *output_name)
{
    std::ifstream input_file{ input_name, std::ios::binary | std::ios::in };
    input_file.exceptions(std::ios::badbit | std::ios::failbit);

    std::array<char, 44> wav_header = read_header(input_file);

    if (input_file.gcount() != 44)
    {
        // Not enough bytes available to read the header.
        return 2;
    }
    std::istreambuf_iterator<char> wav_iterator(input_file), input_end;
    std::list<char> context = read_first_context(wav_iterator);

    std::ofstream output_file{ output_name, std::ios::binary | std::ios::trunc | std::ios::in };
    output_file.exceptions(std::ios::badbit | std::ios::failbit);

    output_file.write(wav_header.data(), wav_header.size());
    output_file.write(&context.front(), 1);
    output_file.write(&context.back(), 1);

    std::vector<std::uint8_t> encoding;

    while (wav_iterator != input_end)
    {
        int guess = fc(context);
        int error = *wav_iterator - guess;

        // Update the context.
        context.pop_back();
        context.push_back(*wav_iterator);

        encoding.push_back(error >> 8);
        encoding.push_back(error);

        ++wav_iterator;
    }
    compress(encoding, output_file);
    return 0;
}

int unarchive(const char *input_name, const char *output_name)
{
    std::ifstream input_file{ input_name, std::ios::binary | std::ios::in };
    input_file.exceptions(std::ios::badbit | std::ios::failbit);

    std::array<char, 44> wav_header = read_header(input_file);

    if (input_file.gcount() != 44)
    {
        // Not enough bytes available to read the header.
        return 2;
    }
    std::istreambuf_iterator<char> wav_iterator(input_file), input_end;
    std::list<char> context = read_first_context(wav_iterator);
    std::stringstream encoding_stream;
    std::vector<std::uint8_t> encoded_wav{ wav_iterator, input_end };

    decompress(encoded_wav, encoding_stream);

    std::ofstream output_file{ output_name, std::ios::binary | std::ios::trunc | std::ios::in };
    output_file.exceptions(std::ios::badbit | std::ios::failbit);

    output_file.write(wav_header.data(), wav_header.size());
    output_file.write(&context.front(), 1);
    output_file.write(&context.back(), 1);

    std::string encoding = encoding_stream.str();
    auto current_encoding = std::cbegin(encoding);

    while (current_encoding != std::cend(encoding))
    {
        std::uint8_t first_byte = *current_encoding;
        ++current_encoding;
        std::uint8_t second_byte = *current_encoding;
        ++current_encoding;
        int error = (first_byte << 8) | second_byte;

        int guess = fc(context);
        char element = error + guess;

        // Update the context.
        context.pop_back();
        context.push_back(element);

        output_file.write(&element, 1);
    }

    return 0;
}

int main(int argc, char **argv)
{
    if (argc != 4)
    {
        return show_usage();
    }
    if (std::strcmp(argv[1], "compress") == 0)
    {
        return archive(argv[2], argv[3]);
    }
    else if (std::strcmp(argv[1], "decompress") == 0)
    {
        return unarchive(argv[2], argv[3]);
    }
    else
    {
        return show_usage();
    }
    return 0;
}