#pragma once

#include "ImageBitStream.h"
#include "BitStream.h"

#include "FileLogger.h"

#include <iostream>

class PngFilter
{
public:
    enum class FilterType : unsigned char
    {
        NONE,
        SUB,
        UP,
        AVERAGE,
        PAETH
    };

    PngFilter(BitStream* inputStream, BitStream* outputStream)
        : mInputStream(inputStream),
          mOutputStream(outputStream)

    {

    }

    void encode()
    {
        auto image_stream = dynamic_cast<ImageBitStream*>(mInputStream);
        if (!image_stream)
        {
            MLOG_ERROR("Expected ImageStream in PngFilter encode - aborting.");
            return;
        }

        const auto bytes_per_scanline = image_stream->getBytesPerScanline();
        unsigned count{0};
        if (mFilterType == FilterType::NONE)
        {
            while(true)
            {
                if (const auto byte = image_stream->readNextByte())
                {
                    if (count % bytes_per_scanline == 0)
                    {
                        mOutputStream->writeByte(0);
                    }

                    mOutputStream->writeByte(*byte);
                }
                else
                {
                    break;
                }
                count++;
            }
        }
    }

    void decode()
    {
        auto image_stream = dynamic_cast<ImageBitStream*>(mOutputStream);
        if (!image_stream)
        {
            MLOG_ERROR("Expected ImageStream in PngFilter decode - aborting.");
            return;
        }

        const auto bytes_per_scanline = image_stream->getBytesPerScanline();
        unsigned count{0};

        FilterType working_filter_type = FilterType::NONE;
        while(auto byte = mInputStream->readNextByte())
        {
            if (count % bytes_per_scanline == 0)
            {
                working_filter_type = static_cast<FilterType>(*byte);
            }
            else
            {
                if (working_filter_type == FilterType::NONE)
                {
                    image_stream->writeByte(*byte);
                }
                else
                {
                    //std::cout << "Got filter type " << static_cast<int>(working_filter_type) << std::endl;
                }
            }
            count++;
        }
    }

private:

    FilterType mFilterType{FilterType::NONE};
    BitStream* mInputStream{nullptr};
    BitStream* mOutputStream{nullptr};
};
