#include "core.h"
#include "cbor.h"
#include "chacha20_poly1305.h"
#include "memory.h"
#include "packet.h"
#include "peer_table.h"
#include <stdbool.h>
#include <stdint.h>

#define CIRCE_CORE_ASSERT(x, err)                                             \
  do                                                                          \
    {                                                                         \
      if (!(x))                                                               \
        {                                                                     \
          return err;                                                         \
        }                                                                     \
    }                                                                         \
  while (false)

static struct circe_peer peer_table[CIRCE_PEER_TABLE_LEN];

static bool
peer_table_lookup (uint8_t identity[CIRCE_IDENTITY_LEN],
                   struct circe_peer **out)
{
  uint8_t start_idx = identity[0] % CIRCE_PEER_TABLE_LEN;
  for (size_t i = 0; i < CIRCE_PEER_TABLE_LEN; ++i)
    {
      if (circe_memory_equal (
              identity,
              peer_table[(i + start_idx) % CIRCE_PEER_TABLE_LEN].identity,
              CIRCE_IDENTITY_LEN))
        {
          *out = &(peer_table[(i + start_idx) % CIRCE_PEER_TABLE_LEN]);
          return true;
        }
    }
  return false;
}

static enum circe_result
unseal_payload (struct circe_peer *peer, struct circe_cbor_decoder *cd,
                uint8_t *out, size_t *out_len)
{
  enum circe_cbor_major_type major_type;
  uint64_t field;
  uint64_t sequence_number;

  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (cd, &field),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (field == CIRCE_PACKET_FIELD_SEQ_NUMBER,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (cd, &sequence_number),
                     CIRCE_RESULT_INVALID_PACKET);

  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (cd, &field),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (field == CIRCE_PACKET_FIELD_PAYLOAD,
                     CIRCE_RESULT_INVALID_PACKET);

  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_BYTESTR,
                     CIRCE_RESULT_INVALID_PACKET);
  struct circe_cbor_buffer sealed;
  CIRCE_CORE_ASSERT (circe_cbor_extract_bytestr (cd, &sealed),
                     CIRCE_RESULT_INVALID_PACKET);
  // Make sure we got the full payload
  CIRCE_CORE_ASSERT (sealed.buf_len
                         <= cd->buf_len - (cd->cursor - cd->buf_start),
                     CIRCE_RESULT_INVALID_PACKET);
  uint8_t nonce[sizeof (uint64_t)];
  for (unsigned int i = 0; i < sizeof (uint64_t); ++i)
    {
      nonce[i] = (sequence_number >> (i * 8)) & 0xff;
    }
  if (!tct_aead_chacha20_poly1305_decrypt_and_verify (
          cd->buf_start, cd->cursor - cd->buf_start, peer->rx_key, nonce,
          sealed.buf_start, sealed.buf_len, out))
    {
      return CIRCE_RESULT_INVALID_MAC;
    }
  *out_len = sealed.buf_len - 16; // Don't include the MAC's length
  return CIRCE_RESULT_SUCCESS;
}

static bool
locations_equal (const struct circe_location a, const struct circe_location b)
{
  if (a.ip_type != b.ip_type)
    {
      return false;
    }
  else if (a.port != b.port)
    {
      return false;
    }
  else if (a.ip_type == CIRCE_ADDRESS_IPV4)
    {
      return circe_memory_equal (a.ip_addr.ipv4, b.ip_addr.ipv4,
                                 sizeof (a.ip_addr.ipv4));
    }
  else if (a.ip_type == CIRCE_ADDRESS_IPV6)
    {
      return circe_memory_equal (a.ip_addr.ipv6, b.ip_addr.ipv6,
                                 sizeof (a.ip_addr.ipv6));
    }
  // What?
  return false;
}

