#include "STM32BurnProtocol.h"
#include "port/BurnPort.h"
#include "ProgressLog.h"
#include "STM32BurnProtocolErrors.h"
#include "protocol/BurnDeviceInfo.h"
#include "protocol/BurnProtocolErrors.h"
#include "LQError.h"
#include <QThread>


static const uint8_t stm_autobaud_req   = 0x7F;
static const uint8_t stm_cmd_resp_ack   = 0x79;
static const uint8_t stm_cmd_resp_nack  = 0x1F;

static const uint8_t stm_cmd_code_get                 = 0x00;
static const uint8_t stm_cmd_code_get_ver_rdp         = 0x01;
static const uint8_t stm_cmd_code_get_id              = 0x02;
static const uint8_t stm_cmd_code_rd_mem              = 0x11;
static const uint8_t stm_cmd_code_go                  = 0x21;
static const uint8_t stm_cmd_code_wr_mem              = 0x31;
static const uint8_t stm_cmd_code_erase               = 0x43;
static const uint8_t stm_cmd_code_ext_erase           = 0x44;
static const uint8_t stm_cmd_code_wr_protect          = 0x63;
static const uint8_t stm_cmd_code_wr_unprotect        = 0x73;
static const uint8_t stm_cmd_code_rd_protect          = 0x82;
static const uint8_t stm_cmd_code_rd_unprotect        = 0x92;
static const uint8_t stm_cmd_code_get_checksum        = 0xA1;

static const uint16_t stm_spec_erase_mass              = 0xFFFF;
static const uint16_t stm_spec_erase_bank1             = 0xFFFE;
static const uint16_t stm_spec_erase_bank2             = 0xFFFD;


static const uint16_t stm_pid_stm32g03x = 0x466;
static const uint16_t stm_pid_stm32h74x = 0x450;
static const uint16_t stm_pid_v84xxx = 0x0342;

static const int stm_mem_rd_max_size = 256;
static const int stm_mem_wr_max_size = 256;


static const BurnDeviceTypeInfo f_stm32g03x_info {QStringLiteral("STM32G03x"), QList<BurnMemArea>{BurnMemArea{0x08000000, 0x10000}}};
static const BurnDeviceTypeInfo f_stm32h74x_info {QStringLiteral("STM32H74x"), QList<BurnMemArea>{BurnMemArea{0x08000000, 0x100000}, BurnMemArea{0x08100000, 0x100000}}};
static const BurnDeviceTypeInfo f_v84xxx_info    {QStringLiteral("V84xxx"), QList<BurnMemArea>{BurnMemArea{0x08000000, 0x100000}}};



struct STM32BurnDevInfo {
    uint16_t pid;
    int mass_erase_time;
    uint32_t opts_addr;
    uint32_t opts_size;
    const BurnDeviceTypeInfo *typeInfo;
};

static const STM32BurnDevInfo f_pid_dev_infos [] = {
    {stm_pid_stm32g03x,    41, 0x1FFF7800, 128, &f_stm32g03x_info},
    {stm_pid_stm32h74x, 26000,          0,   0, &f_stm32h74x_info},
    {stm_pid_v84xxx,    15000, 0x1FFFF800,  48, &f_v84xxx_info}
};


int STM32BurnProtocol::maxReadBlockSize() const {
    return stm_mem_rd_max_size;
}

int STM32BurnProtocol::maxWriteBlockSize() const {
    return stm_mem_wr_max_size;
}

void STM32BurnProtocol::setPort(BurnPort &port, LQError &err) {
    m_port = &port;
    m_port->setParity(QSerialPort::Parity::EvenParity, err);
}

void STM32BurnProtocol::autobaud(int tout, LQError &err) {
    checkPortInit(err);

    if (err.isSuccess()) {
        bool autobaud_done {false};
        /* Изначально загрузчик stm ожидает слово stm_autobaud_req для согласования
         * скорости и должен ответить stm_cmd_resp_ack.
         * Однако если он уже эту стадию прошел, то он уже не отвечает на stm_autobaud_req.
         * При этом в режиме ожидания команд два stm_autobaud_req приведут к посылке
         * NACK.
         * Поэтому, чтобы не требовать каждый раз перезагружать устройство, а иметь
         * возможность продолжить работу при сохранении скорости, то посылаем
         * до двух запрос и считаем ответ как ACK, так и NACK признаком того,
         * что загрузчик котов принимать команды */
        for (int i {0}; (i < 4) && !autobaud_done && err.isSuccess(); ++i) {
            m_port->write(stm_autobaud_req, err);
            if (err.isSuccess()) {
                quint8 resp;

                if (m_port->read(resp, tout, err)) {
                    if ((resp == stm_cmd_resp_ack) || (resp == stm_cmd_resp_nack)) {
                        autobaud_done = true;
                    } else {
                        err = STM32BurnProtocolErrors::autobaudInvalidResponse(resp);
                    }
                }
            }
        }

        if (err.isSuccess() && !autobaud_done) {
            err = BurnProtocolErrors::noAutobaudResponse();
        }
    }
}

