/*---------------------------------------------------------------------------*

  Copyright (C) Nintendo.  All rights reserved.

  These coded instructions, statements, and computer programs contain
  proprietary information of Nintendo of America Inc. and/or Nintendo
  Company Ltd., and are protected by Federal copyright law.  They may
  not be disclosed to third parties or copied or duplicated in any form,
  in whole or in part, without the prior written consent of Nintendo.

 *---------------------------------------------------------------------------*/
////===========================================================================
///  demoCapture.c
///
///     This is capture code for the demo library.
///
////===========================================================================

#include <cafe/demo.h>
#include <cafe/gx2.h>

DEMOCaptureDataStore DEMOCaptureData; // global
static void* DEMOCaptureDLMem = NULL; //display list memory
#define DEMO_DL_SIZE_MAX  (8*1024)

typedef struct _DEMOColor
{
    u8 x, y, z, w;
} DEMOColor;

void _DEMOEncodeTGA(const u8* inData, u8* outData, u32* outByteSize,
                    u32 width, u32 height, u32 pitch, GX2SurfaceFormat format);
void _DEMOComputeRunlengths(u32 cols, DEMOColor *pixelrow, s32 *runlength);
void _DEMODumpSerial(const u8* data, const u32 byte);

// GX2_UUENC is a basic 1 character encoding function to make a char printing
#define _DEMO_UUENC(c) ((c) ? ((c) & 077) + ' ': '`')

void _DEMOComputeRunlengths(u32 cols, DEMOColor *pixelrow, s32 *runlength)
{
    s32 col, start;

    // Initialize all run lengths to 0. (This is just an error check.)
    for (col = 0; col < cols; ++col)
        runlength[col] = 0;

    // Find runs of identical pixels.
    for ( col = 0; col < cols; ) {
        start = col;
        do {
            ++col;
        } while ( col < cols &&
                  col - start < 128 &&
                  (pixelrow[col].x == pixelrow[start].x) &&
                  (pixelrow[col].y == pixelrow[start].y) &&
                  (pixelrow[col].z == pixelrow[start].z) &&
                  (pixelrow[col].w == pixelrow[start].w)  );
        runlength[start] = col - start;
    }

    // Now look for runs of length-1 runs, and turn them into negative runs.
    for (col = 0; col < cols; )
    {
        if (runlength[col] == 1)
        {
            start = col;
            while (col < cols &&
                   col - start < 128 &&
                   (runlength[col] == 1) )
            {
                runlength[col] = 0;
                ++col;
            }
            runlength[start] = - ( col - start );
        } else
            col += runlength[col];
    }
}

void _DEMODumpSerial(const u8* data, const u32 byte)
{
    u32 i, j, idx, leftByte, numLoop;
    s32 ch;
    // uuencode allows 6 bits (=63)
    const u32 maxOneLine = 63;
    // 63/3*4 = 84 characters
    char line[90];
    char *pc;

    idx = 0;
    numLoop = (u32)(byte / maxOneLine);

    DEMOPrintf("begin 644 GX2Capture.tga\n");

    for(i=0; i<numLoop; i++)
    {
        pc = &line[0];
        *pc++ = '_'; // '_' is 63+32 = 95 in ascii
        for(j=0; j<maxOneLine/3; j++)
        {
            ch = data[3*idx] >> 2;
            ch = _DEMO_UUENC(ch);
            *pc++ = ch;
            ch = ((data[3*idx] << 4) & 060) | ((data[3*idx+1] >> 4) & 017);
            ch = _DEMO_UUENC(ch);
            *pc++ = ch;
            ch = ((data[3*idx+1] << 2) & 074) | ((data[3*idx+2] >> 6) & 03);
            ch = _DEMO_UUENC(ch);
            *pc++ = ch;
            ch = data[3*idx+2] & 077;
            ch = _DEMO_UUENC(ch);
            *pc++ = ch;
            idx++;
        }
        *pc++ = '\n';
        *pc++ = 0;
        DEMOPrintf("%s", line);
    }

    leftByte = byte % maxOneLine;
    if (leftByte > 0) {
        pc = &line[0];
        *pc++ = leftByte+32;

        if (leftByte >= 3) {
            for(j=0; j<leftByte/3; j++)
            {
                ch = data[3*idx] >> 2;
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = ((data[3*idx] << 4) & 060) | ((data[3*idx+1] >> 4) & 017);
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = ((data[3*idx+1] << 2) & 074) | ((data[3*idx+2] >> 6) & 03);
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = data[3*idx+2] & 077;
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                idx++;
            }
        }

        if(leftByte % 3 != 0)
        {
            ch = data[3*idx] >> 2;
            ch = _DEMO_UUENC(ch);
            *pc++ = ch;
            if(leftByte % 3 == 2){
                ch = ((data[3*idx] << 4) & 060) | ((data[3*idx+1] >> 4) & 017);
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = ((data[3*idx+1] << 2) & 074);
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = 0;
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
            }else{
                ch = ((data[3*idx] << 4) & 060);
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = 0;
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
                ch = data[3*idx+2] & 077;
                ch = _DEMO_UUENC(ch);
                *pc++ = ch;
            }
        }
        *pc++ = '\n';
        *pc++ = 0;
        DEMOPrintf("%s", line);
    }
    DEMOPrintf("`\nend\n");
}

