#include "PngWriter.h"

#include "Image.h"
#include "File.h"
#include "BufferBitStream.h"
#include "OutputBitStream.h"
#include "ImageBitStream.h"
#include "StringUtils.h"

#include "PngFilter.h"
#include "Lz77Encoder.h"
#include "ZlibEncoder.h"
#include "HuffmanEncoder.h"
#include "CyclicRedundancyChecker.h"

#include "ByteUtils.h"

#include <iostream>

PngWriter::PngWriter()
{

}

PngWriter::~PngWriter()
{

}

std::unique_ptr<PngWriter> PngWriter::Create()
{
    return std::make_unique<PngWriter>();
}

void PngWriter::setCompressionMethod(Deflate::CompressionMethod method)
{
    mCompressionMethod = method;
}

void PngWriter::setPath(const Path& path)
{
    mPath = path;
}

void PngWriter::setPngInfo(const PngInfo& info)
{
    mPngInfoUserSet = true;
    mPngInfo = info;
}

void PngWriter::writeSignature()
{
    mOutStream->writeByte(mPngHeader.getHighBitCheck());
    mOutStream->writeBytes(StringUtils::toBytes(mPngHeader.getFileName()));
    mOutStream->writeBytes(mPngHeader.getSignature());
}

void PngWriter::writeHeader()
{
    writeSignature();

    if (!mPngInfoUserSet)
    {
        if (mWorkingImage->getNumChannels() == 1)
        {
            mPngInfo.mColorType = PngInfo::ColorType::GREYSCALE;
        }
        else if (mWorkingImage->getNumChannels() == 2)
        {
            mPngInfo.mColorType = PngInfo::ColorType::GREYSCALE_ALPHA;
        }
        else if (mWorkingImage->getNumChannels() == 3)
        {
            mPngInfo.mColorType = PngInfo::ColorType::RGB;
        }
        else if (mWorkingImage->getNumChannels() == 4)
        {
            mPngInfo.mColorType = PngInfo::ColorType::RGB_ALPHA;
        }
    }

    mPngHeader.setPngInfo(mPngInfo);
    mPngHeader.setImageData(mWorkingImage->getWidth(), mWorkingImage->getHeight(), mWorkingImage->getBitDepth());

    auto length = mPngHeader.getLength();
    mOutStream->write(length);

    mOutStream->writeBytes(StringUtils::toBytes(mPngHeader.getChunkName()));

    mPngHeader.updateData();
    mOutStream->writeBytes(mPngHeader.getData());

    auto crc = mPngHeader.getCrc();

    //std::cout << mPngHeader.toString() << "*********" << std::endl;
    mOutStream->write(crc);
}

void PngWriter::writeEndChunk()
{
    //std::cout << "Start writing end chunk" << std::endl;
    unsigned length{0};
    mOutStream->write(length);

    mOutStream->writeBytes(StringUtils::toBytes("IEND"));
    std::vector<unsigned char> char_data = StringUtils::toBytes("IEND");

    CyclicRedundancyChecker crc_check;
    for (auto c : char_data)
    {
        crc_check.addValue(c);
    }
    auto crc = crc_check.getChecksum();
    mOutStream->write(crc);

    //std::cout << "Writing end chunk" << std::endl;
}

void PngWriter::writeDataChunks(const BufferBitStream& buffer)
{
    auto num_bytes = buffer.getBuffer().size();
    auto max_bytes{32000};

    auto num_dat_chunks = num_bytes/max_bytes + 1;
    unsigned offset = 0;
    for(std::size_t idx=0;idx<num_dat_chunks;idx++)
    {
        auto length = max_bytes;
        if (idx == num_dat_chunks - 1)
        {
            length = num_bytes - num_dat_chunks*num_bytes;
        }

        //std::cout << "Writing idat length " << num_bytes << std::endl;
        mOutStream->write(num_bytes);

        std::vector<unsigned char> char_data = StringUtils::toBytes("IDAT");
        mOutStream->writeBytes(char_data);

        CyclicRedundancyChecker crc_check;
        for (auto c : char_data)
        {
            crc_check.addValue(c);
        }

        for(unsigned jdx=0; jdx<num_bytes; jdx++)
        {
            auto val = buffer.getBuffer()[idx*max_bytes + jdx];
            mOutStream->writeByte(val);
            crc_check.addValue(val);
        }

        auto crc = crc_check.getChecksum();
        //std::cout << "Writing idat crc" << crc << std::endl;
        mOutStream->write(crc);
        //std::cout << "Finished Writing idat crc" << crc << std::endl;
    }
}

void PngWriter::write(const std::unique_ptr<Image<unsigned char> >& image)
{
    if (!mPath.empty())
    {
        mWorkingFile = std::make_unique<File>(mPath);
        mWorkingFile->open(File::AccessMode::Write);
        mOutStream = std::make_unique<OutputBitStream>(mWorkingFile->getOutHandle());
    }
    else
    {
        mOutStream = std::make_unique<BufferBitStream>();
    }

    mWorkingImage = image.get();

    auto image_bit_stream = std::make_unique<ImageBitStream>(image.get());
    auto raw_image_stream = image_bit_stream.get();

    mInStream = std::move(image_bit_stream);

    writeHeader();

    auto filter_out_stream = std::make_unique<BufferBitStream>();
    PngFilter filter(raw_image_stream, filter_out_stream.get());
    filter.encode();

    filter_out_stream->resetOffsets();

    std::unique_ptr<BufferBitStream> lz77_out_stream;

    if (mCompressionMethod == Deflate::CompressionMethod::NONE)
    {
        lz77_out_stream = std::move(filter_out_stream);
    }
    else
    {
        lz77_out_stream = std::make_unique<BufferBitStream>();
        Lz77Encoder lz77_encoder(filter_out_stream.get(), lz77_out_stream.get());

        if (mCompressionMethod == Deflate::CompressionMethod::FIXED_HUFFMAN)
        {
            auto huffman_encoder = std::make_unique<HuffmanEncoder>();
            huffman_encoder->setUseFixedCode(true);
            lz77_encoder.setPrefixCodeGenerator(std::move(huffman_encoder));
        }
        lz77_encoder.encode();
        lz77_out_stream->resetOffsets();
    }

    BufferBitStream zlib_out_stream;
    ZlibEncoder zlib_encoder(lz77_out_stream.get(), &zlib_out_stream);
    zlib_encoder.setDeflateCompressionMethod(mCompressionMethod);
    zlib_encoder.encode();
    zlib_out_stream.resetOffsets();

    writeDataChunks(zlib_out_stream);

    writeEndChunk();

    if (mWorkingFile)
    {
        mWorkingFile->close();
    }
}