void STM32BurnProtocol::getDevInfo(BurnDeviceInfo &info, LQError &err) {
    checkPortInit(err);
    if (err.isSuccess())
        updateDevInfo(err);
    if (err.isSuccess()) {
        info.setBootloaderVersion(QString{"%1.%2"}.arg((m_boot_ver >> 4) & 0xF)
                                  .arg(m_boot_ver & 0xF));
        info.setDevTypeInfo(*m_devInfo->typeInfo);
    }
}

void STM32BurnProtocol::devCleanup(LQError &err) {
    checkDevInfoValid(err);
    if (err.isSuccess()) {
        QByteArray dummy;
        /* посылаем тестовую команду чтения, чтобы проверить, не включен ли RDP */
        LQError checkErr;
        cmdExecMemRd(m_devInfo->typeInfo->defaultBootAddr(), 4,  dummy, checkErr);
        if (checkErr == STM32BurnProtocolErrors::cmdRejectedRDP()) {
            cmdExecRdUnprotect(err);
        }

        if (err.isSuccess()) {
            cmdExecWrUnprotect(err);
            if (err.isSuccess())
                cmdExecMassErase(err);
        }
    }
}

void STM32BurnProtocol::writeMemBlock(burn_mem_addr_t addr, const QByteArray &data, LQError &err) {
    checkDevInfoValid(err);
    if (err.isSuccess())
        cmdExecMemWr(addr, data, err);
}

void STM32BurnProtocol::readMemBlock(burn_mem_addr_t addr, int size, QByteArray &data, LQError &err) {
    checkDevInfoValid(err);
    if (err.isSuccess())
        cmdExecMemRd(addr, size, data, err);
}

void STM32BurnProtocol::startApp(burn_mem_addr_t addr, LQError &err) {
    checkDevInfoValid(err);
    if (err.isSuccess())
        cmdExecGo(addr, err);
}

void STM32BurnProtocol::checkPortInit(LQError &err) {
    if (!m_port)
        err += BurnProtocolErrors::portNotInit();
}

void STM32BurnProtocol::checkDevInfoValid(LQError &err) {
    checkPortInit(err);
    if (err.isSuccess()) {
        if (!m_devInfo) {
            updateDevInfo(err);
        }
    }
}

void STM32BurnProtocol::updateDevInfo(LQError &err) {
    if (err.isSuccess()) {
        uint8_t boot_ver;
        QList<uint8_t> cmdList;
        cmdExecGet(boot_ver, cmdList, err);
        if (err.isSuccess()) {
            m_boot_ver = boot_ver;
            m_supported_cmd_list = cmdList;
        }
    }

    if (err.isSuccess()) {
        quint16 dev_pid {0};
        cmdExecGetID(dev_pid, err);
        if (err.isSuccess()) {
            m_devInfo = nullptr;
            for (size_t i = 0; (i < sizeof(f_pid_dev_infos)/sizeof(f_pid_dev_infos[0]))
                 && (m_devInfo == nullptr); ++i) {
                if (f_pid_dev_infos[i].pid == dev_pid) {
                    m_devInfo = &f_pid_dev_infos[i];
                }
            }
            if (m_devInfo == nullptr) {
                err = STM32BurnProtocolErrors::unknowPID(dev_pid);
            }
        }
    }
}

/* передача байта и его дополнения. так передаются команды и некоторые одиночные байты данных,
 * в то время как в остальных случаях идет расчет контрольной суммы по XOR */
void STM32BurnProtocol::sendComplemented(quint8 data, LQError &err) {
    QByteArray arr{2, Qt::Uninitialized};
    arr[0] = static_cast<char>(data);
    arr[1] = static_cast<char>(data ^ 0xFF);
    m_port->write(arr, err);
}

void STM32BurnProtocol::sendCmd(quint8 cmd, LQError &err) {
    m_port->clearData();
    sendComplemented(cmd, err);
}

void STM32BurnProtocol::recvCmdAck(quint8 cmd, LQError &err, int tout) {
    quint8 resp;
    if (!m_port->read(resp, tout, err)) {
        err = STM32BurnProtocolErrors::cmdNoResponse();
    } else if (resp == stm_cmd_resp_nack) {
        if (cmdRdpModeAvailable(cmd)) {
            err = STM32BurnProtocolErrors::cmdRejected();
        } else {
            err = STM32BurnProtocolErrors::cmdRejectedRDP();
        }
    } else if (resp != stm_cmd_resp_ack)  {
        err = STM32BurnProtocolErrors::cmdInvalidResponse(resp);
    }
}

