#include "StrHelper.h"

#include <filesystem>
#include <fstream>

#include "OBSWConfig.h"
#include "fsfw/timemanager/Countdown.h"
#include "linux/devices/devicedefinitions/StarTrackerDefinitions.h"
#include "mission/utility/Timestamp.h"
#include "mission/utility/ProgressPrinter.h"

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

StrHelper::~StrHelper() {}

ReturnValue_t StrHelper::initialize() {
#ifdef XIPHOS_Q7S
  sdcMan = SdCardManager::instance();
  if (sdcMan == nullptr) {
    sif::warning << "StrHelper::initialize: Invalid SD Card Manager" << std::endl;
    return RETURN_FAILED;
  }
#endif
  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_READ: {
        result = performFlashRead();
        if (result == RETURN_OK) {
          triggerEvent(FLASH_READ_SUCCESSFUL);
        } else {
          triggerEvent(FLASH_READ_FAILED);
        }
        internalState = InternalState::IDLE;
        break;
      }
      case InternalState::FIRMWARE_UPDATE: {
        result = performFirmwareUpdate();
        if (result == RETURN_OK) {
          triggerEvent(FIRMWARE_UPDATE_SUCCESSFUL);
        } else {
          triggerEvent(FIRMWARE_UPDATE_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) {
#ifdef XIPHOS_Q7S
  ReturnValue_t result = checkPath(fullname);
  if (result != RETURN_OK) {
    return result;
  }
#endif
  uploadImage.uploadFile = fullname;
  if (not std::filesystem::exists(fullname)) {
    return FILE_NOT_EXISTS;
  }
  internalState = InternalState::UPLOAD_IMAGE;
  semaphore.release();
  terminate = false;
  return RETURN_OK;
}

ReturnValue_t StrHelper::startImageDownload(std::string path) {
#ifdef XIPHOS_Q7S
  ReturnValue_t result = checkPath(path);
  if (result != RETURN_OK) {
    return result;
  }
#endif
  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; }

ReturnValue_t StrHelper::startFirmwareUpdate(std::string fullname) {
#ifdef XIPHOS_Q7S
  ReturnValue_t result = checkPath(fullname);
  if (result != RETURN_OK) {
    return result;
  }
#endif
  flashWrite.fullname = fullname;
  if (not std::filesystem::exists(flashWrite.fullname)) {
    return FILE_NOT_EXISTS;
  }
  flashWrite.firstRegion = static_cast<uint8_t>(startracker::FirmwareRegions::FIRST);
  flashWrite.lastRegion = static_cast<uint8_t>(startracker::FirmwareRegions::LAST);
  internalState = InternalState::FIRMWARE_UPDATE;
  semaphore.release();
  terminate = false;
  return RETURN_OK;
}

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

void StrHelper::disableTimestamping() { timestamping = false; }

void StrHelper::enableTimestamping() { timestamping = true; }

ReturnValue_t StrHelper::performImageDownload() {
  ReturnValue_t result;
#if OBSW_DEBUG_STARTRACKER == 1
  ProgressPrinter progressPrinter("Image download", ImageDownload::LAST_POSITION);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  struct DownloadActionRequest downloadReq;
  uint32_t size = 0;
  uint32_t retries = 0;
  std::string image = makeFullFilename(downloadImage.path, downloadImage.filename);
  std::ofstream file(image, 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) {
      file.close();
      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),
               CHUNK_SIZE);
#if OBSW_DEBUG_STARTRACKER == 1
    progressPrinter.print(downloadReq.position);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
    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();
#if OBSW_DEBUG_STARTRACKER == 1
  ProgressPrinter progressPrinter("Image upload", imageSize);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  while ((uploadReq.position + 1) * SIZE_IMAGE_PART < imageSize) {
    if (terminate) {
      file.close();
      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) {
      file.close();
      return RETURN_FAILED;
    }
    result = checkActionReply();
    if (result != RETURN_OK) {
      file.close();
      return result;
    }
#if OBSW_DEBUG_STARTRACKER == 1
    progressPrinter.print((uploadReq.position + 1) * SIZE_IMAGE_PART);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
    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;
  }
#if OBSW_DEBUG_STARTRACKER == 1
  progressPrinter.print((uploadReq.position + 1) * SIZE_IMAGE_PART);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  return RETURN_OK;
}

ReturnValue_t StrHelper::performFirmwareUpdate() {
  using namespace startracker;
  ReturnValue_t result = RETURN_OK;
  result = unlockAndEraseRegions(static_cast<uint32_t>(startracker::FirmwareRegions::FIRST),
                                 static_cast<uint32_t>(startracker::FirmwareRegions::LAST));
  if (result != RETURN_OK) {
    return result;
  }
  result = performFlashWrite();
  return result;
}

ReturnValue_t StrHelper::performFlashWrite() {
  ReturnValue_t result = RETURN_OK;
  uint32_t size = 0;
  uint32_t bytesWritten = 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();
  if (fileSize > FLASH_REGION_SIZE * (flashWrite.lastRegion - flashWrite.firstRegion)) {
    sif::warning << "StrHelper::performFlashWrite: Invalid file" << std::endl;
    return RETURN_FAILED;
  }
#if OBSW_DEBUG_STARTRACKER == 1
  ProgressPrinter progressPrinter("Flash write", fileSize);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  uint32_t fileChunks = fileSize / CHUNK_SIZE;
  bytesWritten = 0;
  req.region = flashWrite.firstRegion;
  req.length = CHUNK_SIZE;
  for (uint32_t idx = 0; idx < fileChunks; idx++) {
    if (terminate) {
      file.close();
      return RETURN_OK;
    }
    file.seekg(idx * CHUNK_SIZE, file.beg);
    file.read(reinterpret_cast<char*>(req.data), CHUNK_SIZE);
    if (bytesWritten + CHUNK_SIZE > FLASH_REGION_SIZE) {
      req.region++;
      bytesWritten = 0;
    }
    req.address = bytesWritten;
    arc_pack_write_action_req(&req, commandBuffer, &size);
    result = sendAndRead(size, req.address);
    if (result != RETURN_OK) {
      file.close();
      return result;
    }
    result = checkActionReply();
    if (result != RETURN_OK) {
      file.close();
      return result;
    }
    bytesWritten += CHUNK_SIZE;
#if OBSW_DEBUG_STARTRACKER == 1
    progressPrinter.print(idx * CHUNK_SIZE);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  }
  uint32_t remainingBytes = fileSize - fileChunks * CHUNK_SIZE;
  file.seekg((fileChunks - 1) * CHUNK_SIZE, file.beg);
  file.read(reinterpret_cast<char*>(req.data), remainingBytes);
  file.close();
  if (bytesWritten + CHUNK_SIZE > FLASH_REGION_SIZE) {
    req.region++;
    bytesWritten = 0;
  }
  req.address = bytesWritten;
  req.length = remainingBytes;
  bytesWritten += remainingBytes;
  arc_pack_write_action_req(&req, commandBuffer, &size);
  result = sendAndRead(size, req.address);
  if (result != RETURN_OK) {
    return result;
  }
  result = checkActionReply();
  if (result != RETURN_OK) {
    return result;
  }
#if OBSW_DEBUG_STARTRACKER == 1
  progressPrinter.print(fileSize);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  return RETURN_OK;
}

ReturnValue_t StrHelper::performFlashRead() {
  ReturnValue_t result;
#if OBSW_DEBUG_STARTRACKER == 1
  ProgressPrinter progressPrinter("Flash read", flashRead.size);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  struct ReadActionRequest req;
  uint32_t bytesRead = 0;
  uint32_t size = 0;
  uint32_t retries = 0;
  Timestamp timestamp;
  std::string fullname = makeFullFilename(flashRead.path, 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.startRegion;
  req.address = 0;
  while (bytesRead < flashRead.size) {
    if (terminate) {
      return RETURN_OK;
    }
    if ((flashRead.size - bytesRead) < CHUNK_SIZE) {
      req.length = flashRead.size - bytesRead;
    } else {
      req.length = CHUNK_SIZE;
    }
    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;
    }
    file.write(reinterpret_cast<const char*>(datalinkLayer.getReply() + FLASH_READ_DATA_OFFSET),
               req.length);
    bytesRead += req.length;
    req.address += req.length;
    if (req.address >= FLASH_REGION_SIZE) {
      req.address = 0;
      req.region++;
    }
    retries = 0;
#if OBSW_DEBUG_STARTRACKER == 1
    progressPrinter.print(bytesRead);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  }
  file.close();
  return RETURN_OK;
}

ReturnValue_t StrHelper::sendAndRead(size_t size, uint32_t parameter, uint32_t delayMs) {
  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) {
    Countdown delay(delayMs);
    delay.resetTimer();
    while (delay.isBusy()) {
    }
    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;
}

#ifdef XIPHOS_Q7S
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;
}
#endif

ReturnValue_t StrHelper::unlockAndEraseRegions(uint32_t from, uint32_t to) {
  ReturnValue_t result = RETURN_OK;
#if OBSW_DEBUG_STARTRACKER == 1
  ProgressPrinter progressPrinter("Unlock and erase", to - from);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  struct UnlockActionRequest unlockReq;
  struct EraseActionRequest eraseReq;
  uint32_t size = 0;
  for (uint32_t idx = from; idx <= to; idx++) {
    unlockReq.region = idx;
    unlockReq.code = startracker::region_secrets::secret[idx];
    arc_pack_unlock_action_req(&unlockReq, commandBuffer, &size);
    sendAndRead(size, unlockReq.region);
    result = checkActionReply();
    if (result != RETURN_OK) {
      sif::warning << "StrHelper::unlockAndEraseRegions: Failed to unlock region with id "
                   << static_cast<unsigned int>(unlockReq.region) << std::endl;
      return result;
    }
    eraseReq.region = idx;
    arc_pack_erase_action_req(&eraseReq, commandBuffer, &size);
    result = sendAndRead(size, eraseReq.region, FLASH_ERASE_DELAY);
    if (result != RETURN_OK) {
      sif::warning << "StrHelper::unlockAndEraseRegions: Failed to erase region with id "
                   << static_cast<unsigned int>(eraseReq.region) << std::endl;
      return result;
    }
#if OBSW_DEBUG_STARTRACKER == 1
    progressPrinter.print(idx - from);
#endif /* OBSW_DEBUG_STARTRACKER == 1 */
  }
  return result;
}

std::string StrHelper::makeFullFilename(std::string path, std::string filename) {
  std::string image;
  Timestamp timestamp;
  if (timestamping) {
    image = path + "/" + timestamp.str() + filename;
  } else {
    image = path + "/" + filename;
  }
  return image;
}