#include "udsfirmwareuploaderwindow.h"
#include "ui_udsfirmwareuploaderwindow.h"
#include "mainwindow.h"
#include "connections/canconmanager.h"
#include "utility.h"

#include <QFile>
#include <QFileDialog>
#include <QDateTime>

UDSFirmwareUploaderWindow::UDSFirmwareUploaderWindow(const QVector<CANFrame> *frames, QWidget *parent) :
    QDialog(parent),
    ui(new Ui::UDSFirmwareUploaderWindow)
{
    ui->setupUi(this);
    setWindowFlags(Qt::Window);

    modelFrames = frames;
    currentState = STATE_IDLE;
    firmwareSize = 0;
    blockSequenceCounter = 0;
    maxBlockLength = 0;
    currentDataPos = 0;
    busId = 0;
    targetId = 0x7E0;
    responseId = 0x7E8;
    sessionType = 0x02;
    securityLevel = 1;
    flashAddr = 0x08000000;
    useSecurity = true;

    udsHandler = new UDS_HANDLER();
    timeoutTimer = new QTimer();
    timeoutTimer->setInterval(1000);
    testerTimer = new QTimer();
    testerTimer->setInterval(2000);
    progressTimer = new QElapsedTimer();

    int numBuses = CANConManager::getInstance()->getNumBuses();
    ui->spinBus->setMaximum(numBuses > 0 ? numBuses - 1 : 0);

    connect(ui->btnLoadFile, &QPushButton::clicked, this, &UDSFirmwareUploaderWindow::handleLoadFile);
    connect(ui->btnStartStop, &QPushButton::clicked, this, &UDSFirmwareUploaderWindow::handleStartStop);
    connect(ui->btnAbort, &QPushButton::clicked, this, &UDSFirmwareUploaderWindow::handleAbort);
    connect(udsHandler, &UDS_HANDLER::newUDSMessage, this, &UDSFirmwareUploaderWindow::gotUDSReply);
    connect(timeoutTimer, &QTimer::timeout, this, &UDSFirmwareUploaderWindow::timeoutTriggered);
    connect(testerTimer, &QTimer::timeout, this, &UDSFirmwareUploaderWindow::testerPresentTick);
}

UDSFirmwareUploaderWindow::~UDSFirmwareUploaderWindow()
{
    stopUpload();
    delete testerTimer;
    delete timeoutTimer;
    delete udsHandler;
    delete progressTimer;
    delete ui;
}

void UDSFirmwareUploaderWindow::logMessage(const QString &msg)
{
    QString timestamp = QDateTime::currentDateTime().toString("hh:mm:ss.zzz");
    ui->txtLog->append("[" + timestamp + "] " + msg);
}

void UDSFirmwareUploaderWindow::setState(UPLOAD_STATE newState)
{
    currentState = newState;
    if (newState == STATE_ERROR || newState == STATE_COMPLETE || newState == STATE_IDLE)
    {
        timeoutTimer->stop();
        testerTimer->stop();
        ui->btnStartStop->setEnabled(true);
        ui->btnAbort->setEnabled(false);
        ui->btnLoadFile->setEnabled(true);
    }
}

void UDSFirmwareUploaderWindow::handleLoadFile()
{
    QFileDialog dialog;
    QStringList filters;
    filters.append("Binary Firmware (*.bin)");
    filters.append("All Files (*)");
    dialog.setFileMode(QFileDialog::ExistingFile);
    dialog.setNameFilters(filters);

    if (dialog.exec() == QDialog::Accepted)
    {
        QString filename = dialog.selectedFiles().constFirst();
        QFile inFile(filename);
        if (!inFile.open(QIODevice::ReadOnly))
        {
            logMessage("ERROR: Could not open file: " + filename);
            return;
        }
        firmwareData = inFile.readAll();
        firmwareSize = firmwareData.length();
        inFile.close();
        ui->lblFilename->setText(filename + QString(" (%1 bytes)").arg(firmwareSize));
        logMessage("Loaded firmware: " + filename + " (" + QString::number(firmwareSize) + " bytes)");
    }
}

void UDSFirmwareUploaderWindow::handleStartStop()
{
    if (currentState != STATE_IDLE && currentState != STATE_COMPLETE && currentState != STATE_ERROR)
    {
        stopUpload();
        return;
    }
    startUpload();
}

void UDSFirmwareUploaderWindow::handleAbort()
{
    stopUpload();
    logMessage("Upload aborted by user");
    setState(STATE_IDLE);
}

