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

  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.

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

#include <stdio.h>
#include <string.h>
#include <math.h>

#if defined(WIN32) || defined(WIN64)
#include <pc/gx2.h>
#include <pc/demo.h>
#endif
#include <cafe/os.h>
#include <cafe/mem.h>
#include <cafe/gx2.h>
#include <cafe/demo.h>

#include <cafe/gfd.h>
#include <cafe/gx2ut.h>

//Include generated shaders
#include "shaders/headers/gx2utExpandDepth.h"


// ----- GX2 Shader
static const GX2VertexShader* const VS_SHADERS[] = { &gx2utExpandDepth_VS,};

static const GX2PixelShader* const PS_SHADERS[] = { &gx2utExpandDepth_PS,
};

static const u32 NUM_SHADERS = 1;

typedef struct _ExpandDepthShader {
    // These variables hold the three types of shaders needed for a call to
    // GX2SetShaders. The vertex and pixel shaders are loaded from the
    // header, but since the fetch shader is generated at run-time
    // it must be handled slightly differently.
    const GX2VertexShader *pVertexShader;
    const GX2PixelShader *pPixelShader;

    // The register locations where the offset uniforms are stored for
    // the pixel and vertex shaders.
    u32 u_positionLocation;
    u32 u_texcoordLocation;

} ExpandDepthShader;

static ExpandDepthShader g_expandDepthShader[NUM_SHADERS];
static GX2FetchShader fetchShader;

#define FETCH_SHADER_SIZE 32  //hard code this value for now
ALIGNVAR(GX2_SHADER_ALIGNMENT) static u8 g_GX2UTFetchShader[FETCH_SHADER_SIZE];

// ----- GX2 Texture

typedef struct _VtxFmtF32x2 {
    f32 texcoord[2];
} VtxFmtF32x2;

static const VtxFmtF32x2 EXPAND_SURFACE_RECT_POSITION_DATA[] =
{
    {0.0f,  0.0f},
    {1.0f,  0.0f},
    {1.0f,  1.0f},
    {0.0f,  1.0f}
};

// Same for both rect and triangle strips
static const u32 VERTEX_COUNT = sizeof(EXPAND_SURFACE_RECT_POSITION_DATA)
                         / sizeof(EXPAND_SURFACE_RECT_POSITION_DATA[0]);

// Initializes how surfaces will be copied
void GX2UTExpandDepthInit()
{
    // Setup shaders
    u32 i;

    for (i = 0; i < NUM_SHADERS; ++i)
    {
        g_expandDepthShader[i].pVertexShader = VS_SHADERS[i];
        g_expandDepthShader[i].pPixelShader = PS_SHADERS[i];

        GX2Invalidate(GX2_INVALIDATE_CPU_SHADER,
            g_expandDepthShader[i].pVertexShader->shaderPtr,
            g_expandDepthShader[i].pVertexShader->shaderSize);

        GX2NotifyMemAlloc(g_expandDepthShader[i].pVertexShader->shaderPtr,
            g_expandDepthShader[i].pVertexShader->shaderSize,
            GX2_SHADER_ALIGNMENT);

        GX2Invalidate(GX2_INVALIDATE_CPU_SHADER,
            g_expandDepthShader[i].pPixelShader->shaderPtr,
            g_expandDepthShader[i].pPixelShader->shaderSize);

        GX2NotifyMemAlloc(g_expandDepthShader[i].pPixelShader->shaderPtr,
            g_expandDepthShader[i].pPixelShader->shaderSize,
            GX2_SHADER_ALIGNMENT);

        // Lookup the uniform locations in the vertex shader.
        // The shader author chose the name "u_positions"
        g_expandDepthShader[i].u_positionLocation =
            (u32)GX2GetVertexUniformVarOffset(g_expandDepthShader[i].pVertexShader, "u_positions");
        ASSERT((g_expandDepthShader[i].u_positionLocation != GX2_UNIFORM_VAR_INVALID_OFFSET)
            && "Couldn't find the correct vertex shader uniforms.");

    }

    ASSERT(GX2CalcFetchShaderSize(0) <= sizeof(g_GX2UTFetchShader) && "g_GX2UTFetchShader too small!\n");
    GX2InitFetchShader(&fetchShader, g_GX2UTFetchShader, 0, NULL);
}