enum circe_result
circe_handle_packet (const struct circe_location remote, const uint8_t *in,
                     size_t len, uint64_t now_ms, uint8_t *out,
                     size_t *out_len, struct circe_event *events,
                     size_t *n_events)
{
  *n_events = 0;
  *out_len = 0;
  struct circe_cbor_decoder cd = {
    .buf_start = in,
    .buf_len = len,
    .cursor = in,
  };
  enum circe_cbor_major_type major_type;
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_MAPPING,
                     CIRCE_RESULT_INVALID_PACKET);
  uint64_t num_fields;
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &num_fields),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (num_fields >= 2, CIRCE_RESULT_INVALID_PACKET);

  // Version
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  uint64_t field;
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &field),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (field == CIRCE_PACKET_FIELD_VERSION,
                     CIRCE_RESULT_INVALID_PACKET);
  uint64_t version;
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &version),
                     CIRCE_RESULT_INVALID_PACKET);
  // If it's anything but v1 we can't handle it (there is no v2)
  CIRCE_CORE_ASSERT (field == CIRCE_VERSION_V1, CIRCE_RESULT_INVALID_PACKET);

  // Remote's identity
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &field),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (field == CIRCE_PACKET_FIELD_IDENTITY,
                     CIRCE_RESULT_INVALID_PACKET);
  struct circe_cbor_buffer remote_id_buf;
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_BYTESTR,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_bytestr (&cd, &remote_id_buf),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (remote_id_buf.buf_len == CIRCE_IDENTITY_LEN,
                     CIRCE_RESULT_INVALID_PACKET);
  struct circe_peer *remote_pt_entry;
  bool peer_in_ram
      = peer_table_lookup (remote_id_buf.buf_start, &remote_pt_entry);

  // Packet type
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &field),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (field == CIRCE_PACKET_FIELD_TYPE,
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (circe_cbor_get_major_type (&cd, &major_type),
                     CIRCE_RESULT_INVALID_PACKET);
  CIRCE_CORE_ASSERT (major_type == CIRCE_CBOR_MAJOR_UNSIGNED,
                     CIRCE_RESULT_INVALID_PACKET);
  uint64_t packet_type;
  CIRCE_CORE_ASSERT (circe_cbor_extract_unsigned (&cd, &packet_type),
                     CIRCE_RESULT_INVALID_PACKET);

  if (packet_type != CIRCE_PACKET_ROUTE_CHALLENGE
      && packet_type != CIRCE_PACKET_OPENTUNNEL && peer_in_ram)
    {
      // Verify that this came from our known remote location
      bool good_location;
      switch (remote_pt_entry->current_route)
        {
        case CIRCE_ROUTE_LAN0:
          good_location = locations_equal (remote_pt_entry->lan0_loc, remote);
          break;
        case CIRCE_ROUTE_LAN1:
          good_location = locations_equal (remote_pt_entry->lan1_loc, remote);
          break;
        case CIRCE_ROUTE_WAN:
          good_location = locations_equal (remote_pt_entry->wan_loc, remote);
          break;
        case CIRCE_ROUTE_RELAY:
          // What?
          good_location = false;
          break;
        default:
          return CIRCE_RESULT_INTERNAL_ERROR;
        }
      if (!good_location)
        {
          (*n_events)++;
          events[0] = (struct circe_event){
            .event_type = CIRCE_EVENT_REJECTED_REMOTE,
          };
          return CIRCE_RESULT_SUCCESS;
        }
    }

  switch (packet_type)
    {
    case CIRCE_PACKET_OPENTUNNEL:
    case CIRCE_PACKET_ACKOPEN:
    case CIRCE_PACKET_IDENTIFY:
    case CIRCE_PACKET_DATA:
      (*n_events)++;
      events[0] = (struct circe_event){
        .event_type = CIRCE_EVENT_DATA_RECEIVED,
      };
      return unseal_payload (remote_pt_entry, &cd, out, out_len);
    case CIRCE_PACKET_COMMAND:
    }
  return CIRCE_RESULT_INVALID_PACKET;
}