// 
// (C) Jan de Vaan 2007-2010, all rights reserved. See the accompanying "License.txt" for licensed use. 
// 

#include "config.h"
#include "util.h"
#include "header.h"
#include "jpegmarker.h"
#include "jpegstreamwriter.h"
#include "jpegimagedatasegment.h"
#include "decoderstrategy.h"
#include "encoderstrategy.h"
#include <memory>


// JFIF\0
BYTE jfifID[] = {'J','F','I','F','\0'};


LONG CLAMP(LONG i, LONG j, LONG MAXVAL)
{
    if (i > MAXVAL || i < j)
        return j;

    return i;
}


JlsCustomParameters ComputeDefault(LONG MAXVAL, LONG NEAR)
{
    JlsCustomParameters preset = JlsCustomParameters();

    LONG FACTOR = (MIN(MAXVAL, 4095) + 128)/256;

    preset.T1 = CLAMP(FACTOR * (BASIC_T1 - 2) + 2 + 3*NEAR, NEAR + 1, MAXVAL);
    preset.T2 = CLAMP(FACTOR * (BASIC_T2 - 3) + 3 + 5*NEAR, preset.T1, MAXVAL);
    preset.T3 = CLAMP(FACTOR * (BASIC_T3 - 4) + 4 + 7*NEAR, preset.T2, MAXVAL);
    preset.MAXVAL = MAXVAL;
    preset.RESET = BASIC_RESET;
    return preset;
}


JLS_ERROR CheckParameterCoherent(const JlsParameters* pparams)
{
    if (pparams->bitspersample < 2 || pparams->bitspersample > 16)
        return ParameterValueNotSupported;

    if (pparams->ilv < 0 || pparams->ilv > 2)
        return InvalidCompressedData;

    switch (pparams->components)
    {
        case 4: return pparams->ilv == ILV_SAMPLE ? ParameterValueNotSupported : OK; 
        case 3: return OK;
        case 1: return pparams->ilv != ILV_NONE ? ParameterValueNotSupported : OK;
        case 0: return InvalidJlsParameters;

        default: return pparams->ilv != ILV_NONE ? ParameterValueNotSupported : OK; 
    }
}


void JpegImageDataSegment::Serialize(JpegStreamWriter& streamWriter)
{
    JlsParameters info = _info;
    info.components = _ccompScan;
    std::auto_ptr<EncoderStrategy> qcodec = JlsCodecFactory<EncoderStrategy>().GetCodec(info, _info.custom);
    ProcessLine* processLine = qcodec->CreateProcess(_rawStreamInfo);
    ByteStreamInfo compressedData = streamWriter.OutputStream();
    size_t cbyteWritten = qcodec->EncodeScan(std::auto_ptr<ProcessLine>(processLine), &compressedData, streamWriter._bCompare ? streamWriter.GetPos() : NULL);
    streamWriter.Seek(cbyteWritten);
}


int ReadScanHeader(BYTE* compressedBytes)
{
    BYTE rgbyte[20];
    size_t readBytes = 0;
    ::memcpy(rgbyte, compressedBytes, 4);
    readBytes += 4;

    size_t cbyteScanheader = rgbyte[3] - 2;

    if (cbyteScanheader > sizeof(rgbyte))
        throw JlsException(InvalidCompressedData);

    ::memcpy(rgbyte, compressedBytes, cbyteScanheader);
    readBytes += cbyteScanheader;
    return (int) readBytes;
}


void Assert(bool valid)
{
    if (!valid)
        throw JlsException(InvalidCompressedData);
}


JpegMarkerReader::JpegMarkerReader(ByteStreamInfo byteStreamInfo) :
        _byteStream(byteStreamInfo),
        _bCompare(false),
        _info(),
        _rect()
{
}