void _DEMOEncodeTGA(const u8* inData, u8* outData, u32* outByteSize,
                    u32 width, u32 height, u32 pitch, GX2SurfaceFormat format)
{
    u8 *pbuffer;
    u32 rows = DEMOCaptureData.colorBuffer.surface.height;
    u32 cols = DEMOCaptureData.colorBuffer.surface.width;
    s32* runlength;
    u32 row, col, i, realrow, byte = 0;
    DEMOColor* buffer;

    DEMOAssert( GX2_SURFACE_FORMAT_TCS_R8_G8_B8_A8_SRGB == format ||
                GX2_SURFACE_FORMAT_TCS_R8_G8_B8_A8_UNORM == format );

    pbuffer = outData;
    DEMOAssert(pbuffer);

    // 18-byte TGA header
    // (16-bit values are written in little-endian order)
    // TGA Image format codes:
    //  1: uncompressed, color-mapped      9: RLE-compressed, color-mapped
    //  2: uncompressed, true-color       10: RLE-compressed, true-color
    //  3: uncompressed, BW or gray       11: RLE-compressed, BW or gray
    //
    *pbuffer++ = 0;                  // Image ID length (0: no ID)
    *pbuffer++ = 0;                  // Color map type  (0: no color map)
    *pbuffer++ = 10;                 // Image type (10: RLE true-color image)
    *pbuffer++ = 0; *pbuffer++ = 0;  // Color map first index offset (16 bit)
    *pbuffer++ = 0; *pbuffer++ = 0;  // Color map length (16 bit)
    *pbuffer++ = 0;                  // Color map bpp
    *pbuffer++ = 0; *pbuffer++ = 0;  // Image X origin (16 bit)
    *pbuffer++ = 0; *pbuffer++ = 0;  // Image Y origin (16 bit)
    *pbuffer++ = (u8)(width % 256);  // Image width lo
    *pbuffer++ = (u8)(width / 256);  // Image width hi
    *pbuffer++ = (u8)(height % 256); // Image height lo
    *pbuffer++ = (u8)(height / 256); // Image height hi
    *pbuffer++ = (u8) 32;            // Image pixel bpp (32 bits)
    *pbuffer++ = (u8) 8;             // Image alpha bpp (8 bits), direction code

    buffer = (DEMOColor*)inData;

    // do run length encoding here
    runlength = DEMOAlloc(sizeof(s32)*cols);
    DEMOAssert(runlength && "DEMOAlloc failed");

    // write out pixels
    for( row = 0; row < rows; ++row )
    {
        realrow = rows - row - 1;
        _DEMOComputeRunlengths( cols, &buffer[realrow * pitch], runlength );
        for( col = 0; col < cols; ){
            if( runlength[col] > 0 ){
                // set runlength
                pbuffer[byte] = (u8)(0x80 + runlength[col] - 1); byte++;
                // write pixel
                // TGA expects BGRA order
                pbuffer[byte] = buffer[realrow * pitch + col].z; byte++;
                pbuffer[byte] = buffer[realrow * pitch + col].y; byte++;
                pbuffer[byte] = buffer[realrow * pitch + col].x; byte++;
                pbuffer[byte] = buffer[realrow * pitch + col].w; byte++;
                col += runlength[col];
            }else if( runlength[col] < 0 ){
                // set runlength
                pbuffer[byte] = (u8)(- runlength[col] - 1); byte++;
                // write pixels
                // TGA expects BGRA order
                for ( i = 0; i < - runlength[col]; ++i )
                {
                    pbuffer[byte] = buffer[realrow * pitch + col + i].z; byte++;
                    pbuffer[byte] = buffer[realrow * pitch + col + i].y; byte++;
                    pbuffer[byte] = buffer[realrow * pitch + col + i].x; byte++;
                    pbuffer[byte] = buffer[realrow * pitch + col + i].w; byte++;
                }
                col += -runlength[col];
            }else
                DEMOAssert(!"Internal error: zero run length");
        }
    }

    DEMOFree(runlength);

    *outByteSize = 18+byte;
}