void STM32BurnProtocol::recvDataAck(LQError &err, int tout) {
    quint8 resp;
    if (!m_port->read(resp, tout, err)) {
        err = STM32BurnProtocolErrors::cmdDataNoAck();
    } else if (resp == stm_cmd_resp_nack) {
        err = STM32BurnProtocolErrors::cmdDataRejected();
    } else if (resp != stm_cmd_resp_ack)  {
        err = STM32BurnProtocolErrors::cmdDataInvalidAck(resp);
    }
}

QByteArray STM32BurnProtocol::recvCmdData(int size, LQError &err, int tout) {
    QByteArray data;
    if (!m_port->read(data, size, tout, err)) {
        err = STM32BurnProtocolErrors::cmdInsufRespData();
    }
    return data;
}

QByteArray STM32BurnProtocol::recvCmdVarData(LQError &err, int tout) {
    QByteArray ret;
    quint8 rx_size;
    if (!m_port->read(rx_size, tout, err)) {
        if (err.isSuccess())
            err = STM32BurnProtocolErrors::cmdInsufRespData();
    }

    if (err.isSuccess()) {
        m_port->read(ret, rx_size+1, tout, err);
        if (err.isSuccess() && (ret.size() != (rx_size+1))) {
            err = STM32BurnProtocolErrors::cmdInsufRespData();
        }
    }
    return ret;
}

void STM32BurnProtocol::execCmd(quint8 cmd, LQError &err, int tout) {
    sendCmd(cmd, err);
    if (err.isSuccess()) {
        recvCmdAck(cmd, err, tout);
    }
}

void STM32BurnProtocol::resetAutobaud(LQError &err) {
    QThread::msleep(reset_wait_time);
    autobaud(reset_autobaud_tout, err);
    if (!err.isSuccess()) {
        err.addMsgPrefix("Autobaud after device reset failed");
    }
}

void STM32BurnProtocol::readOptions(QByteArray &options, LQError &err) {
    checkDevInfoValid(err);
    /* не все устройства поддерживают чтение опций через адреса памяти */
    if (err.isSuccess() && (m_devInfo->opts_size == 0)) {
        err = STM32BurnProtocolErrors::deviceUnsupOptsAccess();
    }
    cmdExecMemRd(m_devInfo->opts_addr, m_devInfo->opts_size, options, err);

}

bool STM32BurnProtocol::checkCmdSupport(quint16 cmd) const {
    return m_supported_cmd_list.contains(cmd);
}

bool STM32BurnProtocol::cmdRdpModeAvailable(quint8 cmd) const {
    return (cmd == stm_cmd_code_get)
            || (cmd == stm_cmd_code_get_ver_rdp)
            || (cmd == stm_cmd_code_get_id)
            || (cmd == stm_cmd_code_rd_unprotect);
}

void STM32BurnProtocol::cmdExecGet(quint8 &boot_ver, QList<quint8> &cmd_list, LQError &err) {
    execCmd(stm_cmd_code_get, err);
    if (err.isSuccess()) {
        QByteArray rxData = recvCmdVarData(err);
        if (err.isSuccess()) {
            boot_ver = rxData.at(0);
            for (int i = 1; i < rxData.size(); ++i) {
                cmd_list.append(rxData.at(i));
            }
            recvDataAck(err);
        }
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Get' command failed"));
    }
}

void STM32BurnProtocol::cmdExecGetID(quint16 &id, LQError &err) {
    execCmd(stm_cmd_code_get_id, err);
    if (err.isSuccess()) {
        QByteArray rxData = recvCmdVarData(err);
        if (err.isSuccess()) {
            if (rxData.size() >= 2) {
                id = (rxData.at(0) << 8) | rxData.at(1);
            } else {
                err = STM32BurnProtocolErrors::cmdInvalidRespDataSize();
            }

            recvDataAck(err);
        }
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Get ID' command failed"));
    }
}

void STM32BurnProtocol::cmdExecMemRd(quint32 addr, int size, QByteArray &data, LQError &err) {
    if ((size == 0) || (size > stm_mem_rd_max_size)) {
        err = STM32BurnProtocolErrors::invalidMemBlockSizeParam(size);
    }

    if (err.isSuccess())
        execCmd(stm_cmd_code_rd_mem, err);
    if (err.isSuccess()) {
        QByteArray txData{5, Qt::Initialization::Uninitialized};
        int pos {0};
        txData[pos++] = static_cast<char>((addr >>  24) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>  16) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   8) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   0) & 0xFF);
        txData[pos] = calcCheckSum(txData, pos);
        m_port->write(txData, err);
        if (err.isSuccess()) {
            recvDataAck(err);
        }
    }

    if (err.isSuccess()) {
        sendComplemented(size-1, err);
        if (err.isSuccess()) {
            recvDataAck(err);
        }
    }

    if (err.isSuccess()) {
        data = recvCmdData(size, err);
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Memory Read' command failed"));
    }
}

