#include "PngReader.h"

#include "BinaryStream.h"
#include "BitStream.h"
#include "BufferBitStream.h"
#include "ImageBitStream.h"
#include "PngFilter.h"

#include "ZlibEncoder.h"
#include "CyclicRedundancyChecker.h"
#include "StringUtils.h"

#include <iostream>
#include <sstream>

PngReader::~PngReader()
{

}

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

bool PngReader::checkSignature()
{
    const int highBitCheck = 0x89;
    const auto firstPos = mFile->getInHandle()->get();
    if (firstPos != highBitCheck)
    {
        return false;
    }

    std::string fileType;
    BinaryStream::getNextString(mFile->getInHandle(), fileType, 3);
    if (fileType != "PNG")
    {
        return false;
    }

    std::vector<char> sequence{13, 10, 26, 10};
    for (auto c : sequence)
    {
        if (mFile->getInHandle()->get() != c)
        {
            return false;
        }
    }
    return true;
}

bool PngReader::readChunk()
{
    unsigned length = *BinaryStream::getNextDWord(mFile->getInHandle());

    std::string chunkType;
    BinaryStream::getNextString(mFile->getInHandle(), chunkType, 4);

    //std::cout << "Got chunk with type: " << chunkType << " and length: " << length << std::endl;
    bool lastChunk = false;
    if (chunkType == "IHDR")
    {
        if (!readHeaderChunk())
        {
            return false;
        }
    }
    else if(chunkType == "IEND")
    {
        lastChunk = true;
        if (mProcessingDatablocks)
        {
            decodeData();
        }
        unsigned crcCheck = *BinaryStream::getNextDWord(mFile->getInHandle());
        (void) crcCheck;
    }
    else if(chunkType == "IDAT")
    {
        mProcessingDatablocks = true;
        if(!readIDATChunk(length))
        {
            return false;
        }
    }
    else
    {
        if (mProcessingDatablocks)
        {
            decodeData();
        }

        for(unsigned idx=0;idx<length;idx++)
        {
            mFile->getInHandle()->get();
        }

        unsigned crcCheck = *BinaryStream::getNextDWord(mFile->getInHandle());
        (void) crcCheck;
    }
    return !lastChunk;
}

bool PngReader::readHeaderChunk()
{
    auto width = *BinaryStream::getNextDWord(mFile->getInHandle());
    auto height = *BinaryStream::getNextDWord(mFile->getInHandle());
    auto bitDepth = mFile->getInHandle()->get();

    mHeader.setImageData(width, height, bitDepth);

    PngInfo info;
    info.mColorType = static_cast<PngInfo::ColorType>(mFile->getInHandle()->get());
    info.mCompressionMethod = static_cast<PngInfo::CompressionMethod>(mFile->getInHandle()->get());
    info.mFilterMethod = static_cast<PngInfo::FilterMethod>(mFile->getInHandle()->get());
    info.mInterlaceMethod = static_cast<PngInfo::InterlaceMethod>(mFile->getInHandle()->get());
    mHeader.setPngInfo(info);
    mHeader.updateData();

    uint32_t file_crc = *BinaryStream::getNextDWord(mFile->getInHandle());
    auto crc_calc = mHeader.getCrc();

    //std::cout << mHeader.toString() << "*************\n";
    if (file_crc != crc_calc)
    {
        //std::cout << "Header crc calc does not match file value. File: " << file_crc << " Header: " << crc_calc << std::endl;;
        return false;
    }
    else
    {
        return true;
    }
}

bool PngReader::readIDATChunk(unsigned length)
{
    auto crc_check = std::make_unique<CyclicRedundancyChecker>();

    std::vector<unsigned char> char_data = StringUtils::toBytes("IDAT");
    for (auto c : char_data)
    {
        crc_check->addValue(c);
    }

    mInputStream->setChecksumCalculator(crc_check.get());
    for(unsigned idx=0; idx<length; idx++)
    {
        mInputStream->writeByte(*mFile->readNextByte());
    }
    mInputStream->clearChecksumCalculator();

    uint32_t file_crc = *BinaryStream::getNextDWord(mFile->getInHandle());
    auto crc_calc = crc_check->getChecksum();

    if (file_crc != crc_calc)
    {
        //std::cout << "IDAT crc calc does not match file value. File: " << file_crc << " Header: " << crc_calc << std::endl;;
        return false;
    }
    else
    {
        return true;
    }
}

std::unique_ptr<Image> PngReader::read()
{
    auto image = std::make_unique<Image>(5, 5);

    mFile = std::make_unique<File>(mPath);
    mFile->open(File::AccessMode::Read);

    if (!checkSignature())
    {
        //std::cout << "Signature check failed" << std::endl;
        return image;
    }

    mInputStream = std::make_unique<BufferBitStream>();
    mOutputStream = std::make_unique<BufferBitStream>();
    mEncoder = std::make_unique<ZlibEncoder>(mInputStream.get(), mOutputStream.get());

    while(readChunk())
    {

    }

    image->setWidth(mHeader.getWidth());
    image->setHeight(mHeader.getHeight());
    image->setBitDepth(mHeader.getBitDepth());
    image->setNumChannels(mHeader.getPngInfo().getNumChannels());

    auto image_bit_stream = std::make_unique<ImageBitStream>(image.get());
    PngFilter filter(mOutputStream.get(), image_bit_stream.get());
    filter.decode();

    return image;
}

void PngReader::decodeData()
{
    mEncoder->decode();
    //std::cout << mEncoder->getData() << "***********" << std::endl;
}
