#include "core/core.h"
#include "common/memory.h"
#include "core/peer_table.h"
#include "mldsa/mldsa_native.h"
#include "mlkem/mlkem_native.h"
#include "proto_circe.h"
#include "tinycrypt/kangarootwelve128.h"
#include "tinycrypt/x25519.h"
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>

#define CIRCE_P1305_MAC_LEN 16

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

#define CIRCE_CORE_ASSERT_WITH_CLEANUP(x, err, errvar, label)                 \
  do                                                                          \
    {                                                                         \
      if (!(x))                                                               \
        {                                                                     \
          (errvar) = (err);                                                   \
          goto label;                                                         \
        }                                                                     \
    }                                                                         \
  while (false)

/* GENERAL HELPER FUNCTIONS */

#if defined(__GNUC__) || defined(__clang__)
#define CIRCE_FIRST_FREE_OUTGOING_PACKET_ENTRY(ctx)                           \
  __builtin_ffsll ((ctx)->outgoing_packet_table_free_mask)
#define CIRCE_FIRST_TAKEN_OUTGOING_PACKET_ENTRY(ctx)                          \
  __builtin_ffsll (~((ctx)->outgoing_packet_table_free_mask))
#define CIRCE_FIRST_FREE_PEER_TABLE_ENTRY(ctx)                                \
  __builtin_ffsll ((ctx)->peer_table_free_mask)
#else
#error Circe can currently only be built with GCC or Clang.
#endif

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;
}

static enum circe_result
peer_table_lookup (struct circe_context *ctx,
                   const uint8_t tunnel_id[CIRCE_TUNNEL_ID_LEN],
                   struct circe_peer **const out)
{
  size_t start_idx = tunnel_id[0] % CIRCE_PEER_TABLE_LEN;
  for (size_t i = 0; i < CIRCE_PEER_TABLE_LEN; ++i)
    {
      if (circe_memory_equal (
              tunnel_id,
              ctx->peer_table[(i + start_idx) % CIRCE_PEER_TABLE_LEN]
                  .tunnel_id,
              CIRCE_TUNNEL_ID_LEN))
        {
          *out = &(ctx->peer_table[(i + start_idx) % CIRCE_PEER_TABLE_LEN]);
          return CIRCE_RESULT_SUCCESS;
        }
    }
  size_t out_idx = CIRCE_FIRST_FREE_PEER_TABLE_ENTRY (ctx);
  if (out_idx == 0)
    {
      return CIRCE_RESULT_PEER_TABLE_FULL;
    }
  if (!ctx->recall_peer_callback (ctx->ctx_handle, tunnel_id,
                                  &ctx->peer_table[out_idx - 1]))
    {
      return CIRCE_RESULT_UNKNOWN_REMOTE;
    }
  ctx->peer_table_free_mask &= ~(1ull << (out_idx - 1));
  ctx->peer_handshake_pending_mask |= (1ull << (out_idx - 1));
  *out = &(ctx->peer_table[out_idx - 1]);
  (*out)->last_handshake_stage = CIRCE_HANDSHAKE_STAGE_OPEN;
  proto_circe_v1_rx_init (&(*out)->dgram_rx_state, (*out)->dgram_rx_buffer,
                          sizeof ((*out)->dgram_rx_buffer));
  return CIRCE_RESULT_HANDSHAKE_PENDING;
}

static void
close_without_saying_goodbye (struct circe_context *ctx,
                              struct circe_peer *remote_pt_entry)
{
  size_t start_idx = remote_pt_entry->identity[0] % CIRCE_PEER_TABLE_LEN;
  size_t peer_idx = CIRCE_PEER_TABLE_LEN;
  for (size_t i = 0; i < CIRCE_PEER_TABLE_LEN; ++i)
    {
      if (circe_memory_equal (
              remote_pt_entry->identity,
              ctx->peer_table[(i + start_idx) % CIRCE_PEER_TABLE_LEN].identity,
              CIRCE_IDENTITY_LEN))
        {
          peer_idx = start_idx;
          break;
        }
    }
  if (peer_idx == CIRCE_PEER_TABLE_LEN)
    {
      // Double-free, perhaps; just ignore it
      return;
    }
  ctx->peer_handshake_pending_mask &= ~(1ull << peer_idx);
  ctx->peer_table_free_mask |= 1ull << peer_idx;
  circe_memory_set ((uint8_t *)&(ctx->peer_table[peer_idx]), 0x0,
                    sizeof (struct circe_peer));
}

/* END GENERAL HELPER FUNCTIONS */

/* PUBLIC API */

enum circe_result
circe_initialize_core (
    struct circe_context *ctx, uint8_t identity[CIRCE_IDENTITY_LEN],
    uint8_t mldsa_privkey[MLDSA_SECRETKEYBYTES (44)], uint64_t now_ms,
    void *ctx_handle,
    void (*random_bytes_callback) (void *, uint8_t *, uint32_t),
    bool (*recall_peer_callback) (void *, const uint8_t[CIRCE_IDENTITY_LEN],
                                  struct circe_peer *))
{
  ctx->last_tick_ms = now_ms;
  circe_memory_copy (identity, ctx->my_identity, CIRCE_IDENTITY_LEN);
  circe_memory_copy (mldsa_privkey, ctx->my_mldsa_privkey,
                     MLDSA_SECRETKEYBYTES (44));
  ctx->peer_handshake_pending_mask = 0x0ull;
  ctx->peer_table_free_mask = ~0x0ull;
  ctx->outgoing_packet_table_free_mask = ~0x0ull;
  circe_memory_set ((uint8_t *)ctx->outgoing_packets, 0x0,
                    sizeof (ctx->outgoing_packets));
  circe_memory_set ((uint8_t *)ctx->peer_table, 0x0, sizeof (ctx->peer_table));
  ctx->ctx_handle = ctx_handle;
  ctx->random_bytes_callback = random_bytes_callback;
  ctx->recall_peer_callback = recall_peer_callback;
  return CIRCE_RESULT_SUCCESS;
}

enum circe_result
circe_handle_packet (struct circe_context *ctx,
                     const struct circe_location remote,
                     const uint8_t *restrict in, size_t len,
                     uint8_t *restrict out, size_t *out_len, uint64_t now_ms,
                     struct circe_event *events, size_t *n_events)
{
  proto_circe_v1_frag_t fragment;
  if (proto_circe_v1_peek (in, len, &fragment) != PBA_OK)
    {
      return CIRCE_RESULT_INVALID_PACKET;
    }
  struct circe_peer *remote_pt_entry;
  peer_table_lookup (ctx, fragment.hdr.tunnel_id, &remote_pt_entry);
  uint8_t *frag_key;
  return CIRCE_RESULT_SUCCESS;
}

enum circe_result
circe_send_outgoing_packet (struct circe_context *ctx,
                            const uint8_t remote_identity[CIRCE_IDENTITY_LEN],
                            const uint8_t *in, size_t len, uint64_t now_ms,
                            struct circe_event *events, size_t *n_events);

enum circe_deadline circe_next_deadline (struct circe_context *ctx,
                                         uint64_t *ms_timestamp);

enum circe_result circe_tick (struct circe_context *ctx, uint64_t now_ms,
                              struct circe_outgoing_packet *out,
                              bool *new_packet_out);