/*
 * Copyright 2013, Gerasim Troeglazov, 3dEyes@gmail.com. All rights reserved.
 * Distributed under the terms of the MIT License.
 */ 


#include <File.h>

#include "BaseTranslator.h"
#include "PSDWriter.h"
#include "DataArray.h"


PSDWriter::PSDWriter(BPositionIO *stream)
{       
        fAlphaChannel = -1;
        fStream = stream;
        fReady = false;

        TranslatorBitmap header;
        stream->Seek(0, SEEK_SET);
        status_t err = stream->Read(&header, sizeof(TranslatorBitmap));
        if (err < B_OK)
                return;
        else if (err < (int)sizeof(TranslatorBitmap))
                return;
                
        fBitmapDataPos = stream->Position();

        BRect bounds;
        bounds.left = B_BENDIAN_TO_HOST_FLOAT(header.bounds.left);
        bounds.top = B_BENDIAN_TO_HOST_FLOAT(header.bounds.top);
        bounds.right = B_BENDIAN_TO_HOST_FLOAT(header.bounds.right);
        bounds.bottom = B_BENDIAN_TO_HOST_FLOAT(header.bounds.bottom);
        fInRowBytes = B_BENDIAN_TO_HOST_INT32(header.rowBytes);
        fColorSpace = (color_space)B_BENDIAN_TO_HOST_INT32(header.colors);
        
        switch (fColorSpace) {
                case B_GRAY8:
                case B_CMAP8:
                        fChannels = 1;
                        break;
                case B_RGBA32:
                        fChannels = 4;
                        fAlphaChannel = 3;
                        break;
                case B_RGB32:
                        fChannels = 3;
                        break;
                default:
                        return;
        };

        fWidth = bounds.IntegerWidth() + 1;
        fHeight = bounds.IntegerHeight() + 1;
        
        fVersion = PSD_FILE;
        fCompression = PSD_COMPRESSED_RAW;

        fReady = true;
}


PSDWriter::~PSDWriter()
{
}


bool
PSDWriter::IsReady(void)
{
        return fReady;
}


void
PSDWriter::SetCompression(int16 compression)
{
        fCompression = compression;
}


void
PSDWriter::SetVersion(int16 ver)
{
        fVersion = ver;
}