void JpegMarkerReader::Read(ByteStreamInfo rawPixels)
{
    ReadHeader();

    JLS_ERROR error = CheckParameterCoherent(&_info);
    if (error != OK)
        throw JlsException(error);

    if (_rect.Width <= 0)
    {
        _rect.Width = _info.width;
        _rect.Height = _info.height;
    }

    int64_t bytesPerPlane = (int64_t)(_rect.Width) * _rect.Height * ((_info.bitspersample + 7)/8);

    if (rawPixels.rawData != NULL && int64_t(rawPixels.count) < bytesPerPlane * _info.components)
        throw JlsException(UncompressedBufferTooSmall);

    int componentIndex = 0;
    
    while (componentIndex < _info.components)
    {
        ReadStartOfScan(componentIndex == 0);

        std::auto_ptr<DecoderStrategy> qcodec = JlsCodecFactory<DecoderStrategy>().GetCodec(_info, _info.custom);   
        ProcessLine* processLine = qcodec->CreateProcess(rawPixels);
        qcodec->DecodeScan(std::auto_ptr<ProcessLine>(processLine), _rect, &_byteStream, _bCompare); 
        SkipBytes(&rawPixels, (size_t)bytesPerPlane);       

        if (_info.ilv != ILV_NONE)
            return;

        componentIndex += 1;
    }
}


void JpegMarkerReader::ReadNBytes(std::vector<char>& dst, int byteCount)
{
    for (int i = 0; i < byteCount; ++i)
    {
        dst.push_back((char)ReadByte());
    }
}


void JpegMarkerReader::ReadHeader()
{
    if (ReadByte() != 0xFF)
        throw JlsException(MissingJpegMarkerStart);

    if (ReadByte() != JPEG_SOI)
        throw JlsException(InvalidCompressedData);

    for (;;)
    {
        if (ReadByte() != 0xFF)
            throw JlsException(MissingJpegMarkerStart);

        BYTE marker = (BYTE)ReadByte();

        if (marker == JPEG_SOS)
            return;

        LONG cbyteMarker = ReadWord();

        int bytesRead = ReadMarker(marker) + 2;

        int paddingToRead = cbyteMarker - bytesRead;

        if (paddingToRead < 0)
            throw JlsException(InvalidCompressedData);

        for (int i = 0; i < paddingToRead; ++i)
        {
            ReadByte();
        }
    }
}


int JpegMarkerReader::ReadMarker(BYTE marker)
{
    switch (marker)
    {
        case JPEG_SOF_55: return ReadStartOfFrame();
        case JPEG_COM: return ReadComment();
        case JPEG_LSE: return ReadPresetParameters();
        case JPEG_APP0: return 0;
        case JPEG_APP7: return ReadColorSpace();
        case JPEG_APP8: return ReadColorXForm();
        case JPEG_SOF_0:
        case JPEG_SOF_1:
        case JPEG_SOF_2:
        case JPEG_SOF_3:
        case JPEG_SOF_5:
        case JPEG_SOF_6:
        case JPEG_SOF_7:
        case JPEG_SOF_9:
        case JPEG_SOF_10:
        case JPEG_SOF_11:
            throw JlsException(UnsupportedEncoding);

        // Other tags not supported (among which DNL DRI)
        default: throw JlsException(UnknownJpegMarker);
    }
}


int JpegMarkerReader::ReadPresetParameters()
{
    LONG type = ReadByte();

    switch (type)
    {
    case 1:
        {
            _info.custom.MAXVAL = ReadWord();
            _info.custom.T1 = ReadWord();
            _info.custom.T2 = ReadWord();
            _info.custom.T3 = ReadWord();
            _info.custom.RESET = ReadWord();
            return 11;
        }
    }

    return 1;
}