void DEMOCaptureInit(u32 width, u32 height, DEMOCaptureFormat format)
{
    DEMOCaptureData.captureFormat = format;
    // Setup render buffer
    DEMOCaptureData.colorBuffer.surface.dim = GX2_SURFACE_DIM_2D;
    DEMOCaptureData.colorBuffer.surface.width = width;
    DEMOCaptureData.colorBuffer.surface.height = height;
    DEMOCaptureData.colorBuffer.surface.depth = 1;
    DEMOCaptureData.colorBuffer.surface.numMips = 0;
    if( SRGB8 == format ){
        DEMOCaptureData.colorBuffer.surface.format =
            GX2_SURFACE_FORMAT_TCS_R8_G8_B8_A8_SRGB;
    } else if( RGBA8 == format ){
        DEMOCaptureData.colorBuffer.surface.format =
            GX2_SURFACE_FORMAT_TCS_R8_G8_B8_A8_UNORM;
    } else {
        DEMOAssert("Only RGBA8/SRGB8 formats supported at the moment.");
    }

    DEMOCaptureData.colorBuffer.surface.aa = GX2_AA_MODE_1X;
    DEMOCaptureData.colorBuffer.surface.use =
        GX2_SURFACE_USE_COLOR_BUFFER_TEXTURE;
    DEMOCaptureData.colorBuffer.surface.tileMode = GX2_TILE_MODE_LINEAR_ALIGNED;
    DEMOCaptureData.colorBuffer.surface.swizzle  = 0;
    DEMOCaptureData.colorBuffer.viewMip = 0;
    DEMOCaptureData.colorBuffer.viewFirstSlice = 0;
    DEMOCaptureData.colorBuffer.viewNumSlices = 1;
    GX2CalcSurfaceSizeAndAlignment(&DEMOCaptureData.colorBuffer.surface);
    GX2InitColorBufferRegs(&DEMOCaptureData.colorBuffer);

    if (DEMOCaptureData.colorBuffer.surface.imagePtr) {
        DEMOGfxFreeMEM2(DEMOCaptureData.colorBuffer.surface.imagePtr);
    }
    DEMOCaptureData.colorBuffer.surface.imagePtr
        =  DEMOGfxAllocMEM2(DEMOCaptureData.colorBuffer.surface.imageSize,
                             DEMOCaptureData.colorBuffer.surface.alignment);
    DEMOAssert(DEMOCaptureData.colorBuffer.surface.imagePtr);
    GX2Invalidate(GX2_INVALIDATE_CPU,
                  DEMOCaptureData.colorBuffer.surface.imagePtr,
                  DEMOCaptureData.colorBuffer.surface.imageSize);

    if (DEMOCaptureData.TGAData) {
        DEMOFree(DEMOCaptureData.TGAData);
    }
    // Worst case scenario is image 1/4 bigger than original
    // because run length encoding adds up to 1 byte per 4 byte pixel
    DEMOCaptureData.TGAData = DEMOAlloc((u32)(18+5*width*height));
    DEMOAssert( DEMOCaptureData.TGAData );

    //Setup Display list support
    if(DEMOCaptureDLMem == NULL)
    {
        DEMOCaptureDLMem = DEMOGfxAllocMEM2(DEMO_DL_SIZE_MAX, 
                                            GX2_DISPLAY_LIST_ALIGNMENT);
        DEMOAssert(DEMOCaptureDLMem);
    }
}

void DEMOCaptureShutdown(void)
{
    if (DEMOCaptureData.colorBuffer.surface.imagePtr) {
        DEMOGfxFreeMEM2(DEMOCaptureData.colorBuffer.surface.imagePtr);
        DEMOCaptureData.colorBuffer.surface.imagePtr = NULL;
    }

    if(DEMOCaptureData.TGAData) {
        DEMOGfxFreeMEM2(DEMOCaptureData.TGAData);
        DEMOCaptureData.TGAData = NULL;
    }

    if(DEMOCaptureDLMem) {
        DEMOGfxFreeMEM2(DEMOCaptureDLMem);
        DEMOCaptureDLMem = NULL;
    }
}

