/*****************************************************************************
 *                                                                           *
 * Copyright (c) 2006, 2010 Rowley Associates Limited.                       *
 *                                                                           *
 * This file may be distributed under the terms of the License Agreement     *
 * provided with this software.                                              *
 *                                                                           *
 * THIS FILE IS PROVIDED AS IS WITH NO WARRANTY OF ANY KIND, INCLUDING THE   *
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. *
 *****************************************************************************/

#include <libmem.h>
#include <libmem_loader.h>
#include <string.h>
#include "spi0df.h"

extern unsigned secondary_loader_begin, secondary_loader_end;
extern unsigned __SDRAM_segment_used_end__, __SDRAM_segment_end__;

unsigned SDRAM_START;
unsigned char current_block[2112];
unsigned current_block_size;
unsigned current_block_address;
unsigned first_application_flash_page;
unsigned secondary_loader_size;
unsigned application_size;
libmem_geometry_t geometry[2];
libmem_driver_functions_t driver_functions;
libmem_ext_driver_functions_t ext_driver_functions;
libmem_driver_handle_t libmem_handle;  

static int
dataflash_erase_impl(libmem_driver_handle_t *h, uint8_t *start, size_t size,
                    uint8_t **erase_start, size_t *erase_size)
{   
  unsigned erase_start_flash_page = 0;
  unsigned flash_page;
  if ((unsigned)start == SDRAM_START)  // erase bootloader
    {
      for (flash_page=0; flash_page < first_application_flash_page; flash_page++)
        SPI0DataFlashErasePage(flash_page);    
    }
  while (size)
    {
      flash_page = first_application_flash_page + (start-(uint8_t*)SDRAM_START)/SPI0DataFlashPageSize;
      if (erase_start_flash_page == 0)
        erase_start_flash_page = flash_page;
      SPI0DataFlashErasePage(flash_page);
      if (size > SPI0DataFlashPageSize)
        {                           
          start += SPI0DataFlashPageSize;
          size -= SPI0DataFlashPageSize;
        }
      else
        {                  
          size = 0;
        }
    }  
  if (erase_start)
    *erase_start = (uint8_t*)SDRAM_START+(erase_start_flash_page*SPI0DataFlashPageSize);
  if (erase_size)
    *erase_size = ((flash_page-erase_start_flash_page)+1)*SPI0DataFlashPageSize;
  return LIBMEM_STATUS_SUCCESS;
}

static void
write_current_block()
{
  unsigned flash_page = first_application_flash_page+(current_block_address-SDRAM_START)/SPI0DataFlashPageSize;  
  SPI0DataFlashWritePage(flash_page, current_block);     
  application_size = (current_block_address-SDRAM_START)+SPI0DataFlashPageSize;
  current_block_address = current_block_size = 0;
}

static int
dataflash_write_impl(libmem_driver_handle_t *h, uint8_t *dest, const uint8_t *src, size_t size)
{
  while (size)
    {
      if (current_block_address+current_block_size==(unsigned)dest)
        {
          unsigned s;
          if (size < (SPI0DataFlashPageSize-current_block_size))
            s = size;
          else
            s = SPI0DataFlashPageSize-current_block_size;
          memcpy(current_block+current_block_size, src, s);
          current_block_size += s;
          dest += s;
          src += s;
          size -= s;          
          if (current_block_size == SPI0DataFlashPageSize)
            write_current_block();
        }
      else if (current_block_address)
        write_current_block();
      else        
        current_block_address = (unsigned)dest;
    }
  return LIBMEM_STATUS_SUCCESS;
}

static int
dataflash_flush_impl(libmem_driver_handle_t *h)
{
  if (current_block_size)
    write_current_block();
  return LIBMEM_STATUS_SUCCESS;
}

static int 
dataflash_read_impl(libmem_driver_handle_t *h, uint8_t *dest, const uint8_t *src, size_t size)
{
  unsigned address = src-(uint8_t*)SDRAM_START;
  unsigned flash_page = first_application_flash_page + (address / SPI0DataFlashPageSize);
  unsigned offset = address % SPI0DataFlashPageSize;
  SPI0DataFlashReadBytes(flash_page, offset, size, dest); 
  return LIBMEM_STATUS_SUCCESS;
}