void JpegMarkerReader::ReadStartOfScan(bool firstComponent)
{
    if (!firstComponent)
    {
        Assert(ReadByte() == 0xFF);
        Assert(ReadByte() == JPEG_SOS);
    }
    int length = ReadByte(); //length
    length = length * 256 + ReadByte();
    
    LONG componentCount = ReadByte();
    if (componentCount != 1 && componentCount != _info.components)
        throw JlsException(ParameterValueNotSupported);

    for (LONG i = 0; i < componentCount; ++i)
    {
        ReadByte();
        ReadByte();
    }
    _info.allowedlossyerror = ReadByte();
    _info.ilv = interleavemode(ReadByte());
    Assert(ILV_NONE <= _info.ilv && _info.ilv <= ILV_SAMPLE);
    Assert(ReadByte() == 0);

    if(_info.bytesperline == 0)
    {
        int width = _rect.Width != 0 ? _rect.Width : _info.width;
        int components = _info.ilv == ILV_NONE ? 1 : _info.components;
        _info.bytesperline = components * width * ((_info.bitspersample + 7)/8);
    }
}


int JpegMarkerReader::ReadComment()
{
    return 0;
}


void JpegMarkerReader::ReadJfif()
{
    for(int i = 0; i < (int)sizeof(jfifID); i++)
    {
        if(jfifID[i] != ReadByte())
            return;
    }
    _info.jfif.Ver   = ReadWord();

    // DPI or DPcm
    _info.jfif.units = ReadByte();
    _info.jfif.XDensity = ReadWord();
    _info.jfif.YDensity = ReadWord();

    // thumbnail
    _info.jfif.Xthumb = ReadByte();
    _info.jfif.Ythumb = ReadByte();
    if(_info.jfif.Xthumb > 0 && _info.jfif.pdataThumbnail) 
    {
        std::vector<char> tempbuff((char*)_info.jfif.pdataThumbnail, (char*)_info.jfif.pdataThumbnail+3*_info.jfif.Xthumb*_info.jfif.Ythumb);
        ReadNBytes(tempbuff, 3*_info.jfif.Xthumb*_info.jfif.Ythumb);
    }
}


int JpegMarkerReader::ReadStartOfFrame()
{
    _info.bitspersample = ReadByte();
    int cline = ReadWord();
    int ccol = ReadWord();
    _info.width = ccol;
    _info.height = cline;
    _info.components= ReadByte();
    return 6;
}


BYTE JpegMarkerReader::ReadByte()
{
    if (_byteStream.rawStream != NULL)
        return (BYTE)_byteStream.rawStream->sbumpc();

    if (_byteStream.count <= 0)
        throw JlsException(InvalidCompressedData);

    BYTE value = _byteStream.rawData[0]; 
    
    SkipBytes(&_byteStream, 1);

    return value;
}


int JpegMarkerReader::ReadWord()
{
    int i = ReadByte() * 256;
    return i + ReadByte();
}


int JpegMarkerReader::ReadColorSpace()
{
    return 0;
}


int JpegMarkerReader::ReadColorXForm()
{
    std::vector<char> sourceTag;
    ReadNBytes(sourceTag, 4);

    if(strncmp(&sourceTag[0],"mrfx", 4) != 0)
        return 4;

    int xform = ReadByte();
    switch(xform) 
    {
        case COLORXFORM_NONE:
        case COLORXFORM_HP1:
        case COLORXFORM_HP2:
        case COLORXFORM_HP3:
            _info.colorTransform = xform;
            return 5;
        case COLORXFORM_RGB_AS_YUV_LOSSY:
        case COLORXFORM_MATRIX:
            throw JlsException(ImageTypeNotSupported);
        default:
            throw JlsException(InvalidCompressedData);
    }
}


ByteStreamInfo FromStream(std::basic_streambuf<char>* stream)
{
    ByteStreamInfo info = ByteStreamInfo();
    info.rawStream = stream;    
    return info;
}


void SkipBytes(ByteStreamInfo* streamInfo, size_t count)
{
    if (streamInfo->rawData == NULL)
        return;

    streamInfo->rawData += count;
    streamInfo->count -= count;
}