status_t
PSDWriter::Encode(BPositionIO *target)
{
        if (!fReady)
                return B_NO_TRANSLATOR;

        status_t status = _LoadChannelsFromRGBA32();
        if (status != B_OK)
                return status;

        // PSD header
        BDataArray psdHeader(64);
        psdHeader << "8BPS"; // Signature
        psdHeader << (uint16)fVersion; // Version
        psdHeader.Repeat(0, 6); // Reserved
        psdHeader << fChannels; // Channels
        psdHeader << fHeight << fWidth; // Image size
        psdHeader << (int16)8; // Depth
        psdHeader << (int16)PSD_COLOR_MODE_RGB; // ColorMode

        // Color mode section
        BDataArray psdColorModeSection(16);
        psdColorModeSection << (uint32)0;

        // Image resource section
        BDataArray psdImageResourceSection(64);
        psdImageResourceSection << "8BIM"; // Block signature
        psdImageResourceSection << (uint16)1005;
        psdImageResourceSection << (uint16)0;
        psdImageResourceSection << (uint32)16;
        uint8 resBlock[16] = {0x00, 0x48, 0x00, 0x00,
                0x00, 0x01, 0x00, 0x01,
                0x00, 0x48, 0x00, 0x00,
                0x00, 0x01, 0x00, 0x01};
        psdImageResourceSection.Append(resBlock, 16);   
        // Current layer info
        psdImageResourceSection << "8BIM"; // Block signature
        psdImageResourceSection << (uint16)1024;
        psdImageResourceSection << (uint16)0;
        psdImageResourceSection << (uint32)2;
        psdImageResourceSection << (uint16)0; // Set current layer to 0

        // Layer & mask section
        BDataArray psdLayersSection;
        psdLayersSection << (uint16)1; // Layers count
        psdLayersSection << (uint32)0; // Layer rect
        psdLayersSection << (uint32)0;
        psdLayersSection << (uint32)fHeight;
        psdLayersSection << (uint32)fWidth;     
        psdLayersSection << (uint16)fChannels;
        
        for (int channelIdx = 0; channelIdx < fChannels; channelIdx++) {
                if (channelIdx == 3)
                        psdLayersSection << (int16)-1; // Alpha channel id (-1)
                else
                        psdLayersSection << (int16)channelIdx; // Channel num
                        
                if (fCompression == PSD_COMPRESSED_RAW) {
                        if (fVersion == PSD_FILE) {
                                        psdLayersSection << (uint32)(psdChannel[channelIdx].Length()
                                                + sizeof(int16));
                        } else {
                                        psdLayersSection << (uint64)(psdChannel[channelIdx].Length()
                                                + sizeof(int16));
                        }
                } else {
                        if (fVersion == PSD_FILE) {
                                psdLayersSection << (uint32)(psdChannel[channelIdx].Length()
                                        + psdByteCounts[channelIdx].Length() + sizeof(int16));
                        } else {
                                psdLayersSection << (uint64)(psdChannel[channelIdx].Length()
                                        + psdByteCounts[channelIdx].Length() + sizeof(int16));
                        }
                }
        }

        psdLayersSection << "8BIM";
        psdLayersSection << "norm"; // Blend mode = norm
        psdLayersSection << (uint8)255; // Opacity
        psdLayersSection << (uint8)0; // Clipping
        psdLayersSection << (uint8)1; // Flags
        psdLayersSection << (uint8)0; // Flags
        psdLayersSection << (uint32)24; // Extra data length
        psdLayersSection << (uint32)0; // Mask info
        psdLayersSection << (uint32)0;

        psdLayersSection << (uint8)15; // Layer name length
        uint8 layerName[16] = {"Layer #1       "};
        psdLayersSection.Append(layerName, 15); // Layer name

        if (fCompression == PSD_COMPRESSED_RAW) {
                for (int channelIdx = 0; channelIdx < fChannels; channelIdx++) {
                        psdLayersSection << fCompression; // Compression mode
                        psdLayersSection.Append(psdChannel[channelIdx]); // Channel data
                }
        } else {        
                for (int channelIdx = 0; channelIdx < fChannels; channelIdx++) {
                        psdLayersSection << fCompression; // Compression mode
                        psdLayersSection.Append(psdByteCounts[channelIdx]); // Bytes count
                        psdLayersSection.Append(psdChannel[channelIdx]); // Channel data
                }
        }

        if (fCompression == PSD_COMPRESSED_RLE
                && psdLayersSection.Length() % 2 != 0) {
                psdLayersSection << (uint8)0;
        }

        psdHeader.WriteToStream(target);
        
        psdColorModeSection.WriteToStream(target);
        
        _WriteUInt32ToStream(target, psdImageResourceSection.Length());
        psdImageResourceSection.WriteToStream(target);

        if (fVersion == PSD_FILE) {
                _WriteUInt32ToStream(target, psdLayersSection.Length() + sizeof(int32));
                _WriteUInt32ToStream(target, psdLayersSection.Length());
        } else {
                _WriteUInt64ToStream(target, psdLayersSection.Length() + sizeof(int64));
                _WriteUInt64ToStream(target, psdLayersSection.Length());                
        }
        psdLayersSection.WriteToStream(target);

        // Merged layer
        _WriteUInt16ToStream(target, fCompression); // Compression mode
        
        if (fCompression == PSD_COMPRESSED_RLE) {
                for (int channelIdx = 0; channelIdx < fChannels; channelIdx++)
                        psdByteCounts[channelIdx].WriteToStream(target);
        }

        for (int channelIdx = 0; channelIdx < fChannels; channelIdx++)
                psdChannel[channelIdx].WriteToStream(target);

        return B_OK;
}


BDataArray*
PSDWriter::_PackBits(uint8 *buff, int32  len)
{
        BDataArray *packedBits = new BDataArray();

        int32  count = len;  
        len = 0;

        while (count > 0) {
                int i;
                for (i = 0; (i < 128) && (buff[0] == buff[i]) && (count - i > 0); i++);
                if (i < 2) {
                        for (i = 0; i < 128; i++) {
                                bool b1 = buff[i] != buff[i + 1];
                                bool b3 = buff[i] != buff[i + 2];
                                bool b2 = count - (i + 2) < 1;
                                if (count - (i + 1) <= 0)
                                        break;
                                if (!(b1 || b2 || b3))
                                        break;
                        }

                        if (count == 1)
                                i = 1;

                        if (i > 0) {
                                packedBits->Append((uint8)(i - 1));
                                for (int j = 0; j < i; j++)
                                        packedBits->Append((uint8)buff[j]);
                                buff += i;
                                count -= i;
                                len += (i + 1);
                        }
                } else {
                        packedBits->Append((uint8)(-(i - 1)));
                        packedBits->Append((uint8)(*buff));
                        buff += i;
                        count -= i;
                        len += 2;
                }
        }
        return packedBits;
}


