/*
 * Copyright (C) 2020-2021 Aloysius Indrayanto
 *
 * This file is part of the eCxx library, see LICENSE file for the license details.
 */

/*
 * ----------------------------------------------------------------------------------
 * The code in this file is copied (with a slight modifications) from:
 *     flash_hal.h & flash_hal.cpp - API for accessing raw flash for filesystems
 *     Esp.cpp                     - esp8266-specific APIs
 * ----------------------------------------------------------------------------------
 * Copyright (C) 2015 Ivan Grokhotkov. All rights reserved.
 *
 * This file is part of the esp8266 core for Arduino environment.
 *
 * This code was influenced by NodeMCU and Sming libraries, and first version of
 * Arduino wrapper written by Hristo Gochkov.
 *
 * This library is free software; you can redistribute it and/or modify it under the
 * terms of the GNU Lesser General Public License as published by the Free Software
 * Foundation; either version 2.1 of the License, or (at your option) any later
 * version.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along
 * with this library; if not, write to the Free Software Foundation, Inc., 51
 * Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 * ----------------------------------------------------------------------------------
 */


#include "CoreSys.h"


//
// Uncomment this to enable flash HAL debugging
// (the system-wide debugging macro must be enabled for it to have any effect)
//
// #define FLASH_HAL_DEBUG(FMT, ...) SYSTEM_DEBUG(FMT, ## __VA_ARGS__)

#ifndef FLASH_HAL_DEBUG
    #define NO_FLASH_HAL_DEBUG
    #define FLASH_HAL_DEBUG(FMT, ...) SYSTEM_DEBUG_NOP
#endif


