// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <device_id.h>
#include <inet6.h>
#include <netifc.h>
#include <stdio.h>
#include <string.h>
#include <xefi.h>
#include <zircon/boot/netboot.h>

#include <tftp/tftp.h>

#define TFTP_BUF_SZ 2048
char tftp_session_scratch[TFTP_BUF_SZ];
char tftp_out_scratch[TFTP_BUF_SZ];

// item being downloaded
static nbfile* item;

// TFTP file state
typedef struct {
  nbfile* netboot_file_data;
  size_t file_size;
  unsigned int progress_reported;
} file_info_t;

// TFTP transport state
typedef struct {
  struct ip6_addr_t dest_addr;
  uint16_t dest_port;
} transport_info_t;

static uint32_t last_cookie = 0;
static uint32_t last_cmd = 0;
static uint32_t last_arg = 0;
static uint32_t last_ack_cmd = 0;
static uint32_t last_ack_arg = 0;

static int nb_boot_now = 0;
static int nb_active = 0;

static char advertise_nodename[64] = "";
static char advertise_data[256] = "nodename=zircon";

static void send_query_ack(const ip6_addr* addr, uint16_t port, uint32_t cookie) {
  uint8_t buffer[256];
  nbmsg* msg = (void*)buffer;
  msg->magic = NB_MAGIC;
  msg->cookie = cookie;
  msg->cmd = NB_ACK;
  msg->arg = NB_VERSION_CURRENT;
  memcpy(msg->data, advertise_nodename, sizeof(advertise_nodename));
  udp6_send(buffer, sizeof(nbmsg) + strlen(advertise_nodename) + 1, addr, port, NB_SERVER_PORT);
}

static void advertise(void) {
  uint8_t buffer[sizeof(nbmsg) + sizeof(advertise_data)];
  nbmsg* msg = (void*)buffer;
  msg->magic = NB_MAGIC;
  msg->cookie = 0;
  msg->cmd = NB_ADVERTISE;
  msg->arg = NB_VERSION_CURRENT;
  size_t data_len = strlen(advertise_data) + 1;
  memcpy(msg->data, advertise_data, data_len);
  udp6_send(buffer, sizeof(nbmsg) + data_len, &ip6_ll_all_nodes, NB_ADVERT_PORT, NB_SERVER_PORT);
}

void netboot_recv(void* data, size_t len, const ip6_addr* saddr, uint16_t sport) {
  nbmsg* msg = data;
  nbmsg ack;
  int do_transmit = 1;

  if (len < sizeof(nbmsg))
    return;
  len -= sizeof(nbmsg);

  // printf("netboot: MSG %08x %08x %08x %08x datalen %zu\n",
  //        msg->magic, msg->cookie, msg->cmd, msg->arg, len);

  if ((last_cookie == msg->cookie) && (last_cmd == msg->cmd) && (last_arg == msg->arg)) {
    // host must have missed the ack. resend
    ack.magic = NB_MAGIC;
    ack.cookie = last_cookie;
    ack.cmd = last_ack_cmd;
    ack.arg = last_ack_arg;
    goto transmit;
  }

  ack.cmd = NB_ACK;
  ack.arg = 0;

  switch (msg->cmd) {
    case NB_COMMAND:
      if (len == 0)
        return;
      msg->data[len - 1] = 0;
      break;
    case NB_SEND_FILE:
      if (len == 0)
        return;
      msg->data[len - 1] = 0;
      for (size_t i = 0; i < (len - 1); i++) {
        if ((msg->data[i] < ' ') || (msg->data[i] > 127)) {
          msg->data[i] = '.';
        }
      }
      item = netboot_get_buffer((const char*)msg->data, msg->arg);
      if (item) {
        item->offset = 0;
        ack.arg = msg->arg;
        size_t prefix_len = strlen(NB_FILENAME_PREFIX);
        const char* filename;
        if (!strncmp((char*)msg->data, NB_FILENAME_PREFIX, prefix_len)) {
          filename = &((const char*)msg->data)[prefix_len];
        } else {
          filename = (const char*)msg->data;
        }
        printf("netboot: Receive File '%s'...\n", filename);
      } else {
        printf("netboot: Rejected File '%s'...\n", (char*)msg->data);
        ack.cmd = NB_ERROR_BAD_FILE;
      }
      break;

    case NB_DATA:
    case NB_LAST_DATA:
      if (item == 0) {
        printf("netboot: > received chunk before NB_FILE\n");
        return;
      }
      if (msg->arg != item->offset) {
        // printf("netboot: < received chunk at offset %d but current offset is %zu\n", msg->arg,
        // item->offset);
        ack.arg = item->offset;
        ack.cmd = NB_ACK;
      } else if ((item->offset + len) > item->size) {
        ack.cmd = NB_ERROR_TOO_LARGE;
        ack.arg = msg->arg;
      } else {
        memcpy(item->data + item->offset, msg->data, len);
        item->offset += len;
        ack.cmd = msg->cmd == NB_LAST_DATA ? NB_FILE_RECEIVED : NB_ACK;
        if (msg->cmd != NB_LAST_DATA) {
          do_transmit = 0;
        }
      }
      break;
    case NB_BOOT:
      nb_boot_now = 1;
      printf("netboot: Boot Kernel...\n");
      break;
    case NB_QUERY:
      // Send reply and return w/o getting the netboot state out of sync.
      send_query_ack(saddr, sport, msg->cookie);
      return;
    default:
      ack.cmd = NB_ERROR_BAD_CMD;
      ack.arg = 0;
  }

  last_cookie = msg->cookie;
  last_cmd = msg->cmd;
  last_arg = msg->arg;
  last_ack_cmd = ack.cmd;
  last_ack_arg = ack.arg;

  ack.cookie = msg->cookie;
  ack.magic = NB_MAGIC;
transmit:
  nb_active = 1;
  if (do_transmit) {
    // printf("netboot: MSG %08x %08x %08x %08x\n",
    //   ack.magic, ack.cookie, ack.cmd, ack.arg);

    udp6_send(&ack, sizeof(ack), saddr, sport, NB_SERVER_PORT);
  }
}

