#include "core/core.h"
#include "common/memory.h"
#include "core/peer_table.h"
#include "core/platform.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
create_prk (const uint8_t ek_shared[32],
            const uint8_t qk_shared[MLKEM768_BYTES],
            uint8_t prk_out[TCT_TURBOSHAKE128_STATE_LEN])
{
  tct_turboshake128_init (prk_out);
  const uint8_t SALT[] = "circe/v1";
  tct_turboshake128_absorb (prk_out, SALT, sizeof (SALT) - 1, 0x0);
  tct_turboshake128_absorb (prk_out, ek_shared, 32, 0x1);
  tct_turboshake128_absorb (prk_out, qk_shared, MLKEM768_BYTES, 0x2);
}

static void
derive_k_hs_e (const uint8_t prk[TCT_TURBOSHAKE128_STATE_LEN],
               const uint8_t th2[64], uint8_t out[32])
{
  uint8_t prk_to_expand[TCT_TURBOSHAKE128_STATE_LEN];
  circe_memory_copy (prk, prk_to_expand, TCT_TURBOSHAKE128_STATE_LEN);
  const uint8_t EXPANSION_DATA[] = "hs eurylochus";
  tct_turboshake128_absorb (prk_to_expand, EXPANSION_DATA,
                            sizeof (EXPANSION_DATA) - 1, 0x3);
  tct_turboshake128_absorb (prk_to_expand, th2, 64, 0x4);
  tct_turboshake128_squeeze_destructive (prk_to_expand, out, 32);
  circe_memory_set (prk_to_expand, 0x0, sizeof (prk_to_expand));
}

static void
derive_k_hs_p (const uint8_t prk[TCT_TURBOSHAKE128_STATE_LEN],
               const uint8_t th2[64], uint8_t out[32])
{
  uint8_t prk_to_expand[TCT_TURBOSHAKE128_STATE_LEN];
  circe_memory_copy (prk, prk_to_expand, TCT_TURBOSHAKE128_STATE_LEN);
  const uint8_t EXPANSION_DATA[] = "hs polites";
  tct_turboshake128_absorb (prk_to_expand, EXPANSION_DATA,
                            sizeof (EXPANSION_DATA) - 1, 0x3);
  tct_turboshake128_absorb (prk_to_expand, th2, 64, 0x4);
  tct_turboshake128_squeeze_destructive (prk_to_expand, out, 32);
  circe_memory_set (prk_to_expand, 0x0, sizeof (prk_to_expand));
}

static void
derive_k_d_p2e (const uint8_t prk[TCT_TURBOSHAKE128_STATE_LEN],
                const uint8_t th3[64], uint8_t out[32])
{
  uint8_t prk_to_expand[TCT_TURBOSHAKE128_STATE_LEN];
  circe_memory_copy (prk, prk_to_expand, TCT_TURBOSHAKE128_STATE_LEN);
  const uint8_t EXPANSION_DATA[] = "data p2e";
  tct_turboshake128_absorb (prk_to_expand, EXPANSION_DATA,
                            sizeof (EXPANSION_DATA) - 1, 0x3);
  tct_turboshake128_absorb (prk_to_expand, th3, 64, 0x4);
  tct_turboshake128_squeeze_destructive (prk_to_expand, out, 32);
  circe_memory_set (prk_to_expand, 0x0, sizeof (prk_to_expand));
}

static void
derive_k_d_e2p (const uint8_t prk[TCT_TURBOSHAKE128_STATE_LEN],
                const uint8_t th3[64], uint8_t out[32])
{
  uint8_t prk_to_expand[TCT_TURBOSHAKE128_STATE_LEN];
  circe_memory_copy (prk, prk_to_expand, TCT_TURBOSHAKE128_STATE_LEN);
  const uint8_t EXPANSION_DATA[] = "data e2p";
  tct_turboshake128_absorb (prk_to_expand, EXPANSION_DATA,
                            sizeof (EXPANSION_DATA) - 1, 0x3);
  tct_turboshake128_absorb (prk_to_expand, th3, 64, 0x4);
  tct_turboshake128_squeeze_destructive (prk_to_expand, out, 32);
  circe_memory_set (prk_to_expand, 0x0, sizeof (prk_to_expand));
}

/* END GENERAL HELPER FUNCTIONS */

/* BUSINESS LOGIC */

