#include <PR/bbboot.h>
#include <PR/bbnand.h>
#include <PR/os_bbsa.h>

#include <sha1.h>

#include "sysapp.h"
#include "skerror.h"
#include "gzip.h"

/* Use block links to find next block */
#define _USE_BLOCK_LINKS 1

#define ALIGN_DCACHE __attribute__((aligned(DCACHE_LINESIZE)))

static u8 tik_data[PI_FLASH_PAGES_PER_BLOCK * PI_FLASH_PAGE_DATA_SIZE] ALIGN_DCACHE;

#define CMDBUNDLE_STORAGE_START   tik_data

/* * * * * * *  READING A FLASH PAGE * * * * * * * * * * * */

/* copied from sk/util.h */

#define PI_FLASH_DEV0_BUF0_READ_PAGE 0x9f008a10
#define PI_FLASH_DEV0_BUF1_READ_PAGE \
          (PI_FLASH_DEV0_BUF0_READ_PAGE | PI_FLASH_CTRL_BUF1) 

#define FLASH_MODULE_PRESENT \
             (!((IO_READ(MI_EINTR_REG)&MI_EINTR_MODULE_REMOVED)))

/* copied from sk/util.c */
/*
 * page based app loading
 */
int readFlashPage(u32 pageNum,int piBuf)
{
    IO_WRITE(PI_FLASH_ADDR_REG,pageNum<<PI_FLASH_PAGE_ADDR_SHIFT);
    if(piBuf)
        IO_WRITE(PI_FLASH_CTRL_REG,PI_FLASH_DEV0_BUF1_READ_PAGE);
    else
        IO_WRITE(PI_FLASH_CTRL_REG,PI_FLASH_DEV0_BUF0_READ_PAGE);

    do{
        if(!FLASH_MODULE_PRESENT){
            /* module not present during operation */

            /* stop current operation */
            IO_WRITE(PI_FLASH_CTRL_REG,0);

            return SK_ERROR_FLASH_MODULE_REMOVED;
        }
    }while(IO_READ(PI_FLASH_CTRL_REG)&PI_FLASH_CTRL_BUSY);

    if(IO_READ(PI_FLASH_CTRL_REG) & PI_FLASH_CTRL_DBERR)
        return SK_ERROR_FLASH_CTRL_DOUBLE_BIT;

    return SK_SUCCESS;
}

/* * * * * * *  FINDING GOOD BLOCKS * * * * * * * * * * * */

/* Copied from sk/boot.c */

#define SK_BLOCKS (LOAD_SK_PAGES/PI_FLASH_PAGES_PER_BLOCK)

/* NOTE: best 2of3 for 1B value in spare area bytes 512, 513 and 514.
 *       Making assumption that a 2of3 match exists because:
 *          i) it's simpler code
 *         ii) someone could attack by setting up flash anyhow, so 
 *             we don't have a true protection mechanism
 *        iii) the result of loading the wrong block is just that
 *             bytes pass through the decryption engine (someone
 *             could be attempting to crack the key), but the code
 *             is never run since the hash will fail.
 */
u32 getBlockLink(u32 encoded)
{
    if( ((encoded>>8)&0xff) == ((encoded>>16)&0xff))
        return (encoded>>8)&0xff;
    return (encoded>>24)&0xff;
}

/* Return is error if < 0 (as per def'ns in boot.h), and block
 * index when >=0.
 *
 * NOTE: first page of good block will be in pi buf 0. 
 */