namespace eCxx {

namespace NodeMCU {


//
// Flash information utilities
//
extern FlashChipMode magicFlashChipMode(uint8_t byte)
{
    const FlashChipMode mode = static_cast<FlashChipMode>(byte);
    return (mode > FCM_DOUT) ? FCM_UNKNOWN : mode;
}


extern uint32_t magicFlashChipSize(uint8_t byte)
{
    switch(byte & 0x0F) {
        case 0x0 : return 512UL * 1024UL;          //   4 Mbit (512KB)
        case 0x1 : return 256UL * 1024UL;          //   2 Mbit (256KB)
        case 0x2 : return   1UL * 1024UL * 1024UL; //   8 Mbit (  1MB)
        case 0x3 : return   2UL * 1024UL * 1024UL; //  16 Mbit (  2MB)
        case 0x4 : return   4UL * 1024UL * 1024UL; //  32 Mbit (  4MB)
        case 0x8 : return   8UL * 1024UL * 1024UL; //  64 Mbit (  8MB)
        case 0x9 : return  16UL * 1024UL * 1024UL; // 128 Mbit ( 16MB)
        default  : return   0UL;                   // Error
    }
}


extern uint32_t magicFlashChipSpeed(uint8_t byte)
{
    switch(byte & 0x0F) {
        case 0x0 : return 40UL * 1000000UL; // 40 MHz
        case 0x1 : return 26UL * 1000000UL; // 26 MHz
        case 0x2 : return 20UL * 1000000UL; // 20 MHz
        case 0xF : return 80UL * 1000000UL; // 80 MHz
        default  : return 0;
    }
}


extern uint32_t getFlashChipId()
{
    static uint32_t flashChipID = 0;
    if(flashChipID == 0) flashChipID = spi_flash_get_id();

    return flashChipID;
}


extern uint8_t getFlashChipVendorId()
{ return getFlashChipId() & 0x000000FF; }


extern uint32_t getFlashChipRealSize()
{ return 1UL << ( ( spi_flash_get_id() >> 16 ) & 0xFF ); }


extern uint32_t getFlashChipSize()
{
    uint32_t data;
    uint8_t* bytes = reinterpret_cast<uint8_t*>(&data);

    // Read the first 4 bytes (magic byte + flash configuration)
    if( spi_flash_read(0x0000, &data, 4) == SPI_FLASH_RESULT_OK ) {
        return magicFlashChipSize( ( bytes[3] & 0xF0 ) >> 4 );
    }

    return 0;
}


extern uint32_t getFlashChipSpeed()
{
    uint32_t data;
    uint8_t* bytes = reinterpret_cast<uint8_t*>(&data);

    // Read the first 4 bytes (magic byte + flash configuration)
    if( spi_flash_read(0x0000, &data, 4) == SPI_FLASH_RESULT_OK ) {
        return magicFlashChipSpeed(bytes[3] & 0x0F);
    }

    return 0;
}


extern FlashChipMode getFlashChipMode()
{
    FlashChipMode mode = FCM_UNKNOWN;

    uint32_t data;
    uint8_t* bytes = reinterpret_cast<uint8_t*>(&data);

    // Read the first 4 bytes (magic byte + flash configuration)
    if( spi_flash_read(0x0000, &data, 4) == SPI_FLASH_RESULT_OK ) {
        mode = magicFlashChipMode(bytes[2]);
    }

    return mode;
}


extern bool checkFlashConfig(bool needsEquals)
{
    if(needsEquals) {
        if(getFlashChipRealSize() == getFlashChipSize()) return true;
    }
    else {
        if(getFlashChipRealSize() >= getFlashChipSize()) return true;
    }
    return false;
}


extern uint32_t getFlashUserProgramSize()
{
    static uint32_t result = 0;
    if(result) return result;

    uint32_t       pos = APP_START_OFFSET;
    image_header_t image_header;

    if( spi_flash_read(pos, reinterpret_cast<uint32_t*>(&image_header), sizeof(image_header)) != SPI_FLASH_RESULT_OK ) {
        return 0;
    }
    pos += sizeof(image_header);

    for(uint32_t secIdx = 0; secIdx < image_header.num_segments; ++secIdx)
    {
        section_header_t section_header = { 0, 0 };
        if( spi_flash_read(pos, reinterpret_cast<uint32_t*>(&section_header), sizeof(section_header) ) != SPI_FLASH_RESULT_OK ) {
            return 0;
        }

        pos += sizeof(section_header);
        pos += section_header.size;
    }

    result = (pos + 16) & ~15;
    return result;
}


extern uint32_t getFlashUserProgramFreeSpace()
{
    const uint32_t usedSize = getFlashUserProgramSize();

    // Round one sector up
    const uint32_t freeSpaceStart = ( usedSize + FLASH_SECTOR_SIZE - 1 ) & ( ~( FLASH_SECTOR_SIZE - 1 ) );
    const uint32_t freeSpaceEnd   = reinterpret_cast<uint32_t>(&_FS_start) - 0x40200000UL;

    return freeSpaceEnd - freeSpaceStart;
}


//
// Flash read & write utilities
//
extern bool spiFlashEraseSector(uint32_t sector)
{ return spi_flash_erase_sector(sector) == 0; }


#ifndef ECXX_NODEMCU_DISABLE_PUYA_SUPPORT

static SpiFlashOpResult spiFlashWritePuya(uint32_t offset, const uint32_t* data, size_t size)
{
    if(data == nullptr) return SPI_FLASH_RESULT_ERR;
    if(size % 4       ) return SPI_FLASH_RESULT_ERR;

    // Allocate buffer
    static uint32_t* flashWritePuyaBuf = nullptr;

    if(flashWritePuyaBuf == nullptr) {
        // There is no need to free this memeory because the flash chip will never change at runtime
        flashWritePuyaBuf = reinterpret_cast<uint32_t*>( malloc(FLASH_PAGE_SIZE) );
        if(flashWritePuyaBuf == nullptr) return SPI_FLASH_RESULT_ERR;
    }

    // PUYA flash chips need to read existing data, update in memory, and then write modified data
    SpiFlashOpResult rc        = SPI_FLASH_RESULT_OK;
    const uint32_t*  ptr       = data;
    uint32_t         pos       = offset;
    size_t           bytesLeft = size;

    while(bytesLeft > 0 && rc == SPI_FLASH_RESULT_OK) {
        size_t bytesNow = bytesLeft;

        if(bytesNow > FLASH_PAGE_SIZE) {
            bytesNow   = FLASH_PAGE_SIZE;
            bytesLeft -= FLASH_PAGE_SIZE;
        }
        else {
            bytesLeft = 0;
        }

        rc = spi_flash_read(pos, flashWritePuyaBuf, bytesNow);
        if(rc != SPI_FLASH_RESULT_OK) return rc;

        for(size_t i = 0; i < bytesNow / 4; ++i) {
            flashWritePuyaBuf[i] &= *ptr;
            ++ptr;
        }

        rc   = spi_flash_write(pos, flashWritePuyaBuf, bytesNow);
        pos += bytesNow;
    }

    return rc;
}

#endif


static bool spiFlashReplaceBlock(uint32_t address, const uint8_t* value, uint32_t byteCount)
{
    const uint32_t alignedAddress  = (address & ~3);
    const uint32_t alignmentOffset = address - alignedAddress;

    if( alignedAddress != ( (address + byteCount - 1 ) & ~3) ) {
        // Only one 4 byte block is supported
        return false;
    }

#ifndef ECXX_NODEMCU_DISABLE_PUYA_SUPPORT
    if( getFlashChipVendorId() == SPI_FLASH_VENDOR_PUYA ) {
        uint8_t tempData[4] __ECXX_ATTR_ALIGNED_4;
        if( spi_flash_read(alignedAddress, reinterpret_cast<uint32_t*>(tempData), 4) != SPI_FLASH_RESULT_OK ) {
            return false;
        }
        for(size_t i = 0; i < byteCount; i++) {
            tempData[i + alignmentOffset] &= value[i];
        }
        if( spi_flash_write(alignedAddress, reinterpret_cast<uint32_t*>(tempData), 4) != SPI_FLASH_RESULT_OK ) {
            return false;
        }
    }
    else
#endif
    {
        uint32_t tempData;
        if( spi_flash_read(alignedAddress, &tempData, 4) != SPI_FLASH_RESULT_OK ) {
            return false;
        }
        memcpy( reinterpret_cast<uint8_t*>(&tempData) + alignmentOffset, value, byteCount);
        if( spi_flash_write(alignedAddress, &tempData, 4) != SPI_FLASH_RESULT_OK ) {
            return false;
        }
    }

    return true;
}


static size_t spiFlashWriteUnalignedMemory(uint32_t address, const uint8_t* data, size_t size)
{
    size_t sizeLeft      = (size & ~3);
    size_t currentOffset = 0;

    // Memory is unaligned, so copy it to an aligned buffer
    uint32_t alignedData[ FLASH_PAGE_SIZE / sizeof(uint32_t) ] __ECXX_ATTR_ALIGNED_4;

    // Handle page boundary
    const bool pageBreak = (address % 4) && ( (address / FLASH_PAGE_SIZE) != ( (address + sizeLeft - 1) / FLASH_PAGE_SIZE ) );

    if(pageBreak) {
        const size_t byteCount = 4 - (address % 4);
        if( !spiFlashReplaceBlock(address, data, byteCount) ) {
            return 0;
        }
        // It now has aligned address, so it can cross page boundaries
        currentOffset += byteCount;
        // Realign size to 4
        sizeLeft = (size - byteCount) & ~3;
    }

    while(sizeLeft) {
        const size_t willCopy = eCxx::min( sizeLeft, sizeof(alignedData) );
        memcpy(alignedData, data + currentOffset, willCopy);
        // Use aligned write because now address, data, and size are aligned to 4 bytes
        if( !spiFlashWrite(address + currentOffset, alignedData, willCopy) ) {
            return 0;
        }
        sizeLeft      -= willCopy;
        currentOffset += willCopy;
    }

    return currentOffset;
}


static bool spiFlashWritePageBreak(uint32_t address, const uint8_t* data, size_t size)
{
    if(size > 4) return false;

    const size_t pageLeft = FLASH_PAGE_SIZE - (address % FLASH_PAGE_SIZE);
          size_t offset   = 0;
          size_t sizeLeft = size;

    if(pageLeft > 3) return false;

    if( !spiFlashReplaceBlock(address, data, pageLeft) ) {
        return false;
    }

    offset   += pageLeft;
    sizeLeft -= pageLeft;

    // The last 4-byte block of the page has been replaced, now write the remainder in the next page(s)
    if( !spiFlashReplaceBlock(address + offset, data + offset, sizeLeft) ) {
        return false;
    }

    return true;
}


extern bool spiFlashWrite(uint32_t address, const uint32_t* data, size_t size)
{
    SpiFlashOpResult rc = SPI_FLASH_RESULT_OK;

    const bool pageBreak = ( (address % 4) && (address / FLASH_PAGE_SIZE) != ( (address + size - 1) / FLASH_PAGE_SIZE ) );

    if( ( reinterpret_cast<uintptr_t>(data) % 4 ) || (size % 4) || pageBreak ) {
        return false;
    }

#ifndef ECXX_NODEMCU_DISABLE_PUYA_SUPPORT
    if( getFlashChipVendorId() == SPI_FLASH_VENDOR_PUYA ) {
        rc = spiFlashWritePuya(address, const_cast<uint32_t*>(data), size);
    }
    else
#endif
    {
        rc = spi_flash_write(address, const_cast<uint32_t*>(data), size);
    }

    return rc == SPI_FLASH_RESULT_OK;
}


extern bool spiFlashWrite(uint32_t address, const uint8_t* data, size_t size)
{
    if(!size) return true;

    size_t sizeLeft      = size & ~3;
    size_t currentOffset = 0;

    if(sizeLeft) {
        if( reinterpret_cast<uintptr_t>(data) % 4 ) {
            const size_t written = spiFlashWriteUnalignedMemory(address, data, size);
            if(!written) {
                return false;
            }
            currentOffset += written;
            sizeLeft      -= written;
        }
        else {
            const bool pageBreak = ( (address % 4) && (address / FLASH_PAGE_SIZE) != ( (address + sizeLeft - 1) / FLASH_PAGE_SIZE ) );
            if(pageBreak) {
                while(sizeLeft) {
                    // It cannot cross the page boundary and because the write must be 4-byte aligned,
                    // calculate the maximum amount of bytes that can be written
                    size_t pageBoundary = ( FLASH_PAGE_SIZE - ( (address + currentOffset) % FLASH_PAGE_SIZE ) ) & ~3;
                    if(sizeLeft > pageBoundary) {
                        // Aligned write up to page boundary
                        if( !spiFlashWrite(address + currentOffset, reinterpret_cast<const uint32_t*>(data + currentOffset), pageBoundary) ) {
                            return false;
                        }
                        currentOffset += pageBoundary;
                        sizeLeft      -= pageBoundary;
                        // Cross the page boundary
                        if( !spiFlashWritePageBreak(address + currentOffset, data + currentOffset, 4) ) {
                            return false;
                        }
                        currentOffset += 4;
                        sizeLeft      -= 4;
                    }
                    else {
                        // It does not cross page boundary
                        if( !spiFlashWrite(address + currentOffset, reinterpret_cast<const uint32_t*>(data + currentOffset), sizeLeft) ) {
                            return false;
                        }
                        currentOffset += sizeLeft;
                        sizeLeft       = 0;
                    }
                }
            }
             else {
                // Pointer is properly aligned and write does not cross the page boundary, therefore, use aligned write
                if( !spiFlashWrite(address, reinterpret_cast<const uint32_t*>(data), sizeLeft) ) {
                    return false;
                }
                currentOffset = sizeLeft;
                sizeLeft      = 0;
            }
        }
    }

    sizeLeft = size - currentOffset;

    if(sizeLeft > 0) {
        // The size was not aligned, therefore, there are some bytes left to write
        // (also need to recheck for page boundary crossing)
        const bool pageBreak = ( (address % 4) && (address / FLASH_PAGE_SIZE) != ( (address + sizeLeft - 1) / FLASH_PAGE_SIZE ) );
        if(pageBreak) {
            // Cross the page boundary
            if( !spiFlashWritePageBreak(address + currentOffset, data + currentOffset, sizeLeft) ) {
                return false;
            }
        }
        else {
            // Write partial block
            spiFlashReplaceBlock(address + currentOffset, data + currentOffset, sizeLeft);
        }
    }

    return true;
}


extern bool spiFlashRead(uint32_t address, uint32_t* data, size_t size)
{
    if( ( reinterpret_cast<uintptr_t>(data) % 4 ) || (size % 4) ) {
        return false;
    }

    return spi_flash_read(address, data, size) == SPI_FLASH_RESULT_OK;
}


extern bool spiFlashRead(uint32_t address, uint8_t* data, size_t size)
{
    const size_t sizeAligned   = size & ~3;
          size_t currentOffset = 0;

    if( reinterpret_cast<uintptr_t>(data) % 4 ) {
        uint32_t alignedData[ FLASH_PAGE_SIZE / sizeof(uint32_t) ] __ECXX_ATTR_ALIGNED_4;
        size_t   sizeLeft = sizeAligned;
        while(sizeLeft) {
            const size_t willCopy = eCxx::min(sizeLeft, sizeof(alignedData));
            // Read to aligned buffer and then copy to data
            if( !spiFlashRead(address + currentOffset, alignedData, willCopy) ) {
                return false;
            }
            memcpy(data + currentOffset, alignedData, willCopy);
            sizeLeft      -= willCopy;
            currentOffset += willCopy;
        }
    }
    else {
        // Pointer is properly aligned, use aligned read
        if( !spiFlashRead(address, reinterpret_cast<uint32_t*>(data), sizeAligned) ) {
            return false;
        }
        currentOffset = sizeAligned;
    }

    if(currentOffset < size) {
        uint32_t tempData;
        if( spi_flash_read(address + currentOffset, &tempData, 4) != SPI_FLASH_RESULT_OK ) {
            return false;
        }
        memcpy( reinterpret_cast<uint8_t*>(data) + currentOffset, &tempData, size - currentOffset );
    }

    return true;
}


////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////


extern int32_t flashHalRead(uint32_t addr, uint32_t size, uint8_t* dst)
{
    Application::yield();

    if( spiFlashRead(addr, dst, size) ) {
        return FLASH_HAL_OK;
    }
    else {
        return FLASH_HAL_READ_ERROR;
    }
}


extern int32_t flashHalWrite(uint32_t addr, uint32_t size, const uint8_t* src)
{
    Application::yield();

    if( spiFlashWrite(addr, src, size) ) {
        return FLASH_HAL_OK;
    }
    else {
        return FLASH_HAL_WRITE_ERROR;
    }
}


extern int32_t flashHalErase(uint32_t addr, uint32_t size)
{
    if( (size & (SPI_FLASH_SEC_SIZE - 1)) != 0 || (addr & (SPI_FLASH_SEC_SIZE - 1)) != 0 ) {
        FLASH_HAL_DEBUG("flashHalErase called with addr=%x, size=%d\r\n", addr, size);
        abort();
    }

    const uint32_t sector      = addr / SPI_FLASH_SEC_SIZE;
    const uint32_t sectorCount = size / SPI_FLASH_SEC_SIZE;

    for(uint32_t i = 0; i < sectorCount; ++i) {
        Application::yield();

        if( !spiFlashEraseSector(sector + i) ) {
            FLASH_HAL_DEBUG("flashHalErase addr=%x size=%d i=%d\r\n", addr, size, i);
            return FLASH_HAL_ERASE_ERROR;
        }
    }

    return FLASH_HAL_OK;
}


} // namespace NodeMCU

} // namespace eCxx


