/*
 * 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:
 *     ArduinoOTA.cpp by Ivan Grokhotkov and Miguel Angel Ajo
 * ----------------------------------------------------------------------------------
 * 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 "OTAProgramming.h"
#include "StreamString.h"


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

#ifndef OTA_DEBUG
    #define NO_OTA_DEBUG
    #define OTA_DEBUG(FMT, ...) SYSTEM_DEBUG_NOP
#endif


namespace eCxx {

namespace NodeMCU {


//
// OTA programming support
//
OTAProgramming::OTAProgramming()
: _port           (0)
, _rebootOnSuccess(true)
, _udpContext     (nullptr)
, _initialized    (false)
, _state          (OTAState::Idle)
, _cmd            (0)
, _size           (0)
, _otaPort        (0)
, _otaUDPPort     (0)
, _startCB        (nullptr)
, _progressCB     (nullptr)
, _endCB          (nullptr)
, _errorCB        (nullptr)
{}


OTAProgramming::~OTAProgramming()
{
    if(!_udpContext) return;

    _udpContext->unref();
    _udpContext = nullptr;
}


void OTAProgramming::setHostname(const char* hostname)
{
    if(_initialized || _hostname.length() || !hostname) return;

    _hostname = hostname;
}


void OTAProgramming::setPort(uint16_t port)
{
    if(_initialized || _port || !port) return;

    _port = port;
}


void OTAProgramming::setPassword(const char* password)
{
    if(_initialized || _password.length() || !password) return;

    MD5Builder md5Builder;

    md5Builder.begin();
    md5Builder.add(password);
    md5Builder.calculate();

    _password = md5Builder.toString();
}


void OTAProgramming::setPasswordMD5(const char* password)
{
    if(_initialized || _password.length() || !password) return;

    _password = password;
}


void OTAProgramming::begin()
{
    // Simply return if already initialized
    if(_initialized) return;

    // Initialize variables
    if( !_hostname.length() ) {
        char tmp[17];
        sprintf( tmp, "esp8266-%08x", system_get_chip_id() );
        _hostname = tmp;
    }

    if( !_port ) _port = 8266;

    // Initialize UDP context
    if(_udpContext) {
        _udpContext->unref();
        _udpContext = 0;
    }
    _udpContext = new UdpContext;
    _udpContext->ref();

    if( !_udpContext->listen(IP_ADDR_ANY, _port) ) return;

    _udpContext->onRx( [this]() { this->OTAProgramming::_onRx(); } );

    // Initialize mDNS
    MDNS.begin( _hostname.c_str() );

    if( _password.length() ) MDNS.enableArduino(_port, true );
    else                     MDNS.enableArduino(_port, false);

    // Done
    _initialized = true;
    _state       = OTAState::Idle;

    OTA_DEBUG("OTA server at %s.local:%u\n", _hostname.c_str(), _port);
}


void OTAProgramming::handle()
{
    if(_state == OTAState::RunUpdate) {
        _runUpdate();
        _state = OTAState::Idle;
    }

    MDNS.update();
}


void OTAProgramming::reinitializeMDNS()
{
    // Uninitialize mDNS
    MDNS.close();
    Application::delayMS(500);

    // Reinitialize mDNS
    MDNS.begin( _hostname.c_str() );

    if( _password.length() ) MDNS.enableArduino(_port, true );
    else                     MDNS.enableArduino(_port, false);
}


int32_t OTAProgramming::parseInt()
{
    while(_udpContext->peek() == ' ') _udpContext->read();

    char data[16];

    for(uint8_t i = 0; i < sizeof(data); ++i) {
        const char value = _udpContext->peek();
        if(value < '0' || value > '9') {
            data[i] = '\0';
            return atoi(data);
        }
        data[i] = _udpContext->read();
    }

    return 0;
}


WString OTAProgramming::readStringUntil(char end)
{
    WString res;

    for(;;) {
        const int32_t value = _udpContext->read();
        if(value < 0 || value == '\0' || value == end) return res;
        res += static_cast<char>(value);
    }

    return res;
}


void OTAProgramming::_onRx()
{
    if( !_udpContext->next() ) return;

    if(_state == OTAState::Idle) {
        const int32_t cmd = parseInt();
        if(cmd != U_FLASH && cmd != U_FS) return;

        _cmd        = cmd;
        _otaIP      = _udpContext->getRemoteAddress();
        _otaUDPPort = _udpContext->getRemotePort();
        _otaPort    = parseInt();
        _size       = parseInt();
        _udpContext->read();

        _md5 = readStringUntil('\n');
        _md5.trim();
        if(_md5.length() != 32) return;

        if( _password.length() ) {
            MD5Builder md5Builder;

            md5Builder.begin();
            md5Builder.add( WString( micros() ) );
            md5Builder.calculate();

            _nonce = md5Builder.toString();

            char auth_req[38];
            sprintf( auth_req, "AUTH %s", _nonce.c_str() );

            _udpContext->append( auth_req, strlen(auth_req) );
            _udpContext->send  ( _otaIP, _otaUDPPort        );

            _state = OTAState::WaitAuth;
            return;
        }
        else {
            _state = OTAState::RunUpdate;
        }
    }

    else if(_state == OTAState::WaitAuth) {
        const int32_t cmd = parseInt();
        if(cmd != U_AUTH) {
            _state = OTAState::Idle;
            return;
        }

        _udpContext->read();
        const WString& cnonce   = readStringUntil(' ');
        const WString& response = readStringUntil('\n');
        if(cnonce.length() != 32 || response.length() != 32) {
            _state = OTAState::Idle;
            return;
        }

        MD5Builder    md5Builder;
        const WString challenge = _password + ':' + WString(_nonce) + ':' + cnonce;

        md5Builder.begin();
        md5Builder.add(challenge);
        md5Builder.calculate();

        WString result = md5Builder.toString();

        if( result.equalsConstantTime(response) ) {
            _state = OTAState::RunUpdate;
        }
        else {
            _udpContext->append( "Authentication Failed", 21 );
            _udpContext->send  ( _otaIP, _otaUDPPort         );

            if(_errorCB) _errorCB(OTAError::AuthError);
            _state = OTAState::Idle;
        }
    }

    while( _udpContext->next() ) _udpContext->flush();
}


void OTAProgramming::_runUpdate()
{
    if( !Updater.begin(_size, _cmd) ) {
        OTA_DEBUG("Update begin error\n");

        if(_errorCB) _errorCB(OTAError::BeginError);

        ::StreamString ss;
        Updater.printError(ss);

        _udpContext->append( "ERR: ",    5           );
        _udpContext->append( ss.c_str(), ss.length() );
        _udpContext->send  ( _otaIP,     _otaUDPPort );

        delay(100);
        _udpContext->listen(IP_ADDR_ANY, _port);

        _state = OTAState::Idle;
        return;
    }

    _udpContext->append( "OK",   2           );
    _udpContext->send  ( _otaIP, _otaUDPPort );
    delay(100);

    Updater.setMD5( _md5.c_str() );
    WiFiUDP::stopAll();
    WiFiClient::stopAll();

    if(_startCB) _startCB();

    if(_progressCB) _progressCB(0, _size);

    WiFiClient client;
    if( !client.connect(_otaIP, _otaPort) ) {
        OTA_DEBUG("Connect failed\n");

        _udpContext->listen(IP_ADDR_ANY, _port);

        if(_errorCB) _errorCB(OTAError::ConnectError);
        _state = OTAState::Idle;
    }

    client.setNoDelay(true); // OTA sends little packets

    uint32_t totalWritten = 0;
    while( !Updater.isFinished() && ( client.connected() || client.available() ) ) {
        int32_t waited = 1000;
        while( !client.available() && waited-- ) delay(1);

        if(!waited) {
            OTA_DEBUG("Receive failed\n");

            _udpContext->listen(IP_ADDR_ANY, _port);

            if(_errorCB) _errorCB(OTAError::ReceiveError);
            _state = OTAState::Idle;
        }

        const uint32_t written = Updater.write(client);
        if(written > 0) {
            client.print(written, Print::DEC);
            totalWritten += written;
            if(_progressCB) _progressCB(totalWritten, _size);
        }
    }

    if( Updater.end() ) {
        // Ensure last count packet has been sent out and not combined with the final OK
        client.flush();
        delay(1000);
        client.print("OK");
        client.flush();
        delay(1000);
        client.stop();

        OTA_DEBUG("Update success\n");

        if(_endCB) _endCB();

        if(_rebootOnSuccess) {
            OTA_DEBUG("Rebooting ...\n");

            // Let serial/network finish tasks that might be given in end callback before restarting
            delay(500);
            restart();
        }
    }

    else {
        Updater.printError(client);
#ifndef NO_OTA_DEBUG
        {
            ::StreamString ss;
            Updater.printError(ss);
            OTA_DEBUG("Update error (%s)\n", ss.c_str());
        }
#endif

        _udpContext->listen(IP_ADDR_ANY, _port);

        if(_errorCB) _errorCB(OTAError::EndError);
        _state = OTAState::Idle;
    }
}


//
// Create a global instance for the OTAProgramming class
//
OTAProgramming OTA;


} // namespace NodeMCU

} // namespace eCxx
