//lz11.h----------------------------------------------------------------------------------
#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <stdlib.h>
#include <cstdio>
#include <squish.h>
#ifdef _WIN32
    #include <windows.h>
#endif
#include "FreeImage.h"
#include <list>
#include <cmath>
#include <cstring>
using namespace std;

typedef uint8_t byte;
typedef uint32_t uint;


class Stream
{
protected:
    int curReadLoc;

public:
    Stream();

    byte ReadByte();
    void Read(byte* out, int unused, int amt);

    void WriteByte(byte b);

    vector<byte> data;
};

long LZ11Decompress(Stream& instream, long inLength, Stream& outstream);

//lz11.cpp--------------------------------------------------------------------
int ToNDSs32(byte* buffer, int offset)
        {
            return (int)(buffer[offset]
                        | (buffer[offset + 1] << 8)
                        | (buffer[offset + 2] << 16)
                        | (buffer[offset + 3] << 24));
        }

        int ToNDSu24(byte* buffer, int offset)
        {
            return (int)(buffer[offset]
                        | (buffer[offset + 1] << 8)
                        | (buffer[offset + 2] << 16));
        }


        long LZ11Decompress(Stream& instream, long inLength, Stream& outstream)
        {
            ////#region Format definition in NDSTEK style
            /*  Data header (32bit)
                  Bit 0-3   Reserved
                  Bit 4-7   Compressed type (must be 1 for LZ77)
                  Bit 8-31  Size of decompressed data. if 0, the next 4 bytes are decompressed length
                Repeat below. Each Flag Byte followed by eight Blocks.
                Flag data (8bit)
                  Bit 0-7   Type Flags for next 8 Blocks, MSB first
                Block Type 0 - Uncompressed - Copy 1 Byte from Source to Dest
                  Bit 0-7   One data byte to be copied to dest
                Block Type 1 - Compressed - Copy LEN Bytes from Dest-Disp-1 to Dest
                    If Reserved is 0: - Default
                      Bit 0-3   Disp MSBs
                      Bit 4-7   LEN - 3
                      Bit 8-15  Disp LSBs
                    If Reserved is 1: - Higher compression rates for files with (lots of) long repetitions
                      Bit 4-7   Indicator
                        If Indicator > 1:
                            Bit 0-3    Disp MSBs
                            Bit 4-7    LEN - 1 (same bits as Indicator)
                            Bit 8-15   Disp LSBs
                        If Indicator is 1: A(B CD E)(F GH)
                            Bit 0-3     (LEN - 0x111) MSBs
                            Bit 4-7     Indicator; unused
                            Bit 8-15    (LEN- 0x111) 'middle'-SBs
                            Bit 16-19   Disp MSBs
                            Bit 20-23   (LEN - 0x111) LSBs
                            Bit 24-31   Disp LSBs
                        If Indicator is 0:
                            Bit 0-3     (LEN - 0x11) MSBs
                            Bit 4-7     Indicator; unused
                            Bit 8-11    Disp MSBs
                            Bit 12-15   (LEN - 0x11) LSBs
                            Bit 16-23   Disp LSBs
             */
            //#endregion

            long readBytes = 0;

            byte type = (byte)instream.ReadByte();
            if (type != 0x11)
            {
                //throw new InvalidDataException("The provided stream is not a valid LZ-0x11 "
                //            + "compressed stream (invalid type 0x" + type.ToString("X") + ")");
                cout << "The provided stream is not a valid LZ-0x11 compressed stream (invalid type " << type << ")" << endl;
                return 0;
            }
            byte* sizeBytes = new byte[3];
            instream.Read(sizeBytes, 0, 3);
            int decompressedSize = ToNDSu24(sizeBytes, 0);
            readBytes += 4;
            if (decompressedSize == 0)
            {
                sizeBytes = new byte[4];
                instream.Read(sizeBytes, 0, 4);
                decompressedSize = ToNDSs32(sizeBytes, 0);
                readBytes += 4;
            }

            // the maximum 'DISP-1' is still 0xFFF.
            int bufferLength = 0x1000;
            byte* buffer = new byte[bufferLength];
            int bufferOffset = 0;

            int currentOutSize = 0;
            int flags = 0, mask = 1;
            while (currentOutSize < decompressedSize)
            {
                // (throws when requested new flags byte is not available)
                //#region Update the mask. If all flag bits have been read, get a new set.
                // the current mask is the mask used in the previous run. So if it masks the
                // last flag bit, get a new flags byte.
                if (mask == 1)
                {
                    if (readBytes >= inLength)
                    {
                        cout << "Not enough data" << endl;
                        return 0;
                        //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                    }
                    flags = instream.ReadByte();
                    readBytes++;
                    if (flags < 0)
                    {
                        cout << "Stream too short" << endl;
                        return 0;
                        //throw new StreamTooShortException();
                    }
                    mask = 0x80;
                }
                else
                {
                    mask >>= 1;
                }
                //#endregion

                // bit = 1 <=> compressed.
                if ((flags & mask) > 0)
                {
                    // (throws when not enough bytes are available)
                    //#region Get length and displacement('disp') values from next 2, 3 or 4 bytes

                    // read the first byte first, which also signals the size of the compressed block
                    if (readBytes >= inLength)
                    {
                        //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                        cout << "Not enough data " << currentOutSize << ", " << decompressedSize << endl;
                        return 0;
                    }
                    int byte1 = instream.ReadByte(); readBytes++;
                    if (byte1 < 0)
                    {
                        //throw new StreamTooShortException();
                        cout << "Stream too short" << endl;
                        return 0;
                    }

                    int length = byte1 >> 4;
                    int disp = -1;
                    if (length == 0)
                    {
                        //#region case 0; 0(B C)(D EF) + (0x11)(0x1) = (LEN)(DISP)

                        // case 0:
                        // data = AB CD EF (with A=0)
                        // LEN = ABC + 0x11 == BC + 0x11
                        // DISP = DEF + 1

                        // we need two more bytes available
                        if (readBytes + 1 >= inLength)
                        {
                            //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                            cout << "Not enough data " << currentOutSize << ", " << decompressedSize << endl;
                            return 0;
                        }
                        int byte2 = instream.ReadByte(); readBytes++;
                        int byte3 = instream.ReadByte(); readBytes++;
                        if (byte3 < 0)
                        {
                            //throw new StreamTooShortException();
                            cout << "Stream too short" << endl;
                            return 0;
                        }

                        length = (((byte1 & 0x0F) << 4) | (byte2 >> 4)) + 0x11;
                        disp = (((byte2 & 0x0F) << 8) | byte3) + 0x1;

                        //#endregion
                    }
                    else if (length == 1)
                    {
                        //#region case 1: 1(B CD E)(F GH) + (0x111)(0x1) = (LEN)(DISP)

                        // case 1:
                        // data = AB CD EF GH (with A=1)
                        // LEN = BCDE + 0x111
                        // DISP = FGH + 1

                        // we need three more bytes available
                        if (readBytes + 2 >= inLength)
                        {
                            //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                            cout << "Not enough data " << currentOutSize << ", " << decompressedSize << endl;
                            return 0;
                        }
                        int byte2 = instream.ReadByte(); readBytes++;
                        int byte3 = instream.ReadByte(); readBytes++;
                        int byte4 = instream.ReadByte(); readBytes++;
                        if (byte4 < 0)
                        {
                            //throw new StreamTooShortException();
                            cout << "Stream too short" << endl;
                            return 0;
                        }

                        length = (((byte1 & 0x0F) << 12) | (byte2 << 4) | (byte3 >> 4)) + 0x111;
                        disp = (((byte3 & 0x0F) << 8) | byte4) + 0x1;

                        //#endregion
                    }
                    else
                    {
                        //#region case > 1: (A)(B CD) + (0x1)(0x1) = (LEN)(DISP)

                        // case other:
                        // data = AB CD
                        // LEN = A + 1
                        // DISP = BCD + 1

                        // we need only one more byte available
                        if (readBytes >= inLength)
                        {
                            //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                            cout << "Not enough data " << currentOutSize << ", " << decompressedSize << endl;
                            return 0;
                        }
                        int byte2 = instream.ReadByte();
                        readBytes++;
                        if (byte2 < 0)
                        {
                            //throw new StreamTooShortException();
                            cout << "Stream too short" << endl;
                            return 0;
                        }

                        length = ((byte1 & 0xF0) >> 4) + 0x1;
                        disp = (((byte1 & 0x0F) << 8) | byte2) + 0x1;

                        //#endregion
                    }

                    if (disp > currentOutSize)
                    {
                        //throw new InvalidDataException("Cannot go back more than already written. "
                        //        + "DISP = " + disp + ", #written bytes = 0x" + currentOutSize.ToString("X")
                        //        + " before 0x" + instream.Position.ToString("X") + " with indicator 0x"
                        //        + (byte1 >> 4).ToString("X"));
                        cout << "Invalid data: Cannot go back more than already written." << endl;
                        return 0;
                    }
                    //#endregion

                    int bufIdx = bufferOffset + bufferLength - disp;
                    for (int i = 0; i < length; i++)
                    {
                        byte next = buffer[bufIdx % bufferLength];
                        bufIdx++;
                        outstream.WriteByte(next);
                        buffer[bufferOffset] = next;
                        bufferOffset = (bufferOffset + 1) % bufferLength;
                    }
                    currentOutSize += length;
                }
                else
                {
                    if (readBytes >= inLength)
                    {
                        //throw new NotEnoughDataException(currentOutSize, decompressedSize);
                        cout << "Not enough data " << currentOutSize << ", " << decompressedSize << endl;
                        return 0;
                    }
                    int next = instream.ReadByte(); readBytes++;
                    if (next < 0)
                    {
                        //throw new StreamTooShortException();
                        cout << "Stream too short" << endl;
                        return 0;
                    }

                    outstream.WriteByte((byte)next);
                    currentOutSize++;
                    buffer[bufferOffset] = (byte)next;
                    bufferOffset = (bufferOffset + 1) % bufferLength;
                }
            }

            if (readBytes < inLength)
            {
                // the input may be 4-byte aligned.
                if ((readBytes ^ (readBytes & 3)) + 4 < inLength)
                {
                    //throw new TooMuchInputException(readBytes, inLength);
                    //cout << "Too much input " << readBytes << ", " << inLength << endl;
                    //return 0;
                }
            }

            return decompressedSize;
        }

Stream::Stream()
{
    curReadLoc = 0;
}

byte Stream::ReadByte()
{
    byte ret = data.data()[curReadLoc];
    curReadLoc++;
    return ret;
}

void Stream::Read(byte* out, int unused, int amt)
{
    for(int i = 0; i < amt; i++)
    {
        out[i] = ReadByte();
    }
}

void Stream::WriteByte(byte b)
{
    data.push_back(b);
}



//main.cpp-------------------------------------------------------------------
#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <stdlib.h>
#include <cstdio>
#include <squish.h>
#ifdef _WIN32
    #include <windows.h>
#endif
#include "FreeImage.h"
#include <list>
#include <cmath>
#include <cstring>
#include "lz11.h"
using namespace std;

//extern unsigned char* LZX_Decode(char* input, int inputLength, int* outputLength);
//extern void lz11Decompress(const uint8_t *src, uint8_t *dst, int size);

bool g_bPieceTogether;

typedef struct
{
    uint32_t numImages;
    uint32_t _EOF;
    uint32_t _pieceOffset;
    uint32_t _imageOffset;
    uint16_t imageW;
    uint16_t imageH;
    uint8_t unknown0[8];
    uint16_t numPieces;
    uint8_t unknown1[2];
    uint32_t _unkOffset1;
    uint32_t _unkOffset2;
    uint32_t _unkOffset3;
} anbHeader;

typedef struct
{
    float x;
    float y;
} Vec2;

typedef struct
{
    Vec2 topLeft;
    Vec2 topLeftUV;
    Vec2 topRight;
    Vec2 topRightUV;
    Vec2 bottomRight;
    Vec2 bottomRightUV;
    Vec2 bottomLeft;
    Vec2 bottomLeftUV;
} piece;

typedef struct
{
    uint32_t x;
    uint32_t y;
} fakeVec2;

typedef struct
{
    fakeVec2 topLeft;
    fakeVec2 topLeftUV;
    fakeVec2 topRight;
    fakeVec2 topRightUV;
    fakeVec2 bottomRight;
    fakeVec2 bottomRightUV;
    fakeVec2 bottomLeft;
    fakeVec2 bottomLeftUV;
} fakePiece;

uint16_t byteSwap16(uint16_t byte)
{
    uint16_t hibyte = (byte & 0xff00) >> 8;
    uint16_t lobyte = (byte & 0x00ff);
    return lobyte << 8 | hibyte;
}

uint32_t byteSwap32(uint32_t val)
{
    val = ((val << 8) & 0xFF00FF00 ) | ((val >> 8) & 0xFF00FF );
    return (val << 16) | (val >> 16);
}

float toFloat(uint32_t val)
{
    float f;
    memcpy(&f, &val, sizeof(float));
    return f;
}

float byteSwapFloat(uint32_t val)
{
    val = byteSwap32(val);
    return toFloat(val);
}

Vec2 byteSwapVec2(fakeVec2 v)
{
    Vec2 ret;
    ret.x = byteSwapFloat(v.x);
    ret.y = byteSwapFloat(v.y);
    return ret;
}

piece byteSwapPiece(fakePiece p)
{
    piece ret;
    ret.topLeft = byteSwapVec2(p.topLeft);
    ret.topLeftUV = byteSwapVec2(p.topLeftUV);
    ret.topRight = byteSwapVec2(p.topRight);
    ret.topRightUV = byteSwapVec2(p.topRightUV);
    ret.bottomRight = byteSwapVec2(p.bottomRight);
    ret.bottomRightUV = byteSwapVec2(p.bottomRightUV);
    ret.bottomLeft = byteSwapVec2(p.bottomLeft);
    ret.bottomLeftUV = byteSwapVec2(p.bottomLeftUV);
    return ret;
}

FIBITMAP* imageFromPixels(uint8_t* imgData, uint32_t width, uint32_t height)
{
    //return FreeImage_ConvertFromRawBits(imgData, width, height, width*4, 32, 0xFF0000, 0x00FF00, 0x0000FF, true); //Doesn't seem to work
    FIBITMAP* curImg = FreeImage_Allocate(width, height, 32);
    FREE_IMAGE_TYPE image_type = FreeImage_GetImageType(curImg);
    if(image_type == FIT_BITMAP)
    {
        int curPos = 0;
        unsigned pitch = FreeImage_GetPitch(curImg);
        BYTE* bits = (BYTE*)FreeImage_GetBits(curImg);
        bits += pitch * height - pitch;
        for(int y = height-1; y >= 0; y--)
        {
            BYTE* pixel = (BYTE*)bits;
            for(int x = 0; x < width; x++)
            {
                pixel[FI_RGBA_RED] = imgData[curPos++];
                pixel[FI_RGBA_GREEN] = imgData[curPos++];
                pixel[FI_RGBA_BLUE] = imgData[curPos++];
                pixel[FI_RGBA_ALPHA] = imgData[curPos++];
                pixel += 4;
            }
            bits -= pitch;
        }
    }
    return curImg;
}

RGBQUAD makeColor(uint8_t a, uint8_t r, uint8_t b, uint8_t g)
{
    //RGBQUAD quad = { r, g, b, a };
    RGBQUAD quad = { r, a, g, b };
    return quad;
    //return (g << 24) | (b << 16) | (a << 8) | (r);
}

FIBITMAP* makeTexture(uint8_t* bytes, int byteslen, int width, int height)
{
    //byteslen = byteslen / 4;
    uint8_t* buffer1 = new uint8_t[byteslen / 2];
    uint8_t* buffer2 = new uint8_t[byteslen / 2];

    int i1 = 0, i2 = 0;
    int pointer = 0;

    while(i1 + i2 < byteslen)
    {
        for(int j = 0; j < 32; j++)
        {
            buffer1[i1++] = bytes[pointer++];
        }
        for(int k = 0; k < 32; k++)
        {
            buffer2[i2++] = bytes[pointer++];
        }
    }
    pointer = 0;

    //BufferedImage img = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
    FIBITMAP* img = FreeImage_Allocate(width, height, 32);

    int x = 0, y = 0;
    i1 = 0;
    i2 = 0;

    while(i1 + i2 < byteslen)
    {
        for(int h = 0; h < 4; h++)
        {
            for(int w = 0; w < 4; w++)
            {
                //img.setRGB(x + w, y + h, makeColor(buffer2[i2++], buffer2[i2++], buffer1[i1++ + 1], buffer1[i1++ - 1]));
                RGBQUAD colorQuad = makeColor(buffer2[i2++], buffer2[i2++], buffer1[i1++ + 1], buffer1[i1++ - 1]);
                FreeImage_SetPixelColor(img, x + w, y + h, &colorQuad);
            }
        }
        x += 4;
        if(x >= width)
        {
            x = 0;
            y += 4;
        }
    }
    delete[] buffer1;
    delete[] buffer2;
    return img;
}

void splitImages(const char* cFilename)
{
    vector<uint8_t> vData;
    FILE* fh = fopen( cFilename, "rb" );
    if(fh == NULL)
    {
        cerr << "Unable to open input file " << cFilename << endl;
        return;
    }
    fseek( fh, 0, SEEK_END );
    size_t fileSize = ftell( fh );
    fseek( fh, 0, SEEK_SET );
    vData.reserve( fileSize );
    size_t amt = fread( vData.data(), fileSize, 1, fh );
    fclose( fh );
    cout << "Splitting images from file " << cFilename << endl;

    //Figure out what we'll be naming the images
    string sName = cFilename;
    //First off, strip off filename extension
    size_t namepos = sName.find(".anb");
    if(namepos != string::npos)
        sName.erase(namepos);
    //Next, strip off any file path before it
    namepos = sName.rfind('/');
    if(namepos == string::npos)
        namepos = sName.rfind('\\');
    if(namepos != string::npos)
        sName.erase(0, namepos+1);

    anbHeader ANBh;
    memcpy(&ANBh, &(vData.data()[0]), sizeof(anbHeader));

    ANBh.numImages = byteSwap32(ANBh.numImages);
    ANBh._EOF = byteSwap32(ANBh._EOF);
    ANBh._pieceOffset = byteSwap32(ANBh._pieceOffset);
    ANBh._imageOffset = byteSwap32(ANBh._imageOffset);
    ANBh.imageH = byteSwap16(ANBh.imageH);
    ANBh.imageW = byteSwap16(ANBh.imageW);
    ANBh.numPieces = byteSwap16(ANBh.numPieces);

    //cout << ANBh.numImages << ", " << ANBh._EOF << ", " << ANBh._pieceOffset << ", " << ANBh._imageOffset << ", " << ANBh.imageW << ", " << ANBh.imageH << ", " << ANBh.numPieces << endl;

    //cout << toFloat(0xb0c10000) << endl;

    for(uint32_t j = 0; j < ANBh.numPieces*ANBh.numImages; j++)
    {
        piece p;
        fakePiece f;
        memcpy(&f, &(vData.data()[ANBh._pieceOffset+j*sizeof(fakePiece)]), sizeof(fakePiece));

        p = byteSwapPiece(f);
        //Store our maximum values, so we know how large the image is
        //cout << "Piece: "<< endl;
        //cout << p.topLeft.x << ", " << p.topLeft.y << ", " << p.topLeftUV.x  << ", " << p.topLeftUV.y << ", " << p.bottomRight.x << ", " << p.bottomRight.y << ", " << p.bottomRightUV.x << ", " << p.bottomRightUV.y << endl;

        //Sanity check: Skip over piecing if there's only one piece total that fills the whole thing
        //if(pd.numPieces == 1 && p.topLeftUV.x == 0.0 && p.topLeftUV.y == 0.0 && p.bottomRightUV.x == 1.0 && p.bottomRightUV.y == 1.0) bPieceTogether = false;

        //pieces.push_back(p);
    }

    //Decompress image data
    //int outputLength = ANBh.imageW * ANBh.imageH * 4;
    //uint8_t img_data[outputLength];// = LZX_Decode((char*)&(vData.data()[ANBh._imageOffset]), ANBh._EOF - ANBh._imageOffset, &outputLength);
    //lz11Decompress(&(vData.data()[ANBh._imageOffset]), img_data, ANBh._EOF - ANBh._imageOffset);
    Stream in;
    Stream out;
    //in.data = vData;
    in.data.reserve(ANBh.imageW * ANBh.imageH * 4);
    for(int i = ANBh._imageOffset; i < fileSize; i++)
    {
        in.data.push_back(vData.data()[i]);
    }
    cout << in.data.size() << endl;
    cout << ANBh._EOF - ANBh._imageOffset << endl;
    int outputLength = LZ11Decompress(in, in.data.size(), out);
    if(outputLength != ANBh.imageW * ANBh.imageH * 4)
        cout << "Length mismatch" << endl;

    //FILE* fpp = fopen( "blob_anvil.anb.data", "rb" );
    //if(fpp == NULL)
    //{
    //  cerr << "Unable to open input file " << cFilename << endl;
    //  return;
    //}
    //fseek(fpp, 0, SEEK_END);
    //size_t totalSz = ftell(fpp);
    //fseek(fpp, 0, SEEK_SET);
    //uint8_t* res = new uint8_t[ANBh.imageW * ANBh.imageH * 4];
    //memset(res, 0, ANBh.imageW * ANBh.imageH * 4);
    //vData.reserve(totalSz);
    //fread(res, totalSz, 1, fpp);
    //fclose(fpp);

    //cout << "Length: " << outputLength << endl;
    //FIBITMAP* img = makeTexture(res, ANBh.imageW * ANBh.imageH * 4, ANBh.imageW, ANBh.imageH);
    FIBITMAP* img = makeTexture((uint8_t*)out.data.data(), outputLength, ANBh.imageW, ANBh.imageH);
    //FIBITMAP* img = imageFromPixels((uint8_t*)img_data, ANBh.imageW, ANBh.imageH);
    ostringstream oss;
    oss << "output/" << sName << '_' << 1 << ".png";
    cout << "Saving " << oss.str() << endl;
    FreeImage_Save(FIF_PNG, img, oss.str().c_str());
    //FILE* fOut = fopen(oss.str().c_str(), "wb");
    //fwrite(out.data.data(), outputLength, 1, fOut);
    //fclose(fOut);
    //delete[] res;
    //free(img_data);
    //FreeImage_Unload(img);
}

int main(int argc, char** argv)
{
    g_bPieceTogether = true;
    FreeImage_Initialise();
#ifdef _WIN32
    CreateDirectory(TEXT("output"), NULL);
#else
    int result = system("mkdir -p output");
#endif
    list<string> sFilenames;
    //Parse commandline
    for(int i = 1; i < argc; i++)
    {
        string s = argv[i];
        if(s == "-nopiece")
            g_bPieceTogether = false;
        else
            sFilenames.push_back(s);
    }

    //Decompress ANB files
    for(list<string>::iterator i = sFilenames.begin(); i != sFilenames.end(); i++)
        splitImages((*i).c_str());
    FreeImage_DeInitialise();
    return 0;
}