/*
 * 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 few modifications) from:
 *     Updater.h 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
 * ----------------------------------------------------------------------------------
 */


#ifndef ECXX_PLATFORM_NODEMCU_UPDATER_H
#define ECXX_PLATFORM_NODEMCU_UPDATER_H


#include "CoreLib.h"
#include "CoreSys.h"
#include "MD5Builder.h"


//
// Constants
//
#define UPDATE_ERROR_OK                0
#define UPDATE_ERROR_WRITE             1
#define UPDATE_ERROR_ERASE             2
#define UPDATE_ERROR_READ              3
#define UPDATE_ERROR_SPACE             4
#define UPDATE_ERROR_SIZE              5
#define UPDATE_ERROR_STREAM            6
#define UPDATE_ERROR_MD5               7
#define UPDATE_ERROR_FLASH_CONFIG      8
#define UPDATE_ERROR_NEW_FLASH_CONFIG  9
#define UPDATE_ERROR_MAGIC_BYTE       10
#define UPDATE_ERROR_BOOTSTRAP        11
#define UPDATE_ERROR_SIGN             12
#define UPDATE_ERROR_NO_DATA          13

#define U_FLASH   0
#define U_FS    100
#define U_AUTH  200


//
// Uncomment this to enable binary signing verification for OTA programming
//
// NOTE : If you updated this line, simply execute "make xpk" to delete some generated
//        files so that this header will be recompiled (even you do not need to use the
//        array from the generated file "build/_public_key.h").
//
// #define ECXX_NODEMCU_BINARY_SIGNING_VERIFICATION

//
// Uncomment this to enable atomic file system updating via OTA programming
//
// NOTE : # By default, OTA filesystem update overwrites the target flash directly. This
//          can lead to the file system being corrupted if there is a power outage during
//          the update process. In order to use the same two-step process which is used
//          for OTA application firmware updates, the ECXX_NODEMCU_ATOMIC_FS_UPDATE macro
//          can be defined. However, because there must be enough unused space for the
//          new filesystem image to be stored, this feature is not enabled by default.
//        # In NodeMCU with 4 MiB flash, 'eagle.flash.4m1m.ld' could be used with atomic
//          file system updating. However, 'eagle.flash.4m2m.ld' would most likely not.
//
// #define ECXX_NODEMCU_ATOMIC_FS_UPDATE


namespace eCxx {

namespace NodeMCU {


//
// Abstract class to implement signing hash
//
class UpdaterHashClass {
    public:
        virtual void begin() = 0;
        virtual void add(const void *data, uint32_t len) = 0;
        virtual void end() = 0;
        virtual int len() = 0;
        virtual const void* hash() = 0;
        virtual const unsigned char* oid() = 0;
};


//
// Abstract class to implement a signature verifier
//
class UpdaterVerifyClass {
    public:
        virtual uint32_t length() = 0;
        virtual bool verify(UpdaterHashClass* hash, const void* signature, uint32_t signatureLen) = 0;
};


//
// Program updater class
//
class ProgUpdater {
    public:
        using ProgressCB = eCxx::FunctionalVoidCallback_Size_Size;

    public:
        ProgUpdater();
        ~ProgUpdater();

        // Optionally add a cryptographic signature verification hash and method
        inline void installSignature(UpdaterHashClass* hash, UpdaterVerifyClass* verify)
        {
            _hash   = hash;
            _verify = verify;
        }

        // Run updater from asynchronous callbacs
        void runAsync(bool async)
        { _async = async; }

        // Set callback that will be called when the updater is receiving data
        inline void onProgress(ProgressCB fn)
        { _progressCB = fn; }

        // Begin update - return false if there is not enough space
        bool begin(size_t size, int command = U_FLASH);

        // # Writes a buffer to the flash and increments the address
        // # Returns the number of bytes written
        size_t write(uint8_t *data, size_t len);