void DEMOCaptureCopyHelper(
    const GX2Surface *srcSurface, u32 srcMip, u32 srcSlice,
    GX2Surface *dstSurface, u32 dstMip, u32 dstSlice)
{
    u32 dlSize;
    DEMOAssert((DEMOCaptureDLMem != NULL) && "DEMOCaptureInit must be called first");

    // Use Display lists and disable profiling modes - This prevents the profiling
    // modes from impacting the capture copy.
    // GX2DrawDone() is always called after using the display list below,
    // so the memory should be free to reuse again.
    GX2Invalidate(GX2_INVALIDATE_CPU, DEMOCaptureDLMem, DEMO_DL_SIZE_MAX);
    GX2BeginDisplayListEx(DEMOCaptureDLMem, DEMO_DL_SIZE_MAX, GX2_FALSE);

    GX2UTSetAndCopySurface(srcSurface, srcMip, srcSlice,
                           dstSurface, dstMip, dstSlice);

    dlSize = GX2EndDisplayList(DEMOCaptureDLMem);
    ASSERT(dlSize <= DEMO_DL_SIZE_MAX);
    GX2CallDisplayList(DEMOCaptureDLMem, dlSize);
}

void DEMOCaptureCopy(GX2ColorBuffer *srcBuffer, const char *dstFilePath)
{
    DEMOAssert(DEMOCaptureData.TGAData != NULL && "DEMOCaptureInit must be called first");

    if (srcBuffer != NULL) {

        GX2Boolean ok;

        // Make sure dest buffer isn't in CPU cache
        GX2Invalidate(GX2_INVALIDATE_CPU,
                      DEMOCaptureData.colorBuffer.surface.imagePtr,
                      DEMOCaptureData.colorBuffer.surface.imageSize);

        // Grab a copy
        if (srcBuffer->surface.aa == GX2_AA_MODE_1X) {
            DEMOCaptureCopyHelper(&srcBuffer->surface,
                srcBuffer->viewMip,
                srcBuffer->viewFirstSlice,
                &DEMOCaptureData.colorBuffer.surface, 0, 0);
        } else {
            // AA, so need to resolve into a temporary buffer first
            GX2Surface tempSurf;
            tempSurf = srcBuffer->surface;
            tempSurf.aa = GX2_AA_MODE_1X;
            GX2CalcSurfaceSizeAndAlignment(&tempSurf);
            tempSurf.imagePtr = DEMOGfxAllocMEM2(tempSurf.imageSize,
                                                 tempSurf.alignment);
            DEMOAssert(tempSurf.imagePtr);  
            GX2ResolveAAColorBuffer(srcBuffer,
                                    &tempSurf, 0, 0);
            DEMOGfxSetContextState();
            DEMOCaptureCopyHelper(&tempSurf, 0, 0,
                &DEMOCaptureData.colorBuffer.surface, 0, 0);
            GX2DrawDone();
            DEMOGfxFreeMEM2(tempSurf.imagePtr);
        }

        DEMOGfxSetContextState();

        // Flush out dest caches
        GX2Invalidate(GX2_INVALIDATE_COLOR_BUFFER,
                      DEMOCaptureData.colorBuffer.surface.imagePtr,
                      DEMOCaptureData.colorBuffer.surface.imageSize);

        // Wait for GPU to finish
        ok = GX2DrawDone();

        // If DrawDone timed out, image is not valid, so set to 0xff
        if (!ok) {
            memset(DEMOCaptureData.colorBuffer.surface.imagePtr, 0xff,
                   DEMOCaptureData.colorBuffer.surface.imageSize);
        }

        // Convert to TGA format
        _DEMOEncodeTGA(DEMOCaptureData.colorBuffer.surface.imagePtr,
                       DEMOCaptureData.TGAData,
                       &DEMOCaptureData.TGALength,
                       DEMOCaptureData.colorBuffer.surface.width,
                       DEMOCaptureData.colorBuffer.surface.height,
                       DEMOCaptureData.colorBuffer.surface.pitch,
                       DEMOCaptureData.colorBuffer.surface.format);
    }

    if(NULL != dstFilePath) {
        if(*dstFilePath==0)
        {
            _DEMODumpSerial(DEMOCaptureData.TGAData, DEMOCaptureData.TGALength);
        } else {
            s32 res;
            DEMOFSFileInfo finfo;
            res = DEMOFSOpenFileMode(dstFilePath, &finfo, "w");
            if (res != DEMO_FS_RESULT_OK) {
                DEMOPrintf("Unable to open file %s for capture output\n", dstFilePath);
                return;
            }
            res = DEMOFSWrite(&finfo, DEMOCaptureData.TGAData, DEMOCaptureData.TGALength);
            if (res != DEMO_FS_RESULT_OK) {
                DEMOPrintf("Unable to write to file %s for capture output\n", dstFilePath);
                return;
            }
            res = DEMOFSCloseFile(&finfo);
            if (res != DEMO_FS_RESULT_OK) {
                DEMOPrintf("Error closing file %s for capture output\n", dstFilePath);
                return;
            }
        }
    }
}