static tftp_status buffer_open(const char* filename, size_t size, void* cookie) {
  file_info_t* file_info = cookie;
  file_info->netboot_file_data = netboot_get_buffer(filename, size);
  if (file_info->netboot_file_data == NULL) {
    printf("netboot: unrecognized file %s - rejecting\n", filename);
    return TFTP_ERR_INVALID_ARGS;
  }
  file_info->netboot_file_data->offset = 0;
  const char* base_filename;
  size_t prefix_len = strlen(NB_FILENAME_PREFIX);
  if (!strncmp(filename, NB_FILENAME_PREFIX, prefix_len)) {
    base_filename = &filename[prefix_len];
  } else {
    base_filename = filename;
  }
  printf("Receiving %s [%lu bytes]... ", base_filename, (unsigned long)size);
  file_info->file_size = size;
  file_info->progress_reported = 0;
  return TFTP_NO_ERROR;
}

static tftp_status buffer_write(const void* data, size_t* len, off_t offset, void* cookie) {
  file_info_t* file_info = cookie;
  nbfile* nb_buf_info = file_info->netboot_file_data;
  if (offset > nb_buf_info->size || (offset + *len) > nb_buf_info->size) {
    printf("netboot: attempt to write past end of buffer\n");
    return TFTP_ERR_INVALID_ARGS;
  }
  memcpy(&nb_buf_info->data[offset], data, *len);
  nb_buf_info->offset = offset + *len;
  if (file_info->file_size >= 100) {
    unsigned int progress_pct = offset / (file_info->file_size / 100);
    if ((progress_pct > file_info->progress_reported) &&
        (progress_pct - file_info->progress_reported >= 5)) {
      printf("%u%%... ", progress_pct);
      file_info->progress_reported = progress_pct;
    }
  }
  return TFTP_NO_ERROR;
}

static void buffer_close(void* cookie) {
  file_info_t* file_info = cookie;
  file_info->netboot_file_data = NULL;
  printf("Done\n");
}

static tftp_status udp_send(void* data, size_t len, void* cookie) {
  transport_info_t* transport_info = cookie;
  int bytes_sent = udp6_send(data, len, &transport_info->dest_addr, transport_info->dest_port,
                             NB_TFTP_OUTGOING_PORT);
  return bytes_sent < 0 ? TFTP_ERR_IO : TFTP_NO_ERROR;
}

