#include "dfu_interface.h"
  #include <string.h>
  #include <ion/src/device/flash.h>
  
  namespace Ion {
  namespace USB {
  namespace Device {
  
  static inline uint32_t min(uint32_t x, uint32_t y) { return (x<y ? x : y); }
  
  void DFUInterface::StatusData::push(Channel * c) const {
    c->push(m_bStatus);
    c->push(m_bwPollTimeout[2]);
    c->push(m_bwPollTimeout[1]);
    c->push(m_bwPollTimeout[0]);
    c->push(m_bState);
    c->push(m_iString);
  }
  
  void DFUInterface::StateData::push(Channel * c) const {
    c->push(m_bState);
  }
  
  void DFUInterface::wholeDataReceivedCallback(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength) {
    if (request->bRequest() == (uint8_t) DFURequest::Download) {
      // Handle a download request
      if (request->wValue() == 0) {
        // The request is a special command
        switch (transferBuffer[0]) {
          case (uint8_t) DFUDownloadCommand::SetAddressPointer:
            setAddressPointerCommand(request, transferBuffer, *transferBufferLength);
            return;
          case (uint8_t) DFUDownloadCommand::Erase:
            eraseCommand(transferBuffer, *transferBufferLength);
            return;
          default:
            m_state = State::dfuERROR;
            m_status = Status::errSTALLEDPKT;
            return;
        }
      }
      if (request->wValue() == 1) {
        m_ep0->stallTransaction();
        return;
      }
      if (request->wLength() > 0) {
        // The request is a "real" download. Compute the writing address.
        m_writeAddress = (request->wValue() - 2) * Endpoint0::MaxTransferSize + m_addressPointer;
        // Store the received data until we copy it on the flash.
        memcpy(m_largeBuffer, transferBuffer, *transferBufferLength);
        m_largeBufferLength = *transferBufferLength;
        m_state = State::dfuDNLOADSYNC;
      }
    }
  }
  
  void DFUInterface::wholeDataSentCallback(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength) {
    if (request->bRequest() == (uint8_t) DFURequest::GetStatus) {
      // Do any needed action after the GetStatus request.
      if (m_state == State::dfuMANIFEST) {
        // Leave DFU routine: Leave DFU, reset device, jump to application code
        leaveDFUAndReset();
      } else if (m_state == State::dfuDNBUSY) {
        if (m_largeBufferLength != 0) {
          // Here, copy the data from the transfer buffer to the flash memory
          writeOnMemory();
        }
        changeAddressPointerIfNeeded();
        eraseMemoryIfNeeded();
        m_state = State::dfuDNLOADIDLE;
      }
    }
  }
  
  bool DFUInterface::processSetupInRequest(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength, uint16_t transferBufferMaxLength) {
    if (Interface::processSetupInRequest(request, transferBuffer, transferBufferLength, transferBufferMaxLength)) {
      return true;
    }
    switch (request->bRequest()) {
      case (uint8_t) DFURequest::Detach:
        m_device->detach();
        return true;
      case (uint8_t) DFURequest::Download:
        return processDownloadRequest(request->wLength(), transferBufferLength);
      case (uint8_t) DFURequest::Upload:
        return processUploadRequest(request, transferBuffer, transferBufferLength, transferBufferMaxLength);
      case (uint8_t) DFURequest::GetStatus:
        return getStatus(request, transferBuffer, transferBufferLength, transferBufferMaxLength);
      case (uint8_t) DFURequest::ClearStatus:
        return clearStatus(request, transferBuffer, transferBufferLength, transferBufferMaxLength);
      case (uint8_t) DFURequest::GetState:
        return getState(transferBuffer, transferBufferLength, transferBufferMaxLength);
      case (uint8_t) DFURequest::Abort:
        return dfuAbort(transferBufferLength);
    }
    return false;
  }
  
  bool DFUInterface::processDownloadRequest(uint16_t wLength, uint16_t * transferBufferLength) {
    if (m_state != State::dfuIDLE && m_state != State::dfuDNLOADIDLE) {
      m_state = State::dfuERROR;
      m_status = Status::errNOTDONE;
      m_ep0->stallTransaction();
      return false;
    }
    if (wLength == 0) {
      // Leave DFU routine: Reset the device and jump to application code
      m_state = State::dfuMANIFESTSYNC;
    } else {
      // Prepare to receive the download data
      m_ep0->clearForOutTransactions(wLength);
      m_state = State::dfuDNLOADSYNC;
    }
    return true;
  }
  
  bool DFUInterface::processUploadRequest(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength, uint16_t transferBufferMaxLength) {
    if (m_state != State::dfuIDLE && m_state !=  State::dfuUPLOADIDLE) {
      m_ep0->stallTransaction();
      return false;
    }
    if (request->wValue() == 0) {
      /* The host requests to read the commands supported by the bootloader. After
       * receiving this command, the device  should returns N bytes representing
       * the command codes for :
       * Get command / Set Address Pointer / Erase / Read Unprotect
       * We no not need it for now. */
       return false;
    } else if (request->wValue() == 1) {
      m_ep0->stallTransaction();
      return false;
    } else {
      /* We decided to never protect Read operation. Else we would have to check
       * here it is not protected before reading. */
  
      // Compute the reading address
      uint32_t readAddress = (request->wValue() - 2) * Endpoint0::MaxTransferSize + m_addressPointer;
      // Copy the requested memory zone into the transfer buffer.
      uint16_t copySize = min(transferBufferMaxLength, request->wLength());
      memcpy(transferBuffer, (void *)readAddress, copySize);
      *transferBufferLength = copySize;
    }
    m_state = State::dfuUPLOADIDLE;
    return true;
  }
  
  void DFUInterface::setAddressPointerCommand(SetupPacket * request, uint8_t * transferBuffer, uint16_t transferBufferLength) {
    assert(transferBufferLength == 5);
    // Compute the new address but change it after the next getStatus request.
    m_potentialNewAddressPointer = transferBuffer[1]
      + (transferBuffer[2] << 8)
      + (transferBuffer[3] << 16)
      + (transferBuffer[4] << 24);
    m_state = State::dfuDNLOADSYNC;
  }
  
  void DFUInterface::changeAddressPointerIfNeeded() {
    if (m_potentialNewAddressPointer == 0) {
      // There was no address change waiting.
      return;
    }
    // If there is a new address pointer waiting, change the pointer address.
    m_addressPointer = m_potentialNewAddressPointer;
    m_potentialNewAddressPointer = 0;
    m_state = State::dfuDNLOADIDLE;
    m_status = Status::OK;
  }
  
  void DFUInterface::eraseCommand(uint8_t * transferBuffer, uint16_t transferBufferLength) {
    /* We determine whether the commands asks for a mass erase or which sector to
     * erase. The erase must be done after the next getStatus request. */
    m_state = State::dfuDNLOADSYNC;
  
    if (transferBufferLength == 1) {
      // Mass erase
      m_erasePage = Flash::Device::NumberOfSectors;
      return;
    }
  
    // Sector erase
    assert(transferBufferLength == 5);
  
    uint32_t eraseAddress = transferBuffer[1]
      + (transferBuffer[2] << 8)
      + (transferBuffer[3] << 16)
      + (transferBuffer[4] << 24);
  
    m_erasePage = Flash::Device::SectorAtAddress(eraseAddress);
  
    if (m_erasePage < 0) {
      // Unrecognized sector
      m_state = State::dfuERROR;
      m_status = Status::errTARGET;
    }
  }
  
  
  void DFUInterface::eraseMemoryIfNeeded() {
    if (m_erasePage < 0) {
      // There was no erase waiting.
      return;
    }
  
    if (m_erasePage == Ion::Flash::Device::NumberOfSectors) {
      Flash::Device::MassErase();
    } else {
      Flash::Device::EraseSector(m_erasePage);
    }
  
    /* Put an out of range value in m_erasePage to indicate that no erase is
     * waiting. */
    m_erasePage = -1;
    m_state = State::dfuDNLOADIDLE;
    m_status = Status::OK;
  }
  
  void DFUInterface::writeOnMemory() {
    if (m_writeAddress >= k_flashStartAddress && m_writeAddress <= k_flashEndAddress) {
      // Write to the Flash memory
      Flash::Device::WriteMemory(m_largeBuffer, reinterpret_cast<uint8_t *>(m_writeAddress), m_largeBufferLength);
    } else if (m_writeAddress >= k_sramStartAddress && m_writeAddress <= k_sramEndAddress) {
      // Write on SRAM
      // FIXME We should check that we are not overriding the current instructions.
      memcpy((void *)m_writeAddress, m_largeBuffer, m_largeBufferLength);
    } else {
      // Invalid write address
      m_largeBufferLength = 0;
      m_state = State::dfuERROR;
      m_status = Status::errTARGET;
      return;
    }
  
    // Reset the buffer length
    m_largeBufferLength = 0;
    // Change the interface state and status
    m_state = State::dfuDNLOADIDLE;
    m_status = Status::OK;
  }
  
  
  bool DFUInterface::getStatus(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength, uint16_t transferBufferMaxLength) {
    // Change the status if needed
    if (m_state == State::dfuMANIFESTSYNC) {
      m_state = State::dfuMANIFEST;
    } else if (m_state == State::dfuDNLOADSYNC) {
      m_state = State::dfuDNBUSY;
    }
    // Copy the status on the TxFifo
    *transferBufferLength = StatusData(m_status, m_state).copy(transferBuffer, transferBufferMaxLength);
    return true;
  }
  
  bool DFUInterface::clearStatus(SetupPacket * request, uint8_t * transferBuffer, uint16_t * transferBufferLength, uint16_t transferBufferMaxLength) {
    m_status = Status::OK;
    m_state = State::dfuIDLE;
    return getStatus(request, transferBuffer, transferBufferLength, transferBufferMaxLength);
  }
  
  bool DFUInterface::getState(uint8_t * transferBuffer, uint16_t * transferBufferLength, uint16_t maxSize) {
    *transferBufferLength = StateData(m_state).copy(transferBuffer, maxSize);
    return true;
  }
  
  bool DFUInterface::dfuAbort(uint16_t * transferBufferLength) {
    m_status = Status::OK;
    m_state = State::dfuIDLE;
    *transferBufferLength = 0;
    return true;
  }
  
  void DFUInterface::leaveDFUAndReset() {
    m_device->setResetOnDisconnect(true);
    m_device->detach();
  }
  
  }
  }
  }