/* Motenpoche 
 *
 * (c) 2023 Daniele Lacamera <root@danielinux.net>
 *
 *
 * Motenpoche is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * Motenpoche is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
 *
 */

#include <stdint.h>
#include "class/cdc/cdc.h"
#include "class/cdc/cdc_device.h"
#include "bsp/board.h"
#include "wolfssl/wolfcrypt/settings.h"
#include "cryptoengine.h"
#include "flash.h"
#include "fsm.h"
#include "hardware/gpio.h"

#include "display.h"


extern uint16_t flash_info;
extern int thread_control_ui;
extern void system_reboot(void);

extern uint8_t checked_digest[];



// Invoked when cdc when line state changed e.g connected/disconnected
void tud_cdc_line_state_cb(uint8_t itf, bool dtr, bool rts)
{
  (void) itf;
  (void) rts;

  // TODO set some indicator
  if ( dtr )
  {
    // Terminal connected

  }else
  {
    // Terminal disconnected
  }
}


static uint8_t cdc_rxbuf[MAX_CDC_CMD];
static uint32_t cdc_rx_idx = 0;
static uint32_t cdc_rx_expected = 0xFFFFFFFF;
static int rx_busy = 0;

static void send_OK(void)
{
    struct cdc_packet_hdr OKhdr;
    OKhdr.magic = VAULT_MAGIC;
    OKhdr.cmd = CDC_OK;
    OKhdr.len = 0;
    tud_cdc_write((void *)&OKhdr, 8);
    tud_cdc_write_flush();
}

static void send_fail(char *errmsg)
{
    uint8_t failbuf[256 + 8];
    struct cdc_packet_hdr *hdr = (struct cdc_packet_hdr*)failbuf;

    memset(failbuf, 0, 256 + 8);
    snprintf(failbuf + 8, 256, "%s", errmsg);
    hdr->magic = VAULT_MAGIC;
    hdr->cmd = CDC_FAIL;
    hdr->len = strlen(failbuf + 8);
    tud_cdc_write((void *)hdr, 8 + strlen(failbuf + 8));
    tud_cdc_write_flush();
}

