#include <exception>
#include "grey.hpp"
#include <arpa/inet.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

static void push_in_network_order(std::uint32_t input, std::vector<std::uint8_t>& output)
{
    input = htonl(input);
    auto pointer_range = reinterpret_cast<std::uint8_t *>(&input);
    std::copy(pointer_range, pointer_range + 4, std::back_inserter(output));
}

static std::uint32_t pop_in_native_order(std::string::const_iterator& input)
{
    std::uint32_t in_native = ntohl(*reinterpret_cast<const std::uint32_t *>(&*input));
    input += 4;

    return in_native;
}

std::vector<std::uint8_t> to_rle(const std::filesystem::path& input)
{
    int x, y, n;
    unsigned char *data = stbi_load(input.native().data(), &x, &y, &n, 1);

    if (data == nullptr)
    {
        throw std::runtime_error(stbi_failure_reason());
    }
    unsigned char *original_data = data;
    std::uint8_t current_color = 255;
    std::uint32_t current_count = 0;
    std::vector<std::uint8_t> rle_stream;

    push_in_network_order(x, rle_stream);
    push_in_network_order(y, rle_stream);
    push_in_network_order(1, rle_stream);

    for (std::size_t i = 0; i < x; ++i)
    {
        for (std::size_t j = 0; j < y; ++j)
        {
            if (current_color == *data)
            {
                ++current_count;
            }
            else
            {
                current_color = current_color == 255 ? 0 : 255;
                push_in_network_order(current_count, rle_stream);
                current_count = 1;
            }
            ++data;
        }
    }
    stbi_image_free(original_data);
    if (current_count != 0)
    {
        push_in_network_order(current_count, rle_stream);
    }
    return rle_stream;
}

void from_rle(const std::string& input, const std::filesystem::path& output)
{
    std::vector<std::uint8_t> buffer;
    auto input_position = std::cbegin(input);

    int x = pop_in_native_order(input_position);
    int y = pop_in_native_order(input_position);
    int comp = pop_in_native_order(input_position);

    std::uint8_t current_color = 255;
    std::uint32_t current_count = 0;

    while (input_position != std::cend(input))
    {
        auto in_native = pop_in_native_order(input_position);

        for (; in_native > 0; --in_native)
        {
            buffer.push_back(current_color);
        }
        current_color = current_color == 255 ? 0 : 255;
    }
    stbi_write_jpg(output.native().data(), x, y, comp, buffer.data(), 100);
}