void UDSFirmwareUploaderWindow::startUpload()
{
    if (firmwareData.isEmpty())
    {
        logMessage("ERROR: No firmware file loaded");
        return;
    }

    busId = ui->spinBus->value();
    targetId = Utility::ParseStringToNum(ui->txtTargetID->text());
    responseId = Utility::ParseStringToNum(ui->txtResponseID->text());
    flashAddr = Utility::ParseStringToNum(ui->txtFlashAddr->text());
    securityLevel = ui->spinSecurityLevel->value();
    useSecurity = ui->ckUseSecurity->isChecked();

    if (ui->cbSessionType->currentIndex() == 0) sessionType = 0x02;
    else sessionType = 0x03;

    blockSequenceCounter = 0;
    currentDataPos = 0;

    ui->progressBar->setValue(0);
    ui->btnStartStop->setEnabled(false);
    ui->btnAbort->setEnabled(true);
    ui->btnLoadFile->setEnabled(false);

    udsHandler->setReception(true);
    udsHandler->addFilter(busId, responseId, 0x7FF);

    logMessage("Starting UDS firmware upload...");
    logMessage(QString("Target: 0x%1, Response: 0x%2, Bus: %3")
               .arg(targetId, 0, 16).arg(responseId, 0, 16).arg(busId));
    logMessage(QString("Flash Address: 0x%1, Size: %2 bytes")
               .arg(flashAddr, 0, 16).arg(firmwareSize));

    sendDiagSession();
}

void UDSFirmwareUploaderWindow::stopUpload(bool resetState)
{
    timeoutTimer->stop();
    testerTimer->stop();
    udsHandler->clearAllFilters();
    udsHandler->setReception(false);

    ui->btnStartStop->setEnabled(true);
    ui->btnAbort->setEnabled(false);
    ui->btnLoadFile->setEnabled(true);

    if (resetState)
    {
        currentState = STATE_IDLE;
        ui->btnStartStop->setText("Start Upload");
    }
}

void UDSFirmwareUploaderWindow::sendDiagSession()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::DIAG_CONTROL;
    msg.subFuncLen = 1;
    msg.subFunc = sessionType;

    logMessage("Sending Diagnostic Session Control (0x10) - Sub: 0x" + QString::number(sessionType, 16));
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_DIAG_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendSecuritySeed()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::SECURITY_ACCESS;
    msg.subFuncLen = 1;
    msg.subFunc = (securityLevel * 2) - 1;

    logMessage("Requesting Security Seed (0x27) - Level: " + QString::number(securityLevel));
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_SEED);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendSecurityKey(const QByteArray &seed)
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::SECURITY_ACCESS;
    msg.subFuncLen = 1;
    msg.subFunc = securityLevel * 2;
    msg.payload() = seed;

    logMessage("Sending Security Key (0x27) - " + seed.toHex(' '));
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_KEY_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendRequestDownload()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::REQUEST_DOWNLOAD;
    msg.subFuncLen = 0;
    msg.subFunc = 0;

    QByteArray data;
    data.append(static_cast<char>(0x00));
    uint32_t addr = flashAddr;
    uint32_t size = firmwareSize;

    int addrLen = 4;
    int sizeLen = 4;
    data[0] = (addrLen << 4) | sizeLen;

    for (int i = addrLen - 1; i >= 0; i--)
        data.append((addr >> (i * 8)) & 0xFF);
    for (int i = sizeLen - 1; i >= 0; i--)
        data.append((size >> (i * 8)) & 0xFF);

    msg.payload() = data;

    logMessage(QString("Sending Request Download (0x34) - Addr: 0x%1, Size: %2")
               .arg(addr, 0, 16).arg(size));
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_DOWNLOAD_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendDataBlock()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::TRANSFER_DATA;
    msg.subFuncLen = 0;
    msg.subFunc = 0;

    QByteArray data;
    data.append(blockSequenceCounter & 0xFF);

    int bytesToSend = firmwareSize - currentDataPos;
    if (maxBlockLength > 0 && bytesToSend > maxBlockLength)
        bytesToSend = maxBlockLength;

    data.append(firmwareData.constData() + currentDataPos, bytesToSend);
    msg.payload() = data;

    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_DATA_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendTransferExit()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::REQ_TRANS_EXIT;
    msg.subFuncLen = 0;
    msg.subFunc = 0;

    logMessage("Sending Request Transfer Exit (0x37)");
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_EXIT_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendECUReset()
{
    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::ECU_RESET;
    msg.subFuncLen = 1;
    msg.subFunc = 0x01;

    logMessage("Sending ECU Reset (0x11)");
    udsHandler->sendUDSFrame(msg);
    setState(STATE_WAIT_RESET_RESP);
    timeoutTimer->start();
}