status_t
PSDWriter::_LoadChannelsFromRGBA32(void)
{
        if (fColorSpace != B_RGB32 && fColorSpace != B_RGBA32)
                return B_NO_TRANSLATOR;

        int32 channelSize = fWidth * fHeight;

        fStream->Seek(fBitmapDataPos, SEEK_SET);

        if (fCompression == PSD_COMPRESSED_RAW) {
                for (int i = 0; i < channelSize; i++) {
                        uint8 rgba[4];
                        fStream->Read(rgba, sizeof(uint32));
                        psdChannel[0].Append((uint8)rgba[2]); // Red channel
                        psdChannel[1].Append((uint8)rgba[1]); // Green channel
                        psdChannel[2].Append((uint8)rgba[0]); // Blue channel
                        if (fChannels == 4)
                                psdChannel[3].Append((uint8)rgba[3]); // Alpha channel
                }
                return B_OK;
        } else if (fCompression == PSD_COMPRESSED_RLE) {
                for (int32 h = 0; h < fHeight; h++) {
                        BDataArray lineData[4];
                        
                        for (int32 w = 0; w < fWidth; w++) {
                                uint8 rgba[4];
                                fStream->Read(rgba, sizeof(uint32));
                                lineData[0].Append((uint8)rgba[2]); // Red channel
                                lineData[1].Append((uint8)rgba[1]); // Green channel
                                lineData[2].Append((uint8)rgba[0]); // Blue channel
                                if (fChannels == 4)
                                        lineData[3].Append((uint8)rgba[3]); // Alpha channel
                        }
                        
                        for (int channelIdx = 0; channelIdx < fChannels; channelIdx++) {
                                BDataArray *packedLine = _PackBits(lineData[channelIdx].Buffer(),
                                        lineData[channelIdx].Length());

                                if (fVersion == PSD_FILE)
                                        psdByteCounts[channelIdx].Append((uint16)packedLine->Length());
                                else
                                        psdByteCounts[channelIdx].Append((uint32)packedLine->Length());

                                psdChannel[channelIdx].Append(*packedLine);
                                delete packedLine;
                        }
                }
                return B_OK;
        }
        return B_NO_TRANSLATOR;
}


void
PSDWriter::_WriteInt64ToStream(BPositionIO *stream, int64 val)
{
        val = B_HOST_TO_BENDIAN_INT64(val);
        stream->Write(&val, sizeof(int32));
}


void
PSDWriter::_WriteUInt64ToStream(BPositionIO *stream, uint64 val)
{
        val = B_HOST_TO_BENDIAN_INT64(val);
        stream->Write(&val, sizeof(uint64));
}


void
PSDWriter::_WriteInt32ToStream(BPositionIO *stream, int32 val)
{
        val = B_HOST_TO_BENDIAN_INT32(val);
        stream->Write(&val, sizeof(int32));
}


void
PSDWriter::_WriteUInt32ToStream(BPositionIO *stream, uint32 val)
{
        val = B_HOST_TO_BENDIAN_INT32(val);
        stream->Write(&val, sizeof(uint32));
}


void
PSDWriter::_WriteInt16ToStream(BPositionIO *stream, int16 val)
{
        val = B_HOST_TO_BENDIAN_INT16(val);
        stream->Write(&val, sizeof(int16));
}


void
PSDWriter::_WriteUInt16ToStream(BPositionIO *stream, uint16 val)
{
        val = B_HOST_TO_BENDIAN_INT16(val);
        stream->Write(&val, sizeof(int16));
}


void
PSDWriter::_WriteInt8ToStream(BPositionIO *stream, int8 val)
{
        stream->Write(&val, sizeof(int8));
}


void
PSDWriter::_WriteUInt8ToStream(BPositionIO *stream, uint8 val)
{
        stream->Write(&val, sizeof(uint8));
}


void
PSDWriter::_WriteFillBlockToStream(BPositionIO *stream,
        uint8 val, size_t count)
{
        for (size_t i = 0; i < count; i++)
                stream->Write(&val, sizeof(uint8));
}


void
PSDWriter::_WriteBlockToStream(BPositionIO *stream,
        uint8 *val, size_t count)
{
        stream->Write(val, count);
}