#include "BasicPngWriter.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>

BasicPngWriter::BasicPngWriter()
    : IImageWriter(ImgFormat::PNG)
{

}

BasicPngWriter::~BasicPngWriter()
{

}

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

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

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

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

void BasicPngWriter::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 BasicPngWriter::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 BasicPngWriter::writeDataChunks(const BufferBitStream& buffer)
{
    auto num_bytes = buffer.getBuffer().size();
    std::size_t max_bytes{ 32000 };

    auto num_dat_chunks = num_bytes / max_bytes + 1;
    std::size_t 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(static_cast<uint32_t>(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 BasicPngWriter::write(const Path& path, Image* image)
{
    std::unique_ptr<File> out_file;

    if (path.empty())
    {
        out_file = std::make_unique<File>(path);
        out_file->open(File::AccessMode::Write);
        mOutStream = std::make_unique<OutputBitStream>(out_file->getOutHandle());
    }
    else
    {
        mOutStream = std::make_unique<BufferBitStream>();
    }

    mWorkingImage = image;

    auto image_bit_stream = std::make_unique<ImageBitStream>(image);
    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();
}