/*
 * ----------------------------------------------------------------------------------
 * The code in this file is copied (with some modifications) from:
 *     Updater.cpp by Ivan Grokhotkov
 * ----------------------------------------------------------------------------------
 * This file is part of the esp8266 core for Arduino environment.
 *
 * 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 <user_interface.h>
#include <eboot/eboot_command.h>

#include "StackThunk.h"
#include "Updater.h"


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

#ifndef UPDATER_DEBUG
    #define NO_UPDATER_DEBUG
    #define UPDATER_DEBUG(FMT, ...) SYSTEM_DEBUG_NOP
#endif


// Optional binary signing verification for OTA programming
#ifdef ECXX_NODEMCU_BINARY_SIGNING_VERIFICATION
namespace esp8266 {
    extern UpdaterHashClass&   updaterSigningHash;
    extern UpdaterVerifyClass& updaterSigningVerifier;
}
#endif


namespace eCxx {

namespace NodeMCU {


//
// Function prototype from "EBootCommand.cpp"
//
extern void ebootCommandWrite(eboot_command* cmd);


//
// Program updater class
//
ProgUpdater::ProgUpdater()
: _async         (false)
, _error         (0)
, _buffer        (0)
, _bufferLen     (0)
, _size          (0)
, _startAddress  (0)
, _currentAddress(0)
, _command       (U_FLASH)
, _hash          (nullptr)
, _verify        (nullptr)
, _progressCB    (nullptr)
{
#ifdef ECXX_NODEMCU_BINARY_SIGNING_VERIFICATION
    installSignature(&esp8266::updaterSigningHash, &esp8266::updaterSigningVerifier);
    stack_thunk_add_ref();
#endif
}


ProgUpdater::~ProgUpdater()
{
#ifdef ECXX_NODEMCU_BINARY_SIGNING_VERIFICATION
    stack_thunk_del_ref();
#endif
}


bool ProgUpdater::begin(size_t size, int command)
{
    if(_size > 0) {
        UPDATER_DEBUG("begin() already running\n");
        return false;
    }

    // If boot mode is 1 (UART download mode), it will not be able to reset into normal
    // mode once update is done
    const uint8_t bootMode = (GPI >> 16) & 0xF;
    if(bootMode == 1) {
        _setError(UPDATE_ERROR_BOOTSTRAP);
        return false;
    }

    if(command == U_FS) UPDATER_DEBUG("begin() update Filesystem\n");

    if(size == 0) {
        _setError(UPDATE_ERROR_SIZE);
        return false;
    }

    if( !checkFlashConfig(false) ) {
        _setError(UPDATE_ERROR_FLASH_CONFIG);
        return false;
    }

    _reset();
    clearError();

    _md5Expected = emptyString;
    _md5         = MD5Builder();

    wifi_set_sleep_type(NONE_SLEEP_T);

    // Address where we will start writing the update
    uintptr_t updateStartAddress = 0;

    // Size of the current binary rounded to a sector
    const size_t currentSketchSize = ( getFlashUserProgramSize() + FLASH_SECTOR_SIZE - 1 ) & ( ~(FLASH_SECTOR_SIZE - 1) );

    // Size of the update rounded to a sector
    const size_t roundedSize = ( size + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1) );

    if(command == U_FLASH) {
        // Address of the end of the space available for binary and update
        const uintptr_t updateEndAddress = reinterpret_cast<uintptr_t>(&_FS_start) - 0x40200000;

        updateStartAddress = (updateEndAddress > roundedSize) ? (updateEndAddress - roundedSize) : 0;

        UPDATER_DEBUG("begin() roundedSize       : 0x%08zX (%zd)\n", roundedSize, roundedSize);
        UPDATER_DEBUG("begin() updateEndAddress  : 0x%08zX (%zd)\n", updateEndAddress, updateEndAddress);
        UPDATER_DEBUG("begin() currentSketchSize : 0x%08zX (%zd)\n", currentSketchSize, currentSketchSize);

        if(updateStartAddress < currentSketchSize) {
            _setError(UPDATE_ERROR_SPACE);
            return false;
        }
    }
    else if(command == U_FS) {
        if( reinterpret_cast<uintptr_t>(&_FS_start) + roundedSize > reinterpret_cast<uintptr_t>(&_FS_end) ) {
            _setError(UPDATE_ERROR_SPACE);
            return false;
        }
#ifdef ECXX_NODEMCU_ATOMIC_FS_UPDATE
        // Address of the end of the space available for update
        const uintptr_t updateEndAddress = reinterpret_cast<uintptr_t>(&_FS_start) - 0x40200000;

        updateStartAddress = (updateEndAddress > roundedSize) ? (updateEndAddress - roundedSize) : 0;

        if(updateStartAddress < currentSketchSize) {
            _setError(UPDATE_ERROR_SPACE);
            return false;
        }
#else
        updateStartAddress = reinterpret_cast<uintptr_t>(&_FS_start) - 0x40200000;
#endif
    }
    else {
        UPDATER_DEBUG("begin() Unknown update command\n");
        return false;
    }

    // Initialize
    _startAddress   = updateStartAddress;
    _currentAddress = _startAddress;
    _size           = size;

    if(system_get_free_heap_size() > 2 * FLASH_SECTOR_SIZE) {
        _bufferSize = FLASH_SECTOR_SIZE;
    }
    else {
        _bufferSize = 256;
    }

    _buffer = new uint8_t[_bufferSize];
    _command = command;

    UPDATER_DEBUG("begin() _startAddress   : 0x%08X (%d)\n", _startAddress, _startAddress);
    UPDATER_DEBUG("begin() _currentAddress : 0x%08X (%d)\n", _currentAddress, _currentAddress);
    UPDATER_DEBUG("begin() _size           : 0x%08zX (%zd)\n", _size, _size);

    if(!_verify || _command == U_FS) _md5.begin();

    return true;
}


bool ProgUpdater::end(bool evenIfRemaining)
{
    if(_size == 0) {
        UPDATER_DEBUG("end() no update\n");
        _reset();
        return false;
    }

    if( !progress() ) { // Updating without any data is an error
        _setError(UPDATE_ERROR_NO_DATA);
    }

    if( hasError() || (!isFinished() && !evenIfRemaining) ) {
        UPDATER_DEBUG("end() premature end (res:%u, pos:%zu/%zu)\n", getError(), progress(), _size);
        _reset();
        return false;
    }

    if(evenIfRemaining) {
        if(_bufferLen > 0) _writeBuffer();
        _size = progress();
    }

    // Use binary signing verification
    if(_verify && _command != U_FS) {
        // Get and check the signature length
        uint32_t sigLen = 0;

        spiFlashRead( _startAddress + _size - sizeof(sigLen), &sigLen, sizeof(sigLen) );
        UPDATER_DEBUG("end() sigLen: %d\n", sigLen);

        if( sigLen != _verify->length() ) {
            _setError(UPDATE_ERROR_SIGN);
            _reset();
            return false;
        }

        // Calculate the adjusted binary size
        const int32_t binSize = _size - sigLen - sizeof(sigLen);
        UPDATER_DEBUG("end() adjusted binary size: %d\n", binSize);

        // Calculate the MD5 and hash using the adjusted binary size
        uint8_t buff[128] __ECXX_ATTR_ALIGNED_4;

        _hash->begin();
        for( int32_t i = 0; i < binSize; i += sizeof(buff) ) {
            spiFlashRead( _startAddress + i, reinterpret_cast<uint32_t*>(buff), sizeof(buff) );
            const size_t read = eCxx::min<int32_t>(sizeof(buff), binSize - i);
            _hash->add(buff, read);
        }
        _hash->end();

        const uint8_t* ret = static_cast<const uint8_t*>( _hash->hash() );
        eCxx::ignoreUnused(ret);

        UPDATER_DEBUG("end() computed hash:");
        for(int32_t i = 0; i < _hash->len(); ++i) {
            UPDATER_DEBUG(" %02x", ret[i]);
        }
        UPDATER_DEBUG("\n");

        // Read the signature
        uint8_t* sig = static_cast<uint8_t*>( malloc(sigLen) );
        if(!sig) {
            _setError(UPDATE_ERROR_SIGN);
            _reset();
            return false;
        }

        spiFlashRead(_startAddress + binSize, sig, sigLen);

        UPDATER_DEBUG("end() received signature:");
        for(uint32_t i = 0; i < sigLen; i++) {
            UPDATER_DEBUG(" %02x", sig[i]);
        }
        UPDATER_DEBUG("\n");

        // Verify the signature
        if( !_verify->verify(_hash, sig, sigLen) ) {
            free(sig);
            _setError(UPDATE_ERROR_SIGN);
            _reset();
            return false;
        }

        free(sig);
        _size = binSize; // Adjust size to remove signature (not part of binary payload)

        UPDATER_DEBUG("end() signature matches\n");
    }
    // Use MD5 only
    else if( _md5Expected.length() ) {
        _md5.calculate();
        if( strcasecmp( _md5Expected.c_str(), _md5.toString().c_str() ) ) {
            _setError(UPDATE_ERROR_MD5);
            return false;
        }
        else {
            UPDATER_DEBUG( "end() MD5 success: %s\n", _md5Expected.c_str() );
        }
    }

    if( !_verifyEnd() ) {
        _reset();
        return false;
    }

    // Update program
    if(_command == U_FLASH) {
        eboot_command ebcmd;
        ebcmd.action  = ACTION_COPY_RAW;
        ebcmd.args[0] = _startAddress;
        ebcmd.args[1] = 0x00000;
        ebcmd.args[2] = _size;
        ebootCommandWrite(&ebcmd);
        UPDATER_DEBUG("end() staged (address:0x%08X, size:0x%08zX)\n", _startAddress, _size);
    }
    // Update file system
    else if(_command == U_FS) {
#ifdef ECXX_NODEMCU_ATOMIC_FS_UPDATE
        eboot_command ebcmd;
        ebcmd.action  = ACTION_COPY_RAW;
        ebcmd.args[0] = _startAddress;
        ebcmd.args[1] = reinterpret_cast<uintptr_t>(&_FS_start) - 0x40200000;
        ebcmd.args[2] = _size;
        ebootCommandWrite(&ebcmd);
#endif
        UPDATER_DEBUG("end() filesystem (address:0x%08X, size:0x%08zX)\n", _startAddress, _size);
    }

    _reset();
    return true;
}


size_t ProgUpdater::write(uint8_t* data, size_t len)
{
    if( hasError() || !isRunning() ) return 0;

    if( progress() + _bufferLen + len > _size ) {
        _setError(UPDATE_ERROR_SPACE);
        return 0;
    }

    size_t left = len;

    while( (_bufferLen + left) > _bufferSize ) {
        const size_t toBuff = _bufferSize - _bufferLen;
        memcpy(_buffer + _bufferLen, data + (len - left), toBuff);
        _bufferLen += toBuff;
        if( !_writeBuffer() ) return len - left;
        left -= toBuff;
        if(!_async) yield();
    }

    // Check what is left
    memcpy(_buffer + _bufferLen, data + (len - left), left);
    _bufferLen += left;

    if( _bufferLen == remaining() ) {
        // Write what is left to flash
        if( !_writeBuffer() ) return len - left;
    }

    return len;
}


size_t ProgUpdater::writeStream(Stream &data)
{
    if( hasError() || !isRunning() ) return 0;

    if( !_verifyHeader( data.peek() ) ) {
        _reset();
        return 0;
    }

    if(_progressCB) _progressCB(0, _size);

    size_t written = 0;
    size_t toRead  = 0;

    while( remaining() ) {
        size_t bytesToRead = _bufferSize - _bufferLen;
        if( bytesToRead > remaining() ) bytesToRead = remaining();

        toRead = data.readBytes(_buffer + _bufferLen,  bytesToRead);
        if(toRead == 0) { // Timeout
            delay(100);
            toRead = data.readBytes(_buffer + _bufferLen, bytesToRead);
            if(toRead == 0) { // Timeout
                _currentAddress = _startAddress + _size;
                _setError(UPDATE_ERROR_STREAM);
                return written;
            }
        }
        _bufferLen += toRead;

        if( ( _bufferLen == remaining() || _bufferLen == _bufferSize ) && !_writeBuffer() ) return written;
        written += toRead;

        if(_progressCB) _progressCB(progress(), _size);
        yield();
    }

    if(_progressCB) _progressCB(progress(), _size);

    return written;
}


void ProgUpdater::_reset()
{
  if(_buffer) delete[] _buffer;

  _buffer         = 0;
  _bufferLen      = 0;
  _startAddress   = 0;
  _currentAddress = 0;
  _size           = 0;
  _command        = U_FLASH;
}


bool ProgUpdater::_writeBuffer()
{
    constexpr uint8_t FLASH_MODE_PAGE   = 0;
    constexpr uint8_t FLASH_MODE_OFFSET = 2;

    bool eraseResult = true;
    bool writeResult = true;

    if(_currentAddress % FLASH_SECTOR_SIZE == 0) {
        if(!_async) yield();
        eraseResult = spiFlashEraseSector(_currentAddress / FLASH_SECTOR_SIZE);
    }

    // If the flash settings do not match, save the original setting, modify it, and later restore it

    bool          modifyFlashMode = false;
    FlashChipMode flashMode       = FCM_QIO;
    FlashChipMode bufferFlashMode = FCM_QIO;

    if( (_currentAddress == _startAddress + FLASH_MODE_PAGE) && (_buffer[0] != 0x1f) && (_command == U_FLASH) ) {
        UPDATER_DEBUG("Header: 0x%1X %1X %1X %1X\n", _buffer[0], _buffer[1], _buffer[2], _buffer[3]);

        flashMode       = getFlashChipMode();
        bufferFlashMode = magicFlashChipMode(_buffer[FLASH_MODE_OFFSET]);

        if(bufferFlashMode != flashMode) {
            UPDATER_DEBUG("Set flash mode from 0x%1X to 0x%1X\n", bufferFlashMode, flashMode);
            _buffer[FLASH_MODE_OFFSET] = flashMode;
            modifyFlashMode = true;
        }
    }

    // Erase OK
    if(eraseResult) {
        if(!_async) yield();
        writeResult = spiFlashWrite(_currentAddress, _buffer, _bufferLen);
    }
    // Erase failed
    else {
        _currentAddress = _startAddress + _size;
        _setError(UPDATE_ERROR_ERASE);
        return false;
    }

    // Restore the old flash settinh if it has been modified
    // (ensures that the MD5 hash will still match what was sent)
    if(modifyFlashMode) {
        _buffer[FLASH_MODE_OFFSET] = bufferFlashMode;
    }

    // Check if the write failed
    if(!writeResult) {
        _currentAddress = _startAddress + _size;
        _setError(UPDATE_ERROR_WRITE);
        return false;
    }

    // Done for this step
    if(!_verify || _command == U_FS) _md5.add(_buffer, _bufferLen);

    _currentAddress += _bufferLen;
    _bufferLen       = 0;

    return true;
}


bool ProgUpdater::_verifyHeader(uint8_t data)
{
    if(_command == U_FLASH) {
        // Check for valid magic bytes
        if( (data != 0xE9) && (data != 0x1F) ) {
            _currentAddress = _startAddress + _size;
            _setError(UPDATE_ERROR_MAGIC_BYTE);
            return false;
        }
        return true;
    }

    else if(_command == U_FS) {
        // No check of file system update is possible
        return true;
    }

    return false;
}


bool ProgUpdater::_verifyEnd()
{
    if(_command == U_FLASH) {
        uint8_t buf[4] __ECXX_ATTR_ALIGNED_4;

        if( !spiFlashRead(_startAddress, reinterpret_cast<uint32_t*>(&buf[0]), 4) ) {
            _currentAddress = _startAddress;
            _setError(UPDATE_ERROR_READ);
            return false;
        }

        // Check for valid magic byte(s)
        if( (buf[0] == 0x1F) && (buf[1] == 0x8B) ) {
            // GZIP compresses the chip size flag so it cannot be checked here, therefore, simply assume OK
            return true;
        }
        else if(buf[0] != 0xE9) {
            _currentAddress = _startAddress;
            _setError(UPDATE_ERROR_MAGIC_BYTE);
            return false;
        }

        // Check if the new binary fits the flash
        const uint32_t binFlashSize = magicFlashChipSize( (buf[3] & 0xF0) >> 4 );

        if( binFlashSize > getFlashChipRealSize() ) {
            _currentAddress = _startAddress;
            _setError(UPDATE_ERROR_NEW_FLASH_CONFIG);
            return false;
        }

        return true;
    }

    else if(_command == U_FS) {
        // The file system is already overwritten, checking it here make no sense
        return true;
    }

    return false;
}


void ProgUpdater::_setError(uint8_t error)
{
    _error = error;

    // Any error condition invalidates the entire update, therefore, clear partial status
    _reset();
}


PGMCharPointer ProgUpdater::getSimpleErrorString()
{
    switch(_error) {
        case UPDATE_ERROR_OK               : return eCxx::PGM_PTR(Application::PGMStrNoError                                 );
        case UPDATE_ERROR_WRITE            : return       PGM_STR("Flash Write Failed"                                       );
        case UPDATE_ERROR_ERASE            : return       PGM_STR("Flash Erase Failed"                                       );
        case UPDATE_ERROR_READ             : return       PGM_STR("Flash Read Failed"                                        );
        case UPDATE_ERROR_SPACE            : return       PGM_STR("Not Enough Space"                                         );
        case UPDATE_ERROR_SIZE             : return       PGM_STR("Bad Size Given"                                           );
        case UPDATE_ERROR_STREAM           : return       PGM_STR("Stream Read Timeout"                                      );
        case UPDATE_ERROR_MD5              : return       PGM_STR("MD5 Failed"                                               );
        case UPDATE_ERROR_FLASH_CONFIG     : return       PGM_STR("Incorrect Flash Configuration"                            );
        case UPDATE_ERROR_NEW_FLASH_CONFIG : return       PGM_STR("Incorrect New Flash Configuration"                        );
        case UPDATE_ERROR_MAGIC_BYTE       : return       PGM_STR("Incorrect Magic Nyte (Not 0xE9)"                          );
        case UPDATE_ERROR_BOOTSTRAP        : return       PGM_STR("Invalid Bootstrapping State (Please Reset before Updating");
        case UPDATE_ERROR_SIGN             : return       PGM_STR("Signature Verification Failed"                            );
        case UPDATE_ERROR_NO_DATA          : return       PGM_STR("No Data Supplied"                                         );
        default                            : return eCxx::PGM_PTR(Application::PGMStrUnknown                                 );
    }
}


void ProgUpdater::printError(Print &out)
{
    out.printf_P( PGM_STR("ERROR[%u]: ").p, _error );

    const char* pgmErrorString = getSimpleErrorString().p;

    if(_error == UPDATE_ERROR_MD5) {
        out.print_P ( pgmErrorString );
        out.printf_P( PGM_STR(" - Expected = %s ; Calculated = %s\n").p, _md5Expected.c_str(), _md5.toString().c_str() );
    }
    else if(_error == UPDATE_ERROR_FLASH_CONFIG) {
        out.print_P ( pgmErrorString );
        out.printf_P( PGM_STR(" - Real = %d ; Selected = %d\n").p, getFlashChipRealSize(), getFlashChipSize() );
    }
    else if(_error == UPDATE_ERROR_NEW_FLASH_CONFIG) {
        out.print_P ( pgmErrorString );
        out.printf_P( PGM_STR(" - Real = %d\n").p, getFlashChipRealSize() );
    }
    else {
        out.println_P( pgmErrorString );
    }
}


//
// Create a global instance for the ProgUpdater class
//
ProgUpdater Updater;


} // namespace NodeMCU

} // namespace eCxx