static uint32_t
dataflash_crc32_impl(libmem_driver_handle_t *h, const uint8_t *start, size_t size, uint32_t crc)
{
  unsigned address = start-(uint8_t*)SDRAM_START;
  while (size)
    {      
      unsigned flash_page = first_application_flash_page + (address / SPI0DataFlashPageSize);
      unsigned offset = address % SPI0DataFlashPageSize;
      unsigned s;
      if (size > sizeof(current_block))
        s = sizeof(current_block);
      else
        s = size;          
      SPI0DataFlashReadBytes(flash_page, offset, s, current_block);
      crc = libmem_crc32_direct(current_block, s, crc);
      size -= s;
      address += s;      
    }
  return crc;
}

int
main(unsigned flags, unsigned parameter)
{
  int res;
  if (flags & LIBMEM_RPC_LOADER_FLAG_PARAM)
    SDRAM_START = parameter;
  else
#ifdef AT91SAM9G45_EK
    SDRAM_START = 0x70000000+0x4000; // reserve 0x4000 bytes for loader/application TLB area
#else
    SDRAM_START = 0x20000000+0x4000; // reserve 0x4000 bytes for loader/application TLB area
#endif

  SPI0DataFlashInit();  
  if (!SPI0DataFlashIdentify())
    libmem_rpc_loader_exit(LIBMEM_STATUS_ERROR, "Failed to identify SPI0DataFlash");

  secondary_loader_size = (unsigned char *)&secondary_loader_end - (unsigned char *)&secondary_loader_begin;
#ifdef AT91RM9200
  unsigned numblocks = (secondary_loader_size/512)+((secondary_loader_size%512)?1:0);
  unsigned numpages = ((numblocks*512)/SPI0DataFlashPageSize)+(((numblocks*512)%SPI0DataFlashPageSize)?1:0);
  unsigned n;
  for (n=0; (1<<n) < numpages; n++);
  *(volatile unsigned *)((&secondary_loader_begin)+5) = (SPI0DataFlashPageSize<<17) | (n<<13) | numblocks;
#else
  *(volatile unsigned *)((&secondary_loader_begin)+5) = secondary_loader_size;    
#endif   

  for (first_application_flash_page = 0; (first_application_flash_page * SPI0DataFlashPageSize) <= secondary_loader_size; first_application_flash_page++); 
  application_size = 0;

  geometry[0].count = SPI0DataFlashNumPages;
  geometry[0].size = SPI0DataFlashPageSize;

  driver_functions.write = dataflash_write_impl;
  driver_functions.fill = 0;      
  driver_functions.erase = dataflash_erase_impl;   
  driver_functions.lock = 0; 
  driver_functions.unlock = 0;  
  driver_functions.flush = dataflash_flush_impl;  

  ext_driver_functions.inrange = 0;
  ext_driver_functions.read = dataflash_read_impl;
  ext_driver_functions.crc32 = dataflash_crc32_impl;
 
  current_block_size = 0;
  current_block_address = 0;

  libmem_register_driver(&libmem_handle, (void*)SDRAM_START, (SPI0DataFlashNumPages-first_application_flash_page)*SPI0DataFlashPageSize, geometry, 0, &driver_functions, &ext_driver_functions);
  res = libmem_rpc_loader_start(&__SDRAM_segment_used_end__, &__SDRAM_segment_end__ - 1);

  // write the secondary boot loader 
  if (res==LIBMEM_STATUS_SUCCESS && application_size) 
    {
      unsigned char *applicationAddress = (unsigned char *)(SDRAM_START);  
      unsigned char *psrc = (unsigned char *)&secondary_loader_begin;   
      *(volatile unsigned *)((&secondary_loader_begin)+8) = application_size;
      *(volatile unsigned *)((&secondary_loader_begin)+9) = SDRAM_START;
      unsigned flash_page = 0;
      while (secondary_loader_size)
        {
          if (secondary_loader_size > SPI0DataFlashPageSize)      
            memcpy(current_block, psrc, SPI0DataFlashPageSize);
          else
            memcpy(current_block, psrc, secondary_loader_size);
          SPI0DataFlashWritePage(flash_page, current_block);
          flash_page++;
          if (secondary_loader_size > SPI0DataFlashPageSize)
            {
              secondary_loader_size -= SPI0DataFlashPageSize;
              psrc += SPI0DataFlashPageSize;            
            }
          else
            secondary_loader_size = 0;
        }     
      application_size = 0;
    }

  libmem_rpc_loader_exit(res, 0);
}