void STM32BurnProtocol::cmdExecMemWr(quint32 addr, const QByteArray &data, LQError &err) {
    int size = data.size();
    if ((size == 0) || (size > stm_mem_wr_max_size)) {
        err = STM32BurnProtocolErrors::invalidMemBlockSizeParam(size);
    }

    if (err.isSuccess())
        execCmd(stm_cmd_code_wr_mem, err);
    if (err.isSuccess()) {
        QByteArray txData{5, Qt::Initialization::Uninitialized};
        int pos {0};
        txData[pos++] = static_cast<char>((addr >>  24) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>  16) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   8) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   0) & 0xFF);
        txData[pos] = calcCheckSum(txData, pos);
        m_port->write(txData, err);
        if (err.isSuccess()) {
            recvDataAck(err);
        }
    }

    if (err.isSuccess()) {
        QByteArray txData{size + 2, Qt::Initialization::Uninitialized};
        int pos {0};
        txData[pos++] = static_cast<char>(size-1);
        for (int i = 0; i < size; ++i) {
            txData[pos++] = data[i];
        }
        txData[pos] = calcCheckSum(txData, pos);
        m_port->write(txData, err);
        if (err.isSuccess()) {
            recvDataAck(err);
        }
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Memory Write' command failed"));
    }
}

void STM32BurnProtocol::cmdExecGo(quint32 addr, LQError &err) {
    execCmd(stm_cmd_code_go, err);
    if (err.isSuccess()) {
        QByteArray txData{5, Qt::Initialization::Uninitialized};
        int pos {0};
        txData[pos++] = static_cast<char>((addr >>  24) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>  16) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   8) & 0xFF);
        txData[pos++] = static_cast<char>((addr >>   0) & 0xFF);
        txData[pos] = calcCheckSum(txData, pos);
        m_port->write(txData, err);
        if (err.isSuccess()) {
            recvDataAck(err);
        }
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'GO' command failed"));
    }
}

void STM32BurnProtocol::cmdExecMassErase(LQError &err) {
    if (checkCmdSupport(stm_cmd_code_ext_erase)) {
        execCmd(stm_cmd_code_ext_erase, err);
        if (err.isSuccess()) {
            QByteArray txData{3, Qt::Initialization::Uninitialized};
            int pos {0};
            txData[pos++] = static_cast<char>((stm_spec_erase_mass >> 8) & 0xFF);
            txData[pos++] = static_cast<char>(stm_spec_erase_mass & 0xFF);
            txData[pos] = calcCheckSum(txData, 2);
            m_port->write(txData, err);
            recvDataAck(err, massEraseCmdTimeout());
        }

        if (!err.isSuccess()) {
            err.addMsgPrefix(tr("Protocol 'Extended Erase' command failed"));
        }
    } else {
        execCmd(stm_cmd_code_erase, err);
        if (err.isSuccess()) {
            sendComplemented(0xFF, err);
            recvDataAck(err, massEraseCmdTimeout());
        }

        if (!err.isSuccess()) {
            err.addMsgPrefix(tr("Protocol 'Erase' command failed"));
        }
    }
}

void STM32BurnProtocol::cmdExecRdProtect(LQError &err) {
    execCmd(stm_cmd_code_rd_protect, err);
    if (err.isSuccess()) {
        recvDataAck(err);
    }
    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Readout protect' command failed"));
    }
}

void STM32BurnProtocol::cmdExecRdUnprotect(LQError &err) {
    execCmd(stm_cmd_code_rd_unprotect, err);
    if (err.isSuccess()) {
        recvDataAck(err, massEraseCmdTimeout());
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Readout Unprotect' command failed"));
    }

    if (err.isSuccess()) {
        resetAutobaud(err);
    }
}

void STM32BurnProtocol::cmdExecWrUnprotect(LQError &err) {
    execCmd(stm_cmd_code_wr_unprotect, err);
    if (err.isSuccess()) {
        recvDataAck(err);
    }

    if (!err.isSuccess()) {
        err.addMsgPrefix(tr("Protocol 'Write Unprotect' command failed"));
    }

    if (err.isSuccess()) {
        resetAutobaud(err);
    }
}

uint8_t STM32BurnProtocol::calcCheckSum(const QByteArray &data, int size) const {
    uint8_t check = 0;
    for (int i = 0; i < size; ++i) {
        check ^= static_cast<uint8_t>(data.at(i));
    }
    return check;
}

int STM32BurnProtocol::massEraseCmdTimeout() const {
    return default_cmd_tout + m_devInfo->mass_erase_time;
}