//
// Function aliases with "C++" linkage for compatibility with the Arduino library part
//

// make mangle CODE='namespace eCxx { namespace NodeMCU { int flashHalRead(unsigned int addr, unsigned int size, unsigned char* dst) {} } }'
extern eCxx::int32_t flash_hal_read(eCxx::uint32_t addr, eCxx::uint32_t size, eCxx::uint8_t* dst) __ECXX_ATTR_ALIAS("_ZN4eCxx7NodeMCU12flashHalReadEjjPh");

// make mangle CODE='namespace eCxx { namespace NodeMCU { int flashHalWrite(unsigned int addr, unsigned int size, const unsigned char* src) {} } }'
extern eCxx::int32_t flash_hal_write(eCxx::uint32_t addr, eCxx::uint32_t size, const eCxx::uint8_t* src) __ECXX_ATTR_ALIAS("_ZN4eCxx7NodeMCU13flashHalWriteEjjPKh");

// make mangle CODE='namespace eCxx { namespace NodeMCU { int flashHalErase(unsigned int addr, unsigned int size) {} } }'
extern eCxx::int32_t flash_hal_erase(eCxx::uint32_t addr, eCxx::uint32_t size) __ECXX_ATTR_ALIAS("_ZN4eCxx7NodeMCU13flashHalEraseEjj");