static void parse_cmd(int rxbuf_size)
{
    struct cdc_packet_hdr *hdr;
    uint16_t svc_active, svc_avail;
    hdr = (struct cdc_packet_hdr *)(cdc_rxbuf);
    switch(hdr->cmd & 0xFF) {
        case CDC_STATUS:
            int i;
            enum vault_state st = fsm_get();
            struct vault_status *vst;
            uint32_t n_st = (uint32_t)st;
            uint8_t status_buffer[sizeof(struct cdc_packet_hdr) +
                sizeof(struct vault_status)];

            hdr = (struct cdc_packet_hdr *)status_buffer;
            vst = (struct vault_status *)(status_buffer +
                    sizeof(struct cdc_packet_hdr));

            hdr->magic = VAULT_MAGIC;
            hdr->cmd = CDC_STATUS;
            hdr->len = sizeof(struct cdc_packet_hdr) + sizeof(struct vault_status);

            memset(vst, 0, sizeof(struct vault_status));
            vst->state = fsm_get(); 
            if (cryptoengine_check_vault() == 0) {
                cryptoengine_fill_vault_status(vst);
            }
            svc_active = 0;
            svc_avail = 0;
            cryptoengine_service_count(&svc_active, &svc_avail);
            vst->services_active = svc_active;
            vst->first_avail = svc_avail;
            for (i = 0; i < sizeof(status_buffer);) {
                int len = 64;
                if ((sizeof(status_buffer) - i) < len)
                    len = (sizeof(status_buffer) - i);
                tud_cdc_write((void *)status_buffer + i, len);
                tud_cdc_write_flush();
                i += len;
            }
            break;
        case CDC_TOFU_INIT:
            if (fsm_get() != VAULT_TOFU) {
                send_fail("Already provisioned");
                break;
            }
            if (hdr->len == sizeof(struct vault_header)) {
                struct vault_header *vh = (struct vault_header *)
                    (cdc_rxbuf + sizeof(struct cdc_packet_hdr));
                if (rxbuf_size != (hdr->len + sizeof(struct cdc_packet_hdr))) {
                    char msg[64];
                    sprintf(msg, "packet len mismatch exp: %d got %d",
                            (hdr->len + sizeof(struct cdc_packet_hdr)),
                             rxbuf_size);
                    send_fail(msg);
                }
                if (vh->magic != VAULT_MAGIC) {
                    send_fail("TOFU Rejected.");
                    break;
                }
                if (cryptoengine_hdr_sha_check(vh) < 0) {
                    send_fail("TOFU Rejected.");
                    break;
                }

                /*  Take control of display/buttons */
                thread_control_ui = 1;
                sleep_ms(100);
                display_clear();
                display_text(0, "Provisioning...");

                flash_sector_erase(VAULT_FLASH_OFFSET);
                flash_write(VAULT_FLASH_OFFSET, cdc_rxbuf + sizeof(struct cdc_packet_hdr), sizeof(struct vault_header));

                display_text(1, "Confirming...  ");
                send_OK();
                display_text(2, "Rebooting...   ");
                system_reboot();
                break;
            } else {
                send_fail("TOFU Rejected.");
                break;
            }
            break;
        case CDC_CHALLENGE:
            break;

        case CDC_REKEY:
            break;
        case CDC_ADDSERV:
            int ret;
            uint16_t count, idx;
            struct vault_service *vs;
            gpio_put(16,1);
            if (rxbuf_size != SVC_SIZE + sizeof(struct cdc_packet_hdr)) {
                char msg[64];
                sprintf(msg, "packet len mismatch exp: %d got %d", hdr->len, rxbuf_size);
                send_fail(msg);
                break;
            } 
            vs = (struct vault_service *)(cdc_rxbuf +sizeof(struct cdc_packet_hdr));
            /* Retrieve the implicit destination slot (current first slot available) */
            cryptoengine_service_count(&count, &idx);
            if (cryptoengine_check_vault() < 0)
                break;
            ret = cryptoengine_import_service(vs, idx);
            if (ret == 0)
                send_OK();
            else {
                char msg[40];
                sprintf(msg, "failed import: svc %s retval %d", vs->name, ret);
                send_fail(msg);
            }
            break;

        case CDC_DELSERV:
            break;
        default:
            send_fail("Bad cmd");

    }
    rx_busy = 0;

}


// Invoked when CDC interface received data from host
void tud_cdc_rx_cb(uint8_t itf)
{
  (void) itf;
  char buf[64];
  uint32_t count;
  struct cdc_packet_hdr *hdr;

  if (rx_busy) {
      tud_cdc_read_flush();
      return;
  }
  count = tud_cdc_read(buf, sizeof(buf));


  if (count > 0) {
      if (count + cdc_rx_idx > MAX_CDC_CMD) {
          cdc_rx_idx = 0;
      }
      memcpy(cdc_rxbuf + cdc_rx_idx, buf, count);
      hdr = (struct cdc_packet_hdr *)cdc_rxbuf;
      if (cdc_rx_idx == 0) {
          if (hdr->magic != VAULT_MAGIC) {
              send_fail("Bad magic");
              cdc_rx_idx = 0;
              return;
          }
          if (hdr->len > MAX_CDC_CMD) {
              send_fail("Bad cmd: too big");
              cdc_rx_idx = 0;
              return;
          }
          cdc_rx_expected = hdr->len + sizeof(struct cdc_packet_hdr);
      }
      cdc_rx_idx += count;
      if (cdc_rx_idx >= cdc_rx_expected) {
          rx_busy = 1;
          parse_cmd(cdc_rx_expected);
          if (cdc_rx_expected < cdc_rx_idx) {
              memcpy(cdc_rxbuf, cdc_rxbuf + cdc_rx_expected, cdc_rx_idx - cdc_rx_expected);
          }
          cdc_rx_idx-=cdc_rx_expected;
      }
  }
}