static int udp_timeout_set(uint32_t timeout_ms, void* cookie) {
  // TODO
  return 0;
}

static int strcmp8to16(const char* str8, const char16_t* str16) {
  while (*str8 != '\0' && *str8 == *str16) {
    str8++;
    str16++;
  }
  return *str8 - *str16;
}

void tftp_recv(void* data, size_t len, const ip6_addr* daddr, uint16_t dport, const ip6_addr* saddr,
               uint16_t sport) {
  static tftp_session* session = NULL;
  static file_info_t file_info = {.netboot_file_data = NULL};
  static transport_info_t transport_info = {};

  if (dport == NB_TFTP_INCOMING_PORT) {
    if (session != NULL) {
      printf("Aborting to service new connection\n");
    }
    // Start TFTP session
    int ret = tftp_init(&session, tftp_session_scratch, sizeof(tftp_session_scratch));
    if (ret != TFTP_NO_ERROR) {
      printf("netboot: failed to initiate tftp session\n");
      session = NULL;
      return;
    }

    // Override our window size on the Acer tablet
    if (!strcmp8to16("INSYDE Corp.", gSys->FirmwareVendor)) {
      uint16_t window_size = 8;
      tftp_set_options(session, NULL, NULL, &window_size);
    }

    // Initialize file interface
    tftp_file_interface file_ifc = {NULL, buffer_open, NULL, buffer_write, buffer_close};
    tftp_session_set_file_interface(session, &file_ifc);

    // Initialize transport interface
    memcpy(&transport_info.dest_addr, saddr, sizeof(struct ip6_addr_t));
    transport_info.dest_port = sport;
    tftp_transport_interface transport_ifc = {udp_send, NULL, udp_timeout_set};
    tftp_session_set_transport_interface(session, &transport_ifc);
  } else if (!session) {
    // Ignore anything sent to the outgoing port unless we've already established a connection
    return;
  }

  size_t outlen = sizeof(tftp_out_scratch);

  char err_msg[128];
  tftp_handler_opts handler_opts = {.inbuf = data,
                                    .inbuf_sz = len,
                                    .outbuf = tftp_out_scratch,
                                    .outbuf_sz = &outlen,
                                    .err_msg = err_msg,
                                    .err_msg_sz = sizeof(err_msg)};
  tftp_status status = tftp_handle_msg(session, &transport_info, &file_info, &handler_opts);
  if (status < 0) {
    printf("netboot: tftp protocol error: %s\n", err_msg);
    session = NULL;
  } else if (status == TFTP_TRANSFER_COMPLETED) {
    session = NULL;
  }
}

#define FAST_TICK 100
#define SLOW_TICK 1000

int netboot_init(const char* nodename) {
  if (netifc_open()) {
    printf("netboot: Failed to open network interface\n");
    return -1;
  }
  char buf[DEVICE_ID_MAX];
  if (!nodename || (nodename[0] == 0)) {
    device_id(eth_addr(), buf);
    nodename = buf;
  }
  if (nodename) {
    strncpy(advertise_nodename, nodename, sizeof(advertise_nodename) - 1);
    snprintf(advertise_data, sizeof(advertise_data), "version=%s;nodename=%s", BOOTLOADER_VERSION,
             nodename);
  }
  return 0;
}

const char* netboot_nodename() { return advertise_nodename; }

static int nb_fastcount = 0;
static int nb_online = 0;

int netboot_poll(void) {
  if (netifc_active()) {
    if (nb_online == 0) {
      printf("netboot: interface online\n");
      nb_online = 1;
      nb_fastcount = 20;
      netifc_set_timer(FAST_TICK);
      advertise();
    }
  } else {
    if (nb_online == 1) {
      printf("netboot: interface offline\n");
      nb_online = 0;
    }
    return 0;
  }
  if (netifc_timer_expired()) {
    if (nb_fastcount) {
      nb_fastcount--;
      netifc_set_timer(FAST_TICK);
    } else {
      netifc_set_timer(SLOW_TICK);
    }
    if (nb_active) {
      // don't advertise if we're in a transfer
      nb_active = 0;
    } else {
      advertise();
    }
  }

  netifc_poll();

  if (nb_boot_now) {
    nb_boot_now = 0;
    return 1;
  } else {
    return 0;
  }
}

void netboot_close(void) { netifc_close(); }