void GX2UTExpandDepthBufferOp(GX2DepthBuffer* depthBuffer)
{
    static GX2Boolean initDone = GX2_FALSE;
    u32 expandWidth;
    u32 expandHeight;
    u32 shaderIdx = 0;

    GX2UTDebugTagIndent(__func__);

    if (initDone == GX2_FALSE)
    {
        // Initialize the resources needed to expand a depth buffer
        GX2UTExpandDepthInit();

        initDone = GX2_TRUE;
    }

    ASSERT((depthBuffer != NULL));

    //
    // Enable special depth expand state
    //
    // Invalidate the depth buffer to guarantee all writes have been flushed
    if (depthBuffer->viewMip)
        GX2Invalidate(GX2_INVALIDATE_DEPTH_BUFFER, depthBuffer->surface.mipPtr, depthBuffer->surface.mipSize);
    else
        GX2Invalidate(GX2_INVALIDATE_DEPTH_BUFFER, depthBuffer->surface.imagePtr, depthBuffer->surface.imageSize);

    for (int slice = depthBuffer->viewFirstSlice; slice < depthBuffer->viewNumSlices; slice++)
    {
        GX2DepthBuffer db = *depthBuffer;

        // Set the view slice for this pass
        db.viewFirstSlice = slice;
        db.viewNumSlices = 1;
        GX2InitDepthBufferRegs(&db);

        // Copy the dimensions
        expandWidth = db.surface.width;
        expandHeight = db.surface.height;

        // Set shaders
        GX2SetFetchShader(&fetchShader);
        GX2SetVertexShader(g_expandDepthShader[shaderIdx].pVertexShader);
        GX2SetPixelShader(g_expandDepthShader[shaderIdx].pPixelShader);

        // Set the uniforms to be used by the vertex shader
        f32 position_base_scale[] = 
        {
            -1.0f,
             1.0f,
             2.0f,
            -2.0f,
        };

        for (int i = 0; i < 4; i++)
        {
            f32 pos[] = {
                position_base_scale[0] + position_base_scale[2] * EXPAND_SURFACE_RECT_POSITION_DATA[i].texcoord[0],
                position_base_scale[1] + position_base_scale[3] * EXPAND_SURFACE_RECT_POSITION_DATA[i].texcoord[1],
                0.0,
                1.0
            };

            GX2SetVertexUniformReg(g_expandDepthShader[shaderIdx].u_positionLocation + i*4, 1*4, pos);
        }

        // Set the depth buffer
        GX2SetDepthBuffer(&db);

        // Render to destination surface dimensions
        GX2SetViewport(0, 0, (f32)expandWidth, (f32)expandHeight, 0.0f, 1.0f);
        GX2SetScissor(0, 0, expandWidth, expandHeight);

        // Draw a full quad that covers the display
        GX2Draw(GX2_PRIMITIVE_RECTS, VERTEX_COUNT);
    }

    // Invalidate the depth buffer to guarantee all writes have been flushed
    if (depthBuffer->viewMip)
        GX2Invalidate(GX2_INVALIDATE_DEPTH_BUFFER, depthBuffer->surface.mipPtr, depthBuffer->surface.mipSize);
    else
        GX2Invalidate(GX2_INVALIDATE_DEPTH_BUFFER, depthBuffer->surface.imagePtr, depthBuffer->surface.imageSize);


    GX2UTDebugTagUndent();
}


//Setup all of the constant renderstate needed for the copy.
void GX2UTSetExpandDepthState(GX2Boolean enable)
{
    if (enable)
    {
        // If your application's steady state can be set to GX2UT common state
        // using a small number of discrete GX2 calls, then customize here
        // instead of using GX2UTSetCommonState().
        GX2UTSetCommonState();

        // Set additional required GX2 state
        GX2SetColorControl(GX2_LOGIC_OP_COPY, GX2_DISABLE, GX2_DISABLE, GX2_DISABLE);

        // Setup depth/stencil state:
        GX2SetDepthStencilControl(GX2_FALSE,
                                  GX2_FALSE,
                                  GX2_COMPARE_NEVER,
                                  GX2_FALSE,
                                  GX2_FALSE,
                                  GX2_COMPARE_NEVER,
                                  GX2_STENCIL_KEEP,
                                  GX2_STENCIL_KEEP,
                                  GX2_STENCIL_KEEP,
                                  GX2_COMPARE_NEVER,
                                  GX2_STENCIL_KEEP,
                                  GX2_STENCIL_KEEP,
                                  GX2_STENCIL_KEEP);

        // Set Complex GX2 State for our operation
        GX2SetSpecialState(GX2_SPECIAL_STATE_EXPAND_DEPTH, GX2_ENABLE);
    }
    else
    {
        // Disable any complex GX2 state
        GX2SetSpecialState(GX2_SPECIAL_STATE_EXPAND_DEPTH, GX2_DISABLE);

        // Set additional required GX2 state
        GX2SetColorControl(GX2_LOGIC_OP_COPY, GX2_DISABLE, GX2_DISABLE, GX2_ENABLE);

        // The purpose of the following is to return the context to GX2 default
        // state.  If your application uses a different "steady state", then
        // customize
        GX2SetDepthStencilControl(
              GX2_TRUE,            //depthTestEnable
              GX2_TRUE,            //depthWriteEnable
              GX2_COMPARE_LESS,    //depthFunc
              GX2_FALSE,           //stencilTestEnable
              GX2_FALSE,           //backStencilEnable
              GX2_COMPARE_ALWAYS,  //frontStencilFunc
              GX2_STENCIL_REPLACE, //frontStencilZPass
              GX2_STENCIL_REPLACE, //frontStencilZFail
              GX2_STENCIL_REPLACE, //frontStencilFail
              GX2_COMPARE_ALWAYS,  //backStencilFunc
              GX2_STENCIL_REPLACE, //backStencilZPass
              GX2_STENCIL_REPLACE, //backStencilZFail
              GX2_STENCIL_REPLACE);//backStencilFail
    }

}