void UDSFirmwareUploaderWindow::sendTesterPresent()
{
    if (currentState != STATE_WAIT_DATA_RESP && currentState != STATE_WAIT_DOWNLOAD_RESP)
        return;

    UDS_MESSAGE msg;
    msg.bus = busId;
    msg.setFrameId(targetId);
    msg.service = UDS_SERVICES::TESTER_PRESENT;
    msg.subFuncLen = 1;
    msg.subFunc = 0x80;

    udsHandler->sendUDSFrame(msg);
}

void UDSFirmwareUploaderWindow::handleNRC(int nrc)
{
    QString desc = udsHandler->getNegativeResponseLong(nrc);
    logMessage(QString("NEGATIVE RESPONSE: 0x%1 - %2").arg(nrc, 0, 16).arg(desc));
    setState(STATE_ERROR);
    stopUpload();
}

void UDSFirmwareUploaderWindow::gotUDSReply(UDS_MESSAGE msg)
{
    if (msg.frameId() != responseId) return;
    if (msg.bus != busId) return;

    const unsigned char *data = reinterpret_cast<const unsigned char *>(msg.payload().constData());
    int dataLen = msg.payload().size();

    if (msg.isErrorReply)
    {
        qDebug() << "UDS FW Upload: NRC service=" << msg.service << " subfunc=" << msg.subFunc;
        handleNRC(msg.subFunc);
        return;
    }

    timeoutTimer->stop();

    switch (currentState)
    {
    case STATE_WAIT_DIAG_RESP:
    {
        logMessage("Diagnostic Session Control response received");
        if (useSecurity)
        {
            sendSecuritySeed();
        }
        else
        {
            sendRequestDownload();
        }
        break;
    }

    case STATE_WAIT_SEED:
    {
        logMessage("Security seed received (" + QString::number(dataLen) + " bytes)");
        QByteArray seed;
        seed.append(reinterpret_cast<const char *>(data), dataLen);
        QByteArray key = seed;
        for (int i = 0; i < key.size(); i++)
            key[i] = ~key[i];
        sendSecurityKey(key);
        break;
    }

    case STATE_WAIT_KEY_RESP:
    {
        logMessage("Security access granted");
        sendRequestDownload();
        break;
    }

    case STATE_WAIT_DOWNLOAD_RESP:
    {
        if (dataLen > 1)
        {
            int sizeLen = data[0] >> 4;
            maxBlockLength = 0;
            for (int i = 0; i < sizeLen && (i + 1) < dataLen; i++)
            {
                maxBlockLength = (maxBlockLength << 8) | data[1 + i];
            }
            logMessage(QString("Download accepted. Max block size: %1 bytes").arg(maxBlockLength));
        }
        else
        {
            maxBlockLength = 4096;
            logMessage("Download accepted (no block size limit advertised)");
        }
        testerTimer->start();
        sendDataBlock();
        break;
    }

    case STATE_WAIT_DATA_RESP:
    {
        progressTimer->start();
        currentDataPos += maxBlockLength > 0 ?
            qMin(maxBlockLength, firmwareSize - currentDataPos) :
            (firmwareSize - currentDataPos);
        blockSequenceCounter++;

        int progress = (currentDataPos * 100) / firmwareSize;
        ui->progressBar->setValue(progress);

        if (currentDataPos >= firmwareSize)
        {
            logMessage("All data blocks transferred");
            testerTimer->stop();
            sendTransferExit();
        }
        else
        {
            sendDataBlock();
        }
        break;
    }

    case STATE_WAIT_EXIT_RESP:
    {
        logMessage("Transfer exit confirmed");
        sendECUReset();
        break;
    }

    case STATE_WAIT_RESET_RESP:
    {
        logMessage("ECU reset acknowledged");
        setState(STATE_COMPLETE);
        stopUpload();
        logMessage("*** Firmware upload complete! ***");
        ui->progressBar->setValue(100);
        break;
    }

    default:
        break;
    }
}

void UDSFirmwareUploaderWindow::timeoutTriggered()
{
    logMessage("TIMEOUT - No response received");
    handleNRC(0x10);
}

void UDSFirmwareUploaderWindow::testerPresentTick()
{
    sendTesterPresent();
    logMessage("Tester Present keep-alive sent");
}

void UDSFirmwareUploaderWindow::advanceState()
{
}