enum circe_result
dispatch_complete_packet (struct circe_context *ctx, struct circe_peer *remote)
{
  enum circe_result res = CIRCE_RESULT_SUCCESS;
  uint8_t key[32];
  uint8_t signature[MLDSA44_BYTES];
  size_t siglen;

  struct proto_circe_v1_open_tunnel ot_contents;
  uint8_t ek_shared[32];
  uint8_t qk_shared[MLKEM768_BYTES];
  uint8_t mlkem_encaps[MLKEM768_CIPHERTEXTBYTES];

  struct proto_circe_v1_ack_open ao_contents;

  struct proto_circe_v1_identify id_contents;

  struct proto_circe_v1_data dt_contents;

  struct proto_circe_v1_close_tunnel ct_contents;

  struct proto_circe_v1_path_challenge pc_contents;

  switch (proto_circe_v1_rx_type (&(remote->dgram_rx_state)))
    {
    case PROTO_CIRCE_V1_MSG_OPEN_TUNNEL:
      if (proto_circe_v1_open_tunnel_decode (
              remote->dgram_rx_buffer,
              proto_circe_v1_rx_len (&(remote->dgram_rx_state)), &ot_contents,
              remote->my_th2)
          != PBA_OK)
        {
          res = CIRCE_RESULT_INTERNAL_ERROR;
          goto dispatch_complete_packet_cleanup;
        }
      circe_platform_random_bytes_callback (remote->my_ecdh_privkey, 32);
      tct_x25519 (remote->my_ecdh_privkey, ot_contents.ecdh_fragment,
                  ek_shared);
      if (PQCP_MLKEM_NATIVE_MLKEM768_enc (mlkem_encaps, qk_shared,
                                          ot_contents.mlkem_pubkey)
          != 0)
        {
          res = CIRCE_RESULT_INTERNAL_ERROR;
          goto dispatch_complete_packet_cleanup;
        }
      create_prk (ek_shared, qk_shared, remote->prk);
      derive_k_hs_e (remote->prk, remote->my_th2, key);
      ao_contents.hdr.tunnel_id = remote->tunnel_id;
      ao_contents.ecdh_fragment = remote->my_ecdh_privkey;
      ao_contents.identity = ctx->my_identity;
      ao_contents.mlkem_encaps = mlkem_encaps;
      if (PQCP_MLDSA_NATIVE_MLDSA44_signature (
              signature, &siglen, ao_contents.identity, CIRCE_IDENTITY_LEN,
              NULL, 0, ctx->my_mldsa_privkey)
          != 0)
        {
          res = CIRCE_RESULT_INTERNAL_ERROR;
          goto dispatch_complete_packet_cleanup;
        }
      proto_circe_v1_ack_open_encode (&ao_contents, key,
                                      remote->dgram_tx_buffer,
                                      sizeof (remote->dgram_tx_buffer),
                                      &remote->dgram_tx_len, remote->my_th3);
      proto_circe_v1_tx_init (&(remote->dgram_tx_state),
                              remote->dgram_tx_buffer, remote->dgram_tx_len,
                              key, remote->last_seqn_tx);
      remote->last_seqn_tx++;
      break;
    case PROTO_CIRCE_V1_MSG_ACK_OPEN:
      derive_k_hs_e (remote->prk, remote->my_th2, key);
      if (proto_circe_v1_ack_open_decode (
              remote->dgram_rx_buffer,
              proto_circe_v1_rx_len (&(remote->dgram_rx_state)), key,
              &ao_contents, remote->my_th3)
          != PBA_OK)
        {
          res = CIRCE_RESULT_INTERNAL_ERROR;
          goto dispatch_complete_packet_cleanup;
        }

    case PROTO_CIRCE_V1_MSG_IDENTIFY:
    case PROTO_CIRCE_V1_MSG_DATA:
    case PROTO_CIRCE_V1_MSG_CLOSE_TUNNEL:
    case PROTO_CIRCE_V1_MSG_PATH_CHALLENGE:
    }
dispatch_complete_packet_cleanup:
  circe_memory_set (key, 0x0, sizeof (key));
  circe_memory_set (ek_shared, 0x0, sizeof (ek_shared));
  circe_memory_set (qk_shared, 0x0, sizeof (qk_shared));
  circe_memory_set (mlkem_encaps, 0x0, sizeof (mlkem_encaps));
  return res;
}

/* END BUSINESS LOGIC */

/* 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,
    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->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)
{
  enum circe_result res = CIRCE_RESULT_SUCCESS;

  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_buf[32];
  uint8_t *frag_key;
  switch (fragment.type)
    {
    case PROTO_CIRCE_V1_MSG_OPEN_TUNNEL:
    case PROTO_CIRCE_V1_MSG_DATA:
      // Can't fragment
      frag_key = NULL;
      break;
    case PROTO_CIRCE_V1_MSG_ACK_OPEN:
      derive_k_hs_e (remote_pt_entry->prk, remote_pt_entry->my_th2,
                     frag_key_buf);
      frag_key = frag_key_buf;
      break;
    case PROTO_CIRCE_V1_MSG_IDENTIFY:
      derive_k_hs_p (remote_pt_entry->prk, remote_pt_entry->my_th2,
                     frag_key_buf);
      frag_key = frag_key_buf;
      break;
    case PROTO_CIRCE_V1_MSG_CLOSE_TUNNEL:
    case PROTO_CIRCE_V1_MSG_PATH_CHALLENGE:
      if (remote_pt_entry->my_role == CIRCE_ROLE_EURYLOCHUS)
        {
          derive_k_d_p2e (remote_pt_entry->prk, remote_pt_entry->my_th3,
                          frag_key_buf);
        }
      else
        {
          derive_k_d_e2p (remote_pt_entry->prk, remote_pt_entry->my_th3,
                          frag_key_buf);
        }
      frag_key = frag_key_buf;
      break;
    }
  switch (proto_circe_v1_rx_push (&(remote_pt_entry->dgram_rx_state), in, len,
                                  frag_key))
    {
    case PBA_OK_COMPLETE:
      dispatch_complete_packet (
          proto_circe_v1_rx_buf (&(remote_pt_entry->dgram_rx_state)),
          remote_pt_entry);
      break;
    case PBA_OK_PARTIAL:
      break;
    default:
      res = CIRCE_RESULT_INVALID_PACKET;
      goto circe_handle_packet_cleanup;
    }
circe_handle_packet_cleanup:
  circe_memory_set (frag_key_buf, 0x0, sizeof (frag_key_buf));
  return res;
}

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