#include "StrHelper.h"

#include "mission/utility/Timestamp.h"

#include <fstream>
#include <filesystem>

StrHelper::StrHelper(object_id_t objectId) : SystemObject(objectId){

}

StrHelper::~StrHelper() {
}

ReturnValue_t StrHelper::initialize() {
    sdcMan = SdCardManager::instance();
    if (sdcMan == nullptr) {
        sif::warning << "StrHelper::initialize: Invalid SD Card Manager" << std::endl;
        return RETURN_FAILED;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::performOperation(uint8_t operationCode) {
    ReturnValue_t result = RETURN_OK;
    semaphore.acquire();
    while(true) {
        switch(internalState) {
        case InternalState::IDLE: {
            semaphore.acquire();
            break;
        }
        case InternalState::UPLOAD_IMAGE: {
            result = performImageUpload();
            if (result == RETURN_OK){
                triggerEvent(IMAGE_UPLOAD_SUCCESSFUL);
            }
            else {
                triggerEvent(IMAGE_UPLOAD_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        case InternalState::DOWNLOAD_IMAGE: {
            result = performImageDownload();
            if (result == RETURN_OK){
                triggerEvent(IMAGE_DOWNLOAD_SUCCESSFUL);
            }
            else {
                triggerEvent(IMAGE_DOWNLOAD_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        case InternalState::FLASH_WRITE: {
            result = performFlashWrite();
            if (result == RETURN_OK){
                triggerEvent(FLASH_WRITE_SUCCESSFUL);
            }
            else {
                triggerEvent(FLASH_WRITE_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        case InternalState::FLASH_READ: {
            result = performFlashRead();
            if (result == RETURN_OK){
                triggerEvent(FLASH_READ_SUCCESSFUL);
            }
            else {
                triggerEvent(FLASH_READ_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        case InternalState::DOWNLOAD_FPGA_IMAGE: {
            result = performFpgaDownload();
            if (result == RETURN_OK){
                triggerEvent(FPGA_DOWNLOAD_SUCCESSFUL);
            }
            else {
                triggerEvent(FPGA_DOWNLOAD_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        case InternalState::UPLOAD_FPGA_IMAGE: {
            result = performFpgaUpload();
            if (result == RETURN_OK){
                triggerEvent(FPGA_UPLOAD_SUCCESSFUL);
            }
            else {
                triggerEvent(FPGA_UPLOAD_FAILED);
            }
            internalState = InternalState::IDLE;
            break;
        }
        default:
            sif::debug << "StrHelper::performOperation: Invalid state" << std::endl;
            break;
        }
    }
}

ReturnValue_t StrHelper::setComIF(DeviceCommunicationIF* communicationInterface_) {
    uartComIF = dynamic_cast<UartComIF*>(communicationInterface_);
    if (uartComIF == nullptr) {
        sif::warning << "StrHelper::initialize: Invalid uart com if" << std::endl;
        return RETURN_FAILED;
    }
    return RETURN_OK;
}

void StrHelper::setComCookie(CookieIF* comCookie_) {
    comCookie = comCookie_;
}

ReturnValue_t StrHelper::startImageUpload(std::string fullname) {
    ReturnValue_t result = checkPath(fullname);
    if (result != RETURN_OK) {
        return result;
    }
    uploadImage.uploadFile = fullname;
    if(not std::filesystem::exists(fullname)) {
        return FILE_NOT_EXISTS;
    }
    internalState = InternalState::UPLOAD_FPGA_IMAGE;
    semaphore.release();
    terminate = false;
    return RETURN_OK;
}

ReturnValue_t StrHelper::startImageDownload(std::string path) {
    ReturnValue_t result = checkPath(path);
    if (result != RETURN_OK) {
        return result;
    }
    if(not std::filesystem::exists(path)) {
        return PATH_NOT_EXISTS;
    }
    downloadImage.path = path;
    internalState = InternalState::DOWNLOAD_IMAGE;
    terminate = false;
    semaphore.release();
    return RETURN_OK;
}

void StrHelper::stopProcess() {
    terminate = true;
}

void StrHelper::setDownloadImageName(std::string filename) {
    downloadImage.filename = filename;
}

void StrHelper::setFlashReadFilename(std::string filename) {
    flashRead.filename = filename;
}

void StrHelper::setDownloadFpgaImage(std::string filename) {
    fpgaDownload.fileName = filename;
}

ReturnValue_t StrHelper::startFlashWrite(std::string fullname, uint8_t region,
        uint32_t address) {
    ReturnValue_t result = checkPath(fullname);
    if (result != RETURN_OK) {
        return result;
    }
    flashWrite.fullname = fullname;
    if(not std::filesystem::exists(flashWrite.fullname)) {
        return FILE_NOT_EXISTS;
    }
    flashWrite.address = address;
    flashWrite.region = region;
    internalState = InternalState::FLASH_WRITE;
    semaphore.release();
    terminate = false;
    return RETURN_OK;
}

ReturnValue_t StrHelper::startFlashRead(std::string path, uint8_t region,
        uint32_t address, uint32_t length) {
    ReturnValue_t result = checkPath(path);
    if (result != RETURN_OK) {
        return result;
    }
    flashRead.path = path;
    if(not std::filesystem::exists(flashRead.path)) {
        return FILE_NOT_EXISTS;
    }
    flashRead.address = address;
    flashRead.region = region;
    flashRead.size = length;
    internalState = InternalState::FLASH_READ;
    semaphore.release();
    terminate = false;
    return RETURN_OK;
}

ReturnValue_t StrHelper::startFpgaDownload(std::string path, uint32_t startPosition,
            uint32_t length) {
    fpgaDownload.path = path;
    fpgaDownload.startPosition = startPosition;
    fpgaDownload.length = length;
    internalState = InternalState::DOWNLOAD_FPGA_IMAGE;
    semaphore.release();
    terminate = false;
    return RETURN_OK;
}

ReturnValue_t StrHelper::startFpgaUpload(std::string uploadFile) {
    fpgaUpload.uploadFile = uploadFile;
    internalState = InternalState::UPLOAD_FPGA_IMAGE;
    semaphore.release();
    terminate = false;
    return RETURN_OK;
}

ReturnValue_t StrHelper::performImageDownload() {
    ReturnValue_t result;
    struct DownloadActionRequest downloadReq;
    uint32_t size = 0;
    uint32_t retries = 0;
    Timestamp timestamp;
    std::string image = downloadImage.path + "/" + timestamp.str() + downloadImage.filename ;
    std::ofstream file(image, std::ios_base::app | std::ios_base::out);
    if(not std::filesystem::exists(image)) {
        return FILE_CREATION_FAILED;
    }
    downloadReq.position = 0;
    while(downloadReq.position < ImageDownload::LAST_POSITION) {
        if (terminate) {
            return RETURN_OK;
        }
        arc_pack_download_action_req(&downloadReq, commandBuffer, &size);
        result = sendAndRead(size, downloadReq.position);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        result = checkActionReply();
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        result = checkReplyPosition(downloadReq.position);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        file.write(reinterpret_cast<const char*>(datalinkLayer.getReply() + IMAGE_DATA_OFFSET),
                IMAGE_DATA_SIZE);
        downloadReq.position++;
        retries = 0;
    }
    file.close();
    return RETURN_OK;
}

ReturnValue_t StrHelper::performImageUpload() {
    ReturnValue_t result = RETURN_OK;
    uint32_t size = 0;
    uint32_t imageSize = 0;
    struct UploadActionRequest uploadReq;
    uploadReq.position = 0;
    std::memset(&uploadReq.data, 0, sizeof(uploadReq.data));
    if (not std::filesystem::exists(uploadImage.uploadFile)) {
        triggerEvent(STR_HELPER_FILE_NOT_EXISTS, static_cast<uint32_t>(internalState));
        internalState = InternalState::IDLE;
        return RETURN_FAILED;
    }
    std::ifstream file(uploadImage.uploadFile, std::ifstream::binary);
    // Set position of next character to end of file input stream
    file.seekg(0, file.end);
    // tellg returns position of character in input stream
    imageSize = file.tellg();
    while((uploadReq.position + 1) * SIZE_IMAGE_PART < imageSize) {
        if (terminate) {
            return RETURN_OK;
        }
        file.seekg(uploadReq.position * SIZE_IMAGE_PART, file.beg);
        file.read(reinterpret_cast<char*>(uploadReq.data), SIZE_IMAGE_PART);
        arc_pack_upload_action_req(&uploadReq, commandBuffer, &size);
        result = sendAndRead(size, uploadReq.position);
        if (result != RETURN_OK) {
            return RETURN_FAILED;
        }
        result = checkActionReply();
        if (result != RETURN_OK) {
            return result;
        }
        uploadReq.position++;
    }
    std::memset(uploadReq.data, 0, sizeof(uploadReq.data));
    uint32_t remainder = imageSize - uploadReq.position * SIZE_IMAGE_PART;
    file.seekg(uploadReq.position * SIZE_IMAGE_PART, file.beg);
    file.read(reinterpret_cast<char*>(uploadReq.data), remainder);
    file.close();
    uploadReq.position++;
    arc_pack_upload_action_req(&uploadReq, commandBuffer, &size);
    result = sendAndRead(size, uploadReq.position);
    if (result != RETURN_OK) {
        return RETURN_FAILED;
    }
    result = checkActionReply();
    if (result != RETURN_OK) {
        return result;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::performFlashWrite() {
    ReturnValue_t result = RETURN_OK;
    uint32_t size = 0;
    uint32_t remainingBytes = 0;
    uint32_t fileSize = 0;
    struct WriteActionRequest req;
    if (not std::filesystem::exists(flashWrite.fullname)) {
        triggerEvent(STR_HELPER_FILE_NOT_EXISTS, static_cast<uint32_t>(internalState));
        internalState = InternalState::IDLE;
        return RETURN_FAILED;
    }
    std::ifstream file(flashWrite.fullname, std::ifstream::binary);
    file.seekg(0, file.end);
    fileSize = file.tellg();
    remainingBytes = fileSize;
    req.region = flashWrite.region;
    req.address = flashWrite.address;
    req.length = MAX_FLASH_DATA;
    while(remainingBytes >= MAX_FLASH_DATA) {
        if (terminate) {
            return RETURN_OK;
        }
        file.seekg(fileSize - remainingBytes, file.beg);
        file.read(reinterpret_cast<char*>(req.data), MAX_FLASH_DATA);
        arc_pack_write_action_req(&req, commandBuffer, &size);
        result = sendAndRead(size, req.address);
        if (result != RETURN_OK) {
            return RETURN_FAILED;
        }
        result = checkFlashActionReply(req.region, req.address, req.length);
        if (result != RETURN_OK) {
            return result;
        }
        remainingBytes = remainingBytes - MAX_FLASH_DATA;
    }
    file.seekg(fileSize - remainingBytes, file.beg);
    file.read(reinterpret_cast<char*>(req.data), remainingBytes);
    file.close();
    arc_pack_write_action_req(&req, commandBuffer, &size);
    result = sendAndRead(size, req.address);
    if (result != RETURN_OK) {
        return RETURN_FAILED;
    }
    result = checkFlashActionReply(req.region, req.address, req.length);
    if (result != RETURN_OK) {
        return result;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::performFlashRead() {
    ReturnValue_t result;
    struct ReadActionRequest req;
    uint32_t bytesRead = 0;
    uint32_t size = 0;
    uint32_t retries = 0;
    Timestamp timestamp;
    std::string fullname = flashRead.path + "/" + timestamp.str() + flashRead.filename ;
    std::ofstream file(fullname, std::ios_base::app | std::ios_base::out);
    if (not std::filesystem::exists(fullname)) {
        return FILE_CREATION_FAILED;
    }
    req.region = flashRead.region;
    while(bytesRead < flashRead.size) {
        if (terminate) {
            return RETURN_OK;
        }
        if ((flashRead.size - bytesRead) < MAX_FLASH_DATA) {
            req.length = flashRead.size - bytesRead;
        }
        else {
            req.length = MAX_FLASH_DATA;
        }
        req.address = flashRead.address + bytesRead;
        arc_pack_read_action_req(&req, commandBuffer, &size);
        result = sendAndRead(size, req.address);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        result = checkActionReply();
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        result = checkFlashActionReply(req.region, req.address, req.length);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        file.write(reinterpret_cast<const char*>(datalinkLayer.getReply() + FLASH_READ_DATA_OFFSET),
                req.length);
        bytesRead += req.length;
        retries = 0;
    }
    file.close();
    return RETURN_OK;
}

ReturnValue_t StrHelper::performFpgaDownload() {
    ReturnValue_t result;
    struct DownloadFPGAImageActionRequest req;
    uint32_t size = 0;
    uint32_t retries = 0;
    Timestamp timestamp;
    std::string image = fpgaDownload.path + "/" + timestamp.str() + fpgaDownload.fileName;
    std::ofstream file(image, std::ios_base::app | std::ios_base::out);
    if(not std::filesystem::exists(image)) {
        return FILE_CREATION_FAILED;
    }
    req.pos = fpgaDownload.startPosition;
    while(req.pos < fpgaDownload.length) {
        if (terminate) {
            return RETURN_OK;
        }
        if (fpgaDownload.length - req.pos >= FpgaDownload::MAX_DATA) {
            req.length = FpgaDownload::MAX_DATA;
        }
        else {
            req.length = fpgaDownload.length - req.pos;
        }
        arc_pack_downloadfpgaimage_action_req(&req, commandBuffer, &size);
        result = sendAndRead(size, req.pos);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        result = checkFpgaActionReply(req.pos, req.length);
        if (result != RETURN_OK) {
            if (retries < CONFIG_MAX_DOWNLOAD_RETRIES) {
                uartComIF->flushUartRxBuffer(comCookie);
                retries++;
                continue;
            }
            file.close();
            return result;
        }
        file.write(reinterpret_cast<const char*>(datalinkLayer.getReply() + FpgaDownload::DATA_OFFSET),
                req.length);
        req.pos += req.length;
        retries = 0;
    }
    file.close();
    return RETURN_OK;
}

ReturnValue_t StrHelper::performFpgaUpload() {
    ReturnValue_t result = RETURN_OK;
    uint32_t commandSize = 0;
    uint32_t bytesUploaded = 0;
    uint32_t fileSize = 0;
    struct UploadFPGAImageActionRequest req;
    if (not std::filesystem::exists(fpgaUpload.uploadFile)) {
        triggerEvent(STR_HELPER_FILE_NOT_EXISTS, static_cast<uint32_t>(internalState));
        internalState = InternalState::IDLE;
        return RETURN_FAILED;
    }
    std::ifstream file(flashWrite.fullname, std::ifstream::binary);
    file.seekg(0, file.end);
    fileSize = file.tellg();
    req.pos = 0;
    while(bytesUploaded <= fileSize) {
        if (terminate) {
            return RETURN_OK;
        }
        if (fileSize - bytesUploaded > FpgaUpload::MAX_DATA) {
            req.length = FpgaUpload::MAX_DATA;
        }
        else {
            req.length = fileSize - bytesUploaded;
        }
        file.seekg(bytesUploaded, file.beg);
        file.read(reinterpret_cast<char*>(req.data), req.length);
        arc_pack_uploadfpgaimage_action_req(&req, commandBuffer, &commandSize);
        result = sendAndRead(commandSize, req.pos);
        if (result != RETURN_OK) {
            return RETURN_FAILED;
        }
        result = checkFpgaActionReply(req.pos, req.length);
        if (result != RETURN_OK) {
            return result;
        }
        bytesUploaded += req.length;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::sendAndRead(size_t size, uint32_t parameter) {
    ReturnValue_t result = RETURN_OK;
    ReturnValue_t decResult = RETURN_OK;
    size_t receivedDataLen = 0;
    uint8_t *receivedData = nullptr;
    size_t bytesLeft = 0;
    uint32_t missedReplies = 0;
    datalinkLayer.encodeFrame(commandBuffer, size);
    result = uartComIF->sendMessage(comCookie, datalinkLayer.getEncodedFrame(),
        datalinkLayer.getEncodedLength());
    if (result != RETURN_OK) {
        sif::warning << "StrHelper::sendAndRead: Failed to send packet" << std::endl;
        triggerEvent(STR_HELPER_SENDING_PACKET_FAILED, result, parameter);
        return RETURN_FAILED;
    }
    decResult = ArcsecDatalinkLayer::DEC_IN_PROGRESS;
    while (decResult == ArcsecDatalinkLayer::DEC_IN_PROGRESS) {
        result = uartComIF->requestReceiveMessage(comCookie,
                StarTracker::MAX_FRAME_SIZE * 2 + 2);
        if (result != RETURN_OK) {
            sif::warning << "StrHelper::sendAndRead: Failed to request reply" << std::endl;
            triggerEvent(STR_HELPER_REQUESTING_MSG_FAILED, result, parameter);
            return RETURN_FAILED;
        }
        result = uartComIF->readReceivedMessage(comCookie, &receivedData, &receivedDataLen);
        if (result != RETURN_OK) {
            sif::warning << "StrHelper::sendAndRead: Failed to read received message" << std::endl;
            triggerEvent(STR_HELPER_READING_REPLY_FAILED, result, parameter);
            return RETURN_FAILED;
        }
        if (receivedDataLen == 0 && missedReplies < MAX_POLLS) {
            missedReplies++;
            continue;
        }
        else if ((receivedDataLen == 0) && (missedReplies >= MAX_POLLS)) {
            triggerEvent(STR_HELPER_NO_REPLY, parameter);
            return RETURN_FAILED;
        }
        else {
            missedReplies = 0;
        }
        decResult = datalinkLayer.decodeFrame(receivedData, receivedDataLen, &bytesLeft);
        if (bytesLeft != 0) {
            // This should never happen
            sif::warning << "StrHelper::sendAndRead: Bytes left after decoding" << std::endl;
            triggerEvent(STR_HELPER_COM_ERROR, result, parameter);
            return RETURN_FAILED;
        }
    }
    if (decResult != RETURN_OK) {
        triggerEvent(STR_HELPER_DEC_ERROR, decResult, parameter);
        return RETURN_FAILED;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::checkActionReply() {
    uint8_t type = datalinkLayer.getReplyFrameType();
    if (type != TMTC_ACTIONREPLY) {
        sif::warning << "StrHelper::checkActionReply: Received reply with invalid type ID"
                << std::endl;
        return INVALID_TYPE_ID;
    }
    uint8_t status = datalinkLayer.getStatusField();
    if (status != ArcsecDatalinkLayer::STATUS_OK) {
        sif::warning << "StrHelper::checkActionReply: Status failure: "
                << static_cast<unsigned int>(status) << std::endl;
        return STATUS_ERROR;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::checkReplyPosition(uint32_t expectedPosition) {
    uint32_t receivedPosition = 0;
    std::memcpy(&receivedPosition, datalinkLayer.getReply() + POS_OFFSET, sizeof(receivedPosition));
    if (receivedPosition != expectedPosition) {
        triggerEvent(POSITION_MISMATCH, receivedPosition);
        return RETURN_FAILED;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::checkFlashActionReply(uint8_t region_, uint32_t address_,
        uint16_t length_) {
    ReturnValue_t result = RETURN_OK;
    result = checkActionReply();
    if (result != RETURN_OK) {
        return result;
    }
    const uint8_t* data = datalinkLayer.getReply();
    uint8_t region = *(data + REGION_OFFSET);
    uint32_t address;
    const uint8_t* addressData = data + ADDRESS_OFFSET;
    size_t size = sizeof(address);
    result = SerializeAdapter::deSerialize(&address, &addressData, &size,
                SerializeIF::Endianness::LITTLE);
    if (result != RETURN_OK) {
        sif::warning << "StrHelper::checkFlashActionReply: Deserialization of address failed"
                << std::endl;
        return result;
    }
    uint16_t length = 0;
    size = sizeof(length);
    const uint8_t* lengthData = data + LENGTH_OFFSET;
    result = SerializeAdapter::deSerialize(&length, lengthData, &size,
                SerializeIF::Endianness::LITTLE);
    if (result != RETURN_OK) {
        sif::warning << "StrHelper::checkFlashActionReply: Deserialization of length failed"
                << std::endl;
    }
    if (region != region_) {
        return REGION_MISMATCH;
    }
    if (address != address_) {
        return ADDRESS_MISMATCH;
    }
    if (length != length_) {
        return LENGTH_MISMATCH;
    }
    return RETURN_OK;
}

ReturnValue_t StrHelper::checkFpgaActionReply(uint32_t expectedPosition,
        uint32_t expectedLength) {
    ReturnValue_t result = RETURN_OK;
    result = checkActionReply();
    if (result != RETURN_OK) {
        return result;
    }
    const uint8_t* data = datalinkLayer.getReply() + ACTION_DATA_OFFSET;
    uint32_t position;
    size_t size = sizeof(position);
    result = SerializeAdapter::deSerialize(&position, &data, &size,
            SerializeIF::Endianness::LITTLE);
    if (result != RETURN_OK) {
        sif::warning << "StrHelper::checkFpgaActionReply: Deserialization of position failed"
                << std::endl;
        return result;
    }
    uint32_t length;
    size = sizeof(length);
    result = SerializeAdapter::deSerialize(&length, &data, &size,
                SerializeIF::Endianness::LITTLE);
    if (result != RETURN_OK) {
        sif::warning << "StrHelper::checkFpgaActionReply: Deserialization of length failed"
                << std::endl;
        return result;
    }
    return result;
}

ReturnValue_t StrHelper::checkPath(std::string name) {
    if (name.substr(0, sizeof(SdCardManager::SD_0_MOUNT_POINT))
            == std::string(SdCardManager::SD_0_MOUNT_POINT)) {
        if (!sdcMan->isSdCardMounted(sd::SLOT_0)) {
            sif::warning << "StrHelper::checkPath: SD card 0 not mounted" << std::endl;
            return SD_NOT_MOUNTED;
        }
    } else if (name.substr(0, sizeof(SdCardManager::SD_1_MOUNT_POINT))
            == std::string(SdCardManager::SD_1_MOUNT_POINT)) {
        if (!sdcMan->isSdCardMounted(sd::SLOT_0)) {
            sif::warning << "StrHelper::checkPath: SD card 1 not mounted" << std::endl;
            return SD_NOT_MOUNTED;
        }
    }
    return RETURN_OK;
}