        // # Writes the remaining bytes from the Stream to the flash
        // # Uses readBytes() and sets UPDATE_ERROR_STREAM on timeout
        // # Returns the number of bytes written
        size_t writeStream(Stream &data);

        // # If all bytes are written this call will write the config to eboot and return true
        // # If there is already an update running but is not finished !evenIfRemaining() or
        //   there is an error this will clear everything and return false the last error is
        //   available through getError()
        // # Using evenIfRemaining() is helpful when updating without knowing the final size
        bool end(bool evenIfRemaining = false);

        // Get a simple program-memory string from the last error
        PGMCharPointer getSimpleErrorString();

        // Prints the last error to an output stream
        void printError(Print &out);

        // Sets the expected MD5 for the firmware (in hexadecimal string)
        inline bool setMD5(const char* expected_md5)
        {
            if(strlen(expected_md5) != 32) return false;
            _md5Expected = expected_md5;
            return true;
        }

        // Returns the MD5 String of the sucessfully ended firmware
        inline WString md5String()
        { return _md5.toString(); }

        // Populate result with the md5 bytes of the sucessfully ended firmware
        inline void md5(uint8_t* result)
        { return _md5.getBytes(result); }

        // Helper functions
        inline void    clearError() { _error = UPDATE_ERROR_OK;                          }
        inline uint8_t getError  () { return _error;                                     }
        inline bool    hasError  () { return _error != UPDATE_ERROR_OK;                  }
        inline bool    isRunning () { return _size > 0;                                  }
        inline bool    isFinished() { return _currentAddress == (_startAddress + _size); }
        inline size_t  size      () { return _size;                                      }
        inline size_t  progress  () { return _currentAddress - _startAddress;            }
        inline size_t  remaining () { return _size - (_currentAddress - _startAddress);  }

        // # Write from any object that expose available() and read(uint8_t*, size_t) methods
        // # Faster than the writeStream() method
        // # Writes only what is available
        template <typename T>
        inline size_t write(T &data)
        {
            if( hasError() || !isRunning() ) return 0;

            size_t written   = 0;
            size_t available = data.available();

            while(available) {
                if( _bufferLen + available > remaining() ) available = remaining() - _bufferLen;

                if(_bufferLen + available > _bufferSize) {
                    const size_t toBuff = _bufferSize - _bufferLen;
                    data.read(_buffer + _bufferLen, toBuff);
                    _bufferLen += toBuff;
                    if( !_writeBuffer() ) return written;
                    written += toBuff;
                }
                else {
                    data.read(_buffer + _bufferLen, available);
                    _bufferLen += available;
                    written    += available;
                    if(_bufferLen == remaining()) {
                        if( !_writeBuffer() ) return written;
                    }
                }

                if( !remaining() ) return written;

                NodeMCU::delay(1);
                available = data.available();
            }

            return written;
        }

    private:
        void _reset();
        bool _writeBuffer();

        bool _verifyHeader(uint8_t data);
        bool _verifyEnd();

        void _setError(uint8_t error);

        bool       _async;
        uint8_t    _error;
        uint8_t*   _buffer;
        size_t     _bufferLen;  // Amount of data written into _buffer
        size_t     _bufferSize; // Total size of _buffer
        size_t     _size;
        uint32_t   _startAddress;
        uint32_t   _currentAddress;
        uint32_t   _command;

        WString    _md5Expected;
        MD5Builder _md5;

        // Optional signed binary verification
        UpdaterHashClass*   _hash;
        UpdaterVerifyClass* _verify;

        // Optional progress callback function
        ProgressCB _progressCB;
};


//
// Declaration of the global instance for the ProgUpdater class
//
extern ProgUpdater Updater;


} // namespace NodeMCU

} // namespace eCxx


//
// For compatibility with the Arduino library part
//
typedef eCxx::NodeMCU::UpdaterHashClass   UpdaterHashClass;
typedef eCxx::NodeMCU::UpdaterVerifyClass UpdaterVerifyClass;


#endif