int getGoodBlock(int startSearch)
{
    int i, badbits, bitshift;
    u32 blockStatus;
    u32 ret;

    do{
        ret = readFlashPage(startSearch<<PI_FLASH_PAGE_PER_BLOCK_POW2,0);
       if(ret == SK_ERROR_FLASH_MODULE_REMOVED){
            return ret;
        }
        blockStatus = IO_READ(PI_BUFFER_0_OOB_START +
                              (PI_FLASH_OOB_BLOCK_STATUS_OFFSET&(~3)));

        /* Given the current definitions in bbnand.h, blockStatus
         * holds bytes 516, 517, 518 and 519 from the flash page.
         * If more than one bit of byte 517 is 0, the block is bad.
         * The bitshift computation will insure byte 517 is checked
         * (or whatever the defines in bbnand.h dictate).
         */
        bitshift = (3 - (PI_FLASH_OOB_BLOCK_STATUS_OFFSET&3))*8;
        for(badbits=0,i=0;i<8;i++){
            if( !((blockStatus>>(bitshift+i))&1) ){
                PRINTF("Found 0 bit in block status byte.\n");
                badbits++;
            }
        }
        startSearch++;
    }while(badbits>1);

    if(ret == SK_ERROR_FLASH_CTRL_DOUBLE_BIT){
        /* block status was good, but had double-bit ECC error */
        PRINTF("Double bit error in getGoodBlock.\n");
        return ret;
    }

    return startSearch-1;
}

/* skip over "skip" good blocks, starting at block startBlock,
 * and return the number of the first good block past the skip.
 *
 * Return is error if < 0 (as per def'ns in boot.h), and block
 * index when >=0.
 *
 * NOTE: thgoto-e first page of block with the returned index will
 *       be in PI buffer 0.
 */
int findGoodBlock(int startBlock,int skip)
{
    int i;

    for(i=0;i<=skip;i++){
        startBlock = getGoodBlock(startBlock);
        if(startBlock<0) /* error */
            break;
        startBlock++;
    }

    return startBlock-1;
}

/* Macro Copied from sk/boot.c */

/* patch BbCrlBundle pointers from ticket bundle (packaged as
 * offsets from wherever 16KB ticket bundle is loaded).
 */
#define BOOT_PATCH_RL(_name,_off)                                      \
    if(rlBundle->##_name.head != NULL) rlBundle->##_name.head  =       \
                (BbCrlHead *)(((u32)(rlBundle->##_name.head))          \
                + (u32)_off);                                          \
    rlBundle->##_name.list =                                           \
                (BbServerSuffix *)(((u32)(rlBundle->##_name.list))     \
                + (u32)_off);                                          \
    rlBundle->##_name.certChain[0] =                                   \
                (BbCertBase *)(((u32)(rlBundle->##_name.certChain[0])) \
                + (u32)_off);                                          \
    rlBundle->##_name.certChain[1] =                                   \
                (BbCertBase *)(((u32)(rlBundle->##_name.certChain[1])) \
                + (u32)_off);

/* * * * * * *  VERIFYING HASH * * * * * * * * * * * */

/*
 * Verifies the hash for a block of data 
 */
int verify_hash(BbShaHash expected_hash, u8* data, u32 length)
{
    SHA1Context sha;
    BbShaHash hash;

    /* initialize SHA */
    SHA1Reset(&sha);
    SHA1Input(&sha, data, length);
    SHA1Result(&sha, (u8 *) hash);

    PRINTF("hash value = 0x%08x %08x %08x %08x %08x\n",
           hash[0], hash[1], hash[2], hash[3], hash[4]);

    if (memcmp(hash, expected_hash, sizeof(hash))) {
        PRINTF("hash mismatch\n");
        return 0;
    }
    else {
        PRINTF("hash matched\n");
        return 1;
    }

}

/* Verifies a sysapp ticket */
int verify_sysapp_ticket(BbContentMetaDataHead *cmd)
{
    BbCertBase *certs[3];
    BbShaHash hash_data;
    SHA1Context sha;
    BbAppLaunchCrls *rlBundle;

    certs[0] = (BbCertBase *)((u8*) cmd + sizeof(BbContentMetaDataHead));
    certs[1] = (BbCertBase *)((u8 *)certs[0] + sizeof(BbRsaCert));
    certs[2] = NULL;

    /* verify rl and insure sysapp bundle is OK */
    rlBundle = (BbAppLaunchCrls *)(((u8 *)certs[1]) + sizeof(BbRsaCert));
    BOOT_PATCH_RL(carl, cmd);
    BOOT_PATCH_RL(cprl, cmd);

    /* calculate hash data here */
    SHA1Reset(&sha);
    SHA1Input(&sha, (u8 *) cmd, BB_CMD_HEAD_SIGNED_BYTES);
    SHA1Result(&sha, (u8 *) hash_data);

    PRINTF("Verify ticket: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));

    return skVerifyHash(hash_data, (BbGenericSig *) &cmd->contentMetaDataSign,
                        certs, rlBundle);
}

/* * * * * * *  LOADING BLOCKS * * * * * * * * * * * */

osBbLoadBlocks(BbContentMetaDataHead* cmd,
               u16 *blockList, s32 listSize, u32 addr)
{
    s32 rv = 0;
    OSMesgQueue dmaMessageQ;
    OSMesg dmaMessageBuf;
    OSIoMesg dmaMesg;
    OSPiHandle* handler;
    s32 fileSize = listSize * BB_FL_BLOCK_SIZE;

    osCreateMesgQueue(&dmaMessageQ, &dmaMessageBuf, 1);

    if (listSize >= 0) {
        /* set up atb mapping */
        blockList[listSize] = 0;
        if ((rv = osBbAtbSetup(PI_DOM1_ADDR2, blockList, listSize+1)) < 0) {
            PRINTF("osBbLoadBlocks: Error during osBbAtbSetup %d\n", rv);
            return 0;
        }
    }
    handler = osCartRomInit();
    IO_WRITE(PI_FLASH_CTRL_REG, PI_FLASH_CTRL_RDPH |
                                (0xf << PI_FLASH_CTRL_ADPH_SHIFT) |
                                PI_FLASH_CTRL_WRDY |
                                PI_FLASH_CTRL_ECC |
                                0x3ff);
    //osEPiReadIo(handler, 0x8, &addr);
    dmaMesg.hdr.pri      = OS_MESG_PRI_NORMAL;
    dmaMesg.hdr.retQueue = &dmaMessageQ;
    dmaMesg.dramAddr     = (char*)addr;
    dmaMesg.devAddr      = 0;
    dmaMesg.size         = fileSize;

    osEPiStartDma(handler, &dmaMesg, OS_READ);
    (void)osRecvMesg(&dmaMessageQ, NULL, OS_MESG_BLOCK);

    osWritebackDCacheAll();
    osInvalICache((void*)K0BASE, 64*1024);

    /* Initialize to default values: What are these for? */
    osRomBase = (void*) PHYS_TO_K1(PI_DOM1_ADDR2);
    osMemSize = 0x400000;
    osTvType = OS_TV_NTSC;

    /* Verify hash */
    if (cmd != NULL) {
        if (!verify_hash(cmd->hash, (char*) addr, cmd->size)) {
            PRINTF("osBbLoadBlocks: Hash mismatch\n");
            return 0;
        }
    }

    return addr;
}

/* 
 * Load the GZIPPED blocks into memory at DRAM_TARGET (verifies hash),
 * and unzips the blocks to GUNZIP_TARGET.  Finally copies the
 * application to the memory specified in the ROM segment.
 */
u32 osBbLoadGzippedBlocks(BbContentMetaDataHead* cmd,
                          u16 *blockList, s32 listSize, u32 loadAll)
{
    u32 addr;
    u32 outsize;

    PRINTF("Load blocks: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    if (!osBbLoadBlocks(cmd, blockList, listSize, DRAM_TARGET))
    {
        PRINTF("osBbLoadGzippedBlocks: Error loading blocks\n");
        return 0;
    }

    PRINTF("Gunzip blocks: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    outsize = expand_gzip((char*) DRAM_TARGET, (char*) GUNZIP_TARGET,
                          cmd->size, MAX_SYSAPP_SIZE);
    if (outsize < ROM_HEADER_SIZE) {
        PRINTF("osBbLoadGzippedBlocks: GZIP failed (outsize = %u)\n", outsize);
        return 0;
    }

    addr = * ((u32*) (GUNZIP_TARGET + N64_ROM_LOADADDR_OFFSET ));
    PRINTF("osBbLoadGzippedBlocks: outsize = %u, addr = 0x%08x\n", outsize, addr);

    PRINTF("Copy blocks: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    if (loadAll) {
        if ((addr - ROM_HEADER_SIZE) < SA1_SEGMENT_END) {
            PRINTF("osBbLoadGzippedBlocks: Cannot load SA2 into address 0x%08x\n",
                   addr - ROM_HEADER_SIZE);
            PRINTF("SA2 must start after SA1 which is from 0x%08x to 0x%08x\n",
                   SA1_SEGMENT_START, SA1_SEGMENT_END);
            return 0;
        }

        memcpy((void*) (addr - ROM_HEADER_SIZE), (void*) GUNZIP_TARGET, outsize);
    }
    else {
        if (addr < SA1_SEGMENT_END) {
            PRINTF("osBbLoadGzippedBlocks: Cannot load SA2 into address 0x%08x\n", addr);
            PRINTF("SA2 must start after SA1 which is from 0x%08x to 0x%08x\n",
                   SA1_SEGMENT_START, SA1_SEGMENT_END);
            return 0;
        }

        memcpy((void*) (addr), (void*) (GUNZIP_TARGET + ROM_HEADER_SIZE) ,
               outsize - ROM_HEADER_SIZE);
    }

    return addr;
}

/* * * * * * *  LOADING SYSAPP TICKETS * * * * * * * * * * * */

/* loadSysappTicket:
 *   read and verify ticket. 
 *
 * startNandBlock:
 *   out - nand block holding first content block.
 *
 * return:
 *   success, or error codes for invalid ticket and problems 
 *   reading from flash.
 */
int loadSysappTicket(u32 *startNandBlock, u32 ticketBlock)
{
    int i,j,ret;
    BbContentMetaDataHead *cmd =
        (BbContentMetaDataHead *)CMDBUNDLE_STORAGE_START;
    
    u32 *p32 = (u32 *)cmd;
    u32 ticketStartBlock;

    /* skip SK blocks on nand */

    if( (ret = findGoodBlock(*startNandBlock, ticketBlock))<0)
        return ret;
    ticketStartBlock = (u32)ret;

    /* copy entire block into tik_data */
    for(i=0;i<PI_FLASH_PAGES_PER_BLOCK;i++){
        /* read page into pi buffer 0. */
        ret = readFlashPage(ticketStartBlock*PI_FLASH_PAGES_PER_BLOCK+i,0);
        if(ret<0)
            return ret;
        if(i==0){
            *startNandBlock = getBlockLink(IO_READ(PI_BUFFER_0_OOB_START+
                                                   BB_FL_BLOCK_LINK_OFF));
        }
        /* copy to buffer */
        for(j=0;j<PI_FLASH_PAGE_DATA_SIZE;j+=4){
            *(p32++) = IO_READ(PI_BUFFER_BASE_REG+j);
        }
    }

    PRINTF("Ticket Start Block = %u\n", ticketStartBlock);
    PRINTF("SysApp Content Size = %u\n", cmd->size);

    /* verify chain (sysapp case assumes proper certs in place) */
    PRINTF("Checking certs\n");
    if ((ret = verify_sysapp_ticket(cmd)) == SK_API_SUCCESS) {
        PRINTF("verify_sysapp_ticket passed\n");
    }
    else {
        PRINTF("verify_sysapp_ticket failed: %d\n", ret);
        return SK_FAIL;
    }

    return SK_SUCCESS;
}

/* * * * * * *  LOADING SA2 * * * * * * * * * * * */

/* 
 * Derived from osBbExecApp
 * Execute the application at addr with the given entryCode
 * Do not set any special registers.
 */
s32
osBbExecSA(u32 addr, u32 entryCode)
{
    int (*f)(u32);
    f = (void*)addr;
    __osDisableInt();
    return (*f)(entryCode);
}

/* List of blocks for the SA - SA is at most 1MB = 64 16K blocks 
   (actually less) */
#define MAX_SYSAPP_BLOCKS BB_FL_BYTE_TO_BLOCK(MAX_SYSAPP_SIZE)
static u16 blockList[MAX_SYSAPP_BLOCKS+1];

u32 osBbLoadSA2()
{
    u32 loadAddr = 0;
    int ret;
    u32 i;
    BbContentMetaDataHead *cmd=
        (BbContentMetaDataHead *) CMDBUNDLE_STORAGE_START;
    u32 SA1size, SA2size, SA1blocks, SA2blocks;
    u32 SA1fb=0, SA2fb=0; /* Starting blocks of SA1 and SA2 */

    /* load SA1 ticket block */
    PRINTF("Load T1: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    if ((ret = loadSysappTicket(&SA1fb, SK_BLOCKS)) < 0) {
        PRINTF("Error loading T1: %d\n", ret);
        return 0;
    }
    SA1size = cmd->size;
    SA1blocks = BB_FL_BYTE_TO_BLOCK(SA1size);

    SA2fb = SA1fb;
#if _USE_BLOCK_LINKS
    /* Follow block links to where T2 starts */
    for (i=1; i <= SA1blocks; i++) {
        ret = readFlashPage(SA2fb*PI_FLASH_PAGES_PER_BLOCK, 0);
        if( ret < 0 ) {
            PRINTF("Error getting link for SA1 block %d from %d, ret = %d",
                   i, SA2fb, ret);
            return 0;
        }
        SA2fb = getBlockLink(IO_READ(PI_BUFFER_0_OOB_START+
                                     BB_FL_BLOCK_LINK_OFF));
    }

    /* Load SA2 ticket block */
    PRINTF("Load T2: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    /* Flash block for T2 should be after SA1 */
    /* Flash block for SA2 should be after T2 */
    if ((ret = loadSysappTicket(&SA2fb, 0)) < 0) {
#else
    if ((ret = loadSysappTicket(&SA2fb, SA1blocks)) < 0) {
#endif
        PRINTF("Error loading T2: %d\n", ret);
        return 0;
    }

    SA2size = cmd->size;
    SA2blocks = BB_FL_BYTE_TO_BLOCK(SA2size);
    if (SA2blocks > MAX_SYSAPP_BLOCKS) {
        PRINTF("SA2 too big: Occupies %d blocks, max %d blocks\n",
               SA2blocks, MAX_SYSAPP_BLOCKS);
        return 0;
    }

    PRINTF("SA1 content size = %u bytes\n", SA1size);
    PRINTF("SA2 content size = %u bytes\n", SA2size);
    PRINTF("SA1 flash block = %u\n", SA1fb);
    PRINTF("SA2 flash block = %u\n", SA2fb);
    PRINTF("Setup SA2 blocks: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));
    /* Setup the block list and the launch information */
    blockList[0] = SA2fb;
    for (i=1; i < SA2blocks; i++) {
#if _USE_BLOCK_LINKS
        ret = readFlashPage(blockList[i-1]*PI_FLASH_PAGES_PER_BLOCK, 0);
        if( ret < 0 ) {
            PRINTF("Error getting link for SA2 block %d from %d, ret = %d",
                   i, blockList[i-1], ret);
            return 0;
        }
        blockList[i] = getBlockLink(IO_READ(PI_BUFFER_0_OOB_START+
                                            BB_FL_BLOCK_LINK_OFF));
#else
        blockList[i] = findGoodBlock(blockList[i-1], 1);
#endif
    }
    blockList[i] = 0;
    
    PRINTF("osBbLoadApp...");
    if ((loadAddr = osBbLoadGzippedBlocks(cmd, blockList, SA2blocks, 1)) != NULL) {
        PRINTF("Finished loading: %llu\n", OS_CYCLES_TO_NSEC(osGetTime()));

        PRINTF("succeeded - load address = 0x%08x\n", loadAddr);

    }
    else {
        PRINTF("failed\n");
        return 0;
    }

    return loadAddr;
}