#define LOCAL_DEBUG

#include <stdlib.h>
#include "lib/tunnel_encaps.h"

struct format_fn {
  node n;
  int type;
  const char *name;
  format_tunnel_encap_fn cb;
};

list format_fns;

static format_tunnel_encap_fn lookup_format_tunnel_encap_cb(int type);
static const char *lookup_format_tunnel_encap_name(int type);
static int default_format_tunnel_encap(const struct te_encap *encap, byte *buf, uint size);

static
int walk_encap(const void *p, size_t sub_tlv_len, const struct te_visitor *visitor, struct te_context *ctx)
{
  struct te_encap encap;
  memset(&encap, 0, sizeof(encap));

  encap.type = 0;
  encap.length = sub_tlv_len;
  encap.data = (void *)p;
  return visitor->visit_encap ? visitor->visit_encap(&encap, ctx) : 0;
}

static
int walk_color(const void *p, size_t sub_tlv_len, const struct te_visitor *visitor, struct te_context *ctx)
{
  if (sub_tlv_len != 8)
  {
    DBG(L_TRACE "WG: color len error %d", sub_tlv_len);
    return -1;
  }

  if (get_u16(p) != 0x030b)
  {
    DBG(L_TRACE "WG: color error %04x", get_u16(p));
    return -1;
  }

  return visitor->visit_color ? visitor->visit_color(get_u32(p+4), ctx) : 0;
}

static
int walk_udp_dest_port(const void *p, size_t sub_tlv_len, const struct te_visitor *visitor, struct te_context *ctx)
{
  if (sub_tlv_len != 2)
  {
    DBG(L_TRACE "WG: udp dest port len error %d", sub_tlv_len);
    return -1;
  }

  return visitor->visit_udp_dest_port ? visitor->visit_udp_dest_port(get_u16(p), ctx) : 0;
}

static
int walk_tunnel_ep(const void *p, size_t sub_tlv_len, const struct te_visitor *visitor, struct te_context *ctx)
{
  if (sub_tlv_len < 6)
  {
    DBG(L_TRACE "WG: tunnel ep len error");
    return -1;
  }

  struct te_endpoint ep;
  memset(&ep, 0, sizeof(ep));

  ep.reserved = get_u32(p);
  u16 af = get_u16(p + 4);
  ep.af = af;
  switch (af)
  {
  case 0:
    if (sub_tlv_len != 6)
    {
      DBG(L_TRACE "WG: Fam 0 len error %d", sub_tlv_len);
      return -1;
    }
    ep.ip = IP6_NONE;
    break;
  case AFI_IPV4:
    if (sub_tlv_len != 10)
    {
      DBG(L_TRACE "WG: IPv4 len error %d", sub_tlv_len);
      return -1;
    }
    ep.ip = ipa_from_ip4(get_ip4(p + 6));
    break;
  case AFI_IPV6:
    if (sub_tlv_len != 22)
    {
      DBG(L_TRACE "WG: IPv6 len error %d", sub_tlv_len);
      return -1;
    }
    ep.ip = ipa_from_ip6(get_ip6(p + 6));
    break;
  default:
    DBG(L_TRACE "WG: Address family error %d", af);
    return -1;
  }

  return visitor->visit_ep ? visitor->visit_ep(&ep, ctx) : 0;
}

static
int walk_unknown(const void *p, int type, size_t sub_tlv_len, const struct te_visitor *visitor, struct te_context *ctx)
{
  struct te_unknown unknown;
  memset(&unknown, 0, sizeof(unknown));

  unknown.length = sub_tlv_len;
  unknown.data = (void *)p;
  return visitor->visit_unknown ? visitor->visit_unknown(type, &unknown, ctx) : 0;
}

static
int walk_sub_tlv(const u8 *p, size_t len, const struct te_visitor *visitor, struct te_context *ctx)
{
  if (len < 2)
  {
    DBG(L_TRACE "WG: sub_tlv len error %d", len);
    return -1;
  }

  const u8 *first = p;
  const u8 *last = p + len;
  int type = get_u8(p++);
  u16 sub_tlv_len = 0;

  if (type >= 0 && type <= 127)
  {
    sub_tlv_len = get_u8(p);
    p++;
  }
  else if (type >= 128 && type <= 255)
  {
    if (len < 3)
    {
      DBG(L_TRACE "WG: sub_tlv len error %d", len);
      return -1;
    }

    sub_tlv_len = get_u16(p);
    p += 2;
  }
  else
  {
    DBG(L_TRACE "WG: sub_tlv type error %d", type);
    return -1;
  }

  if (p + sub_tlv_len > last)
  {
    DBG(L_TRACE "WG: sub_tlv value len error %d", sub_tlv_len);
    return -1;
  }

  if (visitor->visit_subtlv && visitor->visit_subtlv(type, ctx) < 0)
    return p - first + sub_tlv_len;

  int res = 0;

  switch (type)
  {
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP:
    res = walk_encap(p, sub_tlv_len, visitor, ctx);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_TUNNEL_EP:
    res = walk_tunnel_ep(p, sub_tlv_len, visitor, ctx);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR:
    res = walk_color(p, sub_tlv_len, visitor, ctx);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT:
    res = walk_udp_dest_port(p, sub_tlv_len, visitor, ctx);
    break;
  default:
    res = walk_unknown(p, type, sub_tlv_len, visitor, ctx);
    break;
  }

  if (res < 0)
  {
    DBG(L_TRACE "WG: Decode error");
    return res;
  }

  /* TODO use return value? */
  if (visitor->visit_subtlv_end)
    visitor->visit_subtlv_end(type, ctx);

  return p - first + sub_tlv_len;
}

static
int walk_sub_tlvs(const u8 *p, size_t size, const struct te_visitor *visitor, struct te_context *ctx)
{
  const u8 *pos = p;

  while (size > 0)
  {
    int l = walk_sub_tlv(pos, size, visitor, ctx);

    if (l < 0)
    {
      DBG(L_TRACE "WG: decode error %d", l);
      return l;
    }

    ADVANCE(pos, size, l);
  }

  return pos - p;
}

static
int walk_tlv(const u8 *p, size_t len, const struct te_visitor *visitor, struct te_context *ctx)
{
  const u8 *pos = p;
  int l;

  if (len < 4)
  {
    DBG(L_TRACE "WG: tunnel_encap len error %d", len);
    return -1;
  }

  u16 type = get_u16(pos);
  ADVANCE(pos, len, 2);

  u16 value_length = get_u16(pos);
  ADVANCE(pos, len, 2);

  if (len < value_length)
  {
    DBG(L_ERR "WG: tunnel encap len error %d", value_length);
    return -1;
  }

  ctx->sub_tlvs_pos = pos;
  ctx->sub_tlvs_len = value_length;

  if (visitor->visit_tlv && visitor->visit_tlv(type, ctx) < 0)
    return 0;

  l = walk_sub_tlvs(pos, value_length,  visitor, ctx);
  if (l < 0)
  {
    DBG(L_ERR "WG: sub-TLVs decode error");
    return l;
  }

  if (l != value_length)
  {
    DBG(L_ERR "WG: TLV length error");
    return -1;
  }

  ADVANCE(pos, len, l);

  if (visitor->visit_tlv_end && visitor->visit_tlv_end(ctx) < 0)
    return 0;

  ctx->sub_tlvs_pos = NULL;
  ctx->sub_tlvs_len = 0;

  return pos - p;
}

int walk_tunnel_encap(const struct adata *a, const struct te_visitor *visitor, void *opaque)
{
  const u8 *p = a->data;
  const u8 *pos = p;
  uint size = a->length;
  struct te_context ctx = { .opaque = opaque, };

  while (size > 0)
  {
    int l = walk_tlv(p, size, visitor, &ctx);

    if (l < 0)
    {
      DBG(L_ERR "WG: TLV decode error");
      return l;
    }

    ADVANCE(pos, size, l);
  }

  return 0;
}

struct count_subtlvs {
  uint count;
};

static int
visitor_count_subtlv(int UNUSED type, struct te_context *ctx){
  struct count_subtlvs *data = ctx->opaque;

  data->count++;
  return 0;
}

int te_count_subtlvs(struct te_context *ctx)
{
  struct count_subtlvs data = { .count = 0, };
  struct te_context new_ctx = { .opaque = &data, };
  struct te_visitor visitor = { .visit_subtlv = visitor_count_subtlv, };

  if (walk_sub_tlvs(ctx->sub_tlvs_pos, ctx->sub_tlvs_len, &visitor, &new_ctx) < 0)
    return -1;

  return data.count;
}

struct {
  int type;
  const char *name;
} tunnel_types[] = {
  {0, "Reserved"},
  {1, "L2TPv3 over IP"},
  {2, "GRE"},
  {3, "Transmit tunnel endpoint (DEPRECATED)"},
  {4, "IPsec in Tunnel-mode (DEPRECATED)"},
  {5, "IP in IP tunnel with IPsec Transport Mode (DEPRECATED)"},
  {6, "MPLS-in-IP tunnel with IPsec Transport Mode (DEPRECATED)"},
  {7, "IP in IP"},
  {8, "VXLAN Encapsulation"},
  {9, "NVGRE Encapsulation"},
  {10, "MPLS Encapsulation"},
  {11, "MPLS in GRE Encapsulation"},
  {12, "VXLAN GPE Encapsulation"},
  {13, "MPLS in UDP Encapsulation"},
  {14, "IPv6 Tunnel"},
  {15, "SR TE Policy Type"},
  {16, "Bare"},
  {17, "SR Tunnel (DEPRECATED)"},
  {18, "Cloud Security"},
  {19, "Geneve Encapsulation"},
  {20, "Any-Encapsulation"},
  {21, "GTP Tunnel Type"},
  {22, "Dynamic Path Selection (DPS) Tunnel Encapsulation"},
  {23, "Originating PE (OPE)"},
  {24, "Dynamic Path Selection (DPS) Policy"},
};

const uint num_tunnel_types = sizeof(tunnel_types)/sizeof(tunnel_types[0]);

static const char *
lookup_name(int type)
{
  for (uint i = 0; i < num_tunnel_types; i++)
  {
    if (tunnel_types[i].type == type)
      return tunnel_types[i].name;
  }

  const char *name = lookup_format_tunnel_encap_name(type);
  return name ? name : "Unassigned";
}

static int visitor_format_tlv(int type, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  const char *name = lookup_name(type);

  int l = bsnprintf(buf->pos, buf->size, "%s{type: %s(%d)", buf->tlv?"}, ":"", name, type);
  if (l < 0)
  {
    DBG(L_ERR "WG: bsnprintf error");
    return -1;
  }

  buf->tlv++;
  buf->tlv_type = type;
  buf->subtlv=0;
  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_subtlv(int UNUSED type, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;

  int l = bsnprintf(buf->pos, buf->size, ", ");
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_encap(const struct te_encap *encap, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  int l = bsnprintf(buf->pos, buf->size, "encap: ");
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);

  l = lookup_format_tunnel_encap_cb(buf->tlv_type)(encap, buf->pos, buf->size);
  if (l < 0)
    l = default_format_tunnel_encap(encap, buf->pos, buf->size);

  if (l < 0)
  {
    DBG(L_ERR "WG: Encapsulation format error");
    return l;
  }

  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_ep(const struct te_endpoint *ep, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  int l =  bsnprintf(buf->pos, buf->size, "ep: %I", ep->ip);
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_color(u32 color, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  int l = bsnprintf(buf->pos, buf->size, "color: %d", color);
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_udp_dest_port(u16 udp_dest_port, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  int l = bsnprintf(buf->pos, buf->size, "udp: %d", udp_dest_port);
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);
  return 0;
}

static int visitor_format_unknown(int type, const struct te_unknown *unknown, struct te_context *ctx)
{
  struct te_buffer *buf = ctx->opaque;
  const byte *data = unknown->data;
  int l = bsnprintf(buf->pos, buf->size, "unknown(%d): ", type);
  if (l < 0)
    return -1;
  ADVANCE(buf->pos, buf->size, l);

  int first = 1;
  for (uint i = 0; i < unknown->length; i++)
  {
    if (buf->size < 4) {
      DBG(L_ERR "WG: Format unknown sub-TLV error");
      return -1;
    }

    int l = bsnprintf(buf->pos, buf->size, "%s%02x", first?"":" ", data[i]);
    ADVANCE(buf->pos, buf->size, l);
    first = 0;
  }

  return 0;
}

struct te_visitor format_visitor =
{
  .visit_tlv = visitor_format_tlv,
  .visit_subtlv = visitor_format_subtlv,
  .visit_encap = visitor_format_encap,
  .visit_ep = visitor_format_ep,
  .visit_color = visitor_format_color,
  .visit_udp_dest_port = visitor_format_udp_dest_port,
  .visit_unknown = visitor_format_unknown,
};

const struct te_visitor *te_get_format_visitor(void)
{
  return &format_visitor;
}

int format_tunnel_encap(const struct adata *a, byte *buf, uint size)
{
  struct te_buffer buffer = {.pos=buf, .size=size, .tlv=0, .subtlv=0};
  int l = bsnprintf(buffer.pos, buffer.size, "[");
  if (l < 0)
    return -1;
  ADVANCE(buffer.pos, buffer.size, l);

  walk_tunnel_encap(a, &format_visitor, &buffer);

  l = bsnprintf(buffer.pos, buffer.size, "%s]", buffer.tlv?"}":"");
  if (l < 0)
    return -1;
  ADVANCE(buffer.pos, buffer.size, l);
  return buffer.pos - buf;
}

static int default_format_tunnel_encap(const struct te_encap *encap, byte *buf, uint size)
{
  byte *pos = buf;
  const byte *data = encap->data;
  int first = 1;

  for (uint i = 0; i < encap->length; i++)
  {
    if (size < 4) {
      DBG(L_ERR "WG: Default format tunnel encap error");
      return pos - buf;
    }

    int l = bsnprintf(pos, size, "%s%02x", first?"":" ", data[i]);
    ADVANCE(pos, size, l);

    first = 0;
  }

  return pos - buf;
}

static struct format_fn *lookup_format_tunnel_encap(int type)
{
  struct format_fn *n;
  WALK_LIST(n, format_fns) {
    if (n->type == type)
      return n;
  }

  return NULL;
}

static const char *lookup_format_tunnel_encap_name(int type)
{
  struct format_fn *fn = lookup_format_tunnel_encap(type);

  if (fn && fn->name)
    return fn->name;

  return NULL;
}

static format_tunnel_encap_fn lookup_format_tunnel_encap_cb(int type)
{
  struct format_fn *fn = lookup_format_tunnel_encap(type);;

  if (fn && fn->cb)
    return fn->cb;

  DBG("lookup_format_tunnel_encap not found: %d\n", type);
  return default_format_tunnel_encap;
}

struct tlv_buffer {
  uint len;
  uint size;
  byte *p;
  byte *tlv_len_p;
  byte *subtlv_len_p;
};

static int calc_tlv(int type, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
      Header:
      Tunnel Type (2 Octets)
      Length (2 Octets)
      Value (Variable)
  */
  buf->len += 4;

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  if (type <= 0 || type > 65535)
    return -1;

  put_u16(buf->p, type);
  buf->p += 2;

  buf->tlv_len_p = buf->p;

  put_u16(buf->p, 0xcccc);		/* Update in calc_tlv_end */
  buf->p += 2;
  return 0;
}

static int calc_tlv_end(struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;

  if (!buf->size)
    return 0;

  if (!buf->tlv_len_p)
    return -1;

  put_u16(buf->tlv_len_p, buf->p - buf->tlv_len_p - 2);
  buf->tlv_len_p = NULL;
  return 0;
}

static int calc_subtlv(int type, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Type (1 Octet)
    Sub-TLV Length (1 or 2 Octets)
  */
  if (type >= 0 && type <= 127)
    buf->len += 2;
  else if (type >= 128 && type <= 255)
    buf->len += 3;
  else
  {
    DBG(L_ERR "calc_subtlv bad type %d\n", type );
    return -1;
  }

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  put_u8(buf->p, type);
  buf->p++;

  buf->subtlv_len_p = buf->p;

  if (type >= 0 && type <= 127)
  {
    put_u8(buf->p, 0xdd);	/* Update in calc_subtlv_end */
    buf->p++;
  }
  else
  {
    put_u16(buf->p, 0xeeee);	/* Update in calc_subtlv_end */
    buf->p += 2;
  }

  return 0;
}

static int calc_subtlv_end(int type, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;

  if (!buf->size)
    return 0;

  if (!buf->subtlv_len_p)
    return -1;

  if (type >= 0 && type <= 127)
  {
    put_u8(buf->subtlv_len_p, buf->p - buf->subtlv_len_p - 1);
    buf->subtlv_len_p = NULL;
  }
  else if (type >= 128 && type <= 255)
  {
    put_u16(buf->subtlv_len_p, buf->p - buf->subtlv_len_p - 2);
    buf->subtlv_len_p = NULL;
  }
  else
    return -1;

  return 0;
}

static int calc_encap(const struct te_encap *encap, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Value (Variable)
   */
  if (encap->type <= 0)
  {
    DBG(L_ERR "calc_encap bad type %d\n", encap->type );
    return -1;
  }

  buf->len += encap->length;

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  memcpy(buf->p, encap->data, encap->length);
  buf->p += encap->length;
  return 0;
}

static int calc_ep(const struct te_endpoint *ep, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Value (10 or 22 Octets)
  */
  if (ep->af == AFI_IPV4 && ipa_is_ip4(ep->ip))
    buf->len += 10;
  else if (ep->af == AFI_IPV6 && ipa_is_ip6(ep->ip))
    buf->len += 22;
  else if (ep->af == 0)
    buf->len += 6;
  else
  {
    DBG(L_ERR "calc_encap bad af %04x\n", ep->af );
    return -1;
  }

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  put_u32(buf->p, ep->reserved);
  buf->p += 4;

  put_u16(buf->p, ep->af);
  buf->p += 2;

  switch (ep->af)
  {
  case AFI_IPV4:
    put_ip4(buf->p, ipa_to_ip4(ep->ip));
    buf->p += 4;
    break;
  case AFI_IPV6:
    put_ip6(buf->p, ipa_to_ip6(ep->ip));
    buf->p += 16;
    break;
  default:
    break;
  }

  return 0;
}

static int calc_color(u32 color, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Value (8 Octets)
  */
  buf->len += 8;

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  put_u32(buf->p, 0x030b0000);
  buf->p += 4;

  put_u32(buf->p, color);
  buf->p += 4;
  return 0;
}

static int calc_udp_dest_port(u16 UNUSED udp_dest_port, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Value (2 Octets)
  */
  buf->len += 2;

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  put_u16(buf->p, udp_dest_port);
  buf->p += 2;
  return 0;
}

static int calc_unknown(int UNUSED type, const struct te_unknown *unknown, struct te_context *ctx)
{
  struct tlv_buffer *buf = ctx->opaque;
  /*
    Sub-TLV Value (Variable)
   */

  buf->len += unknown->length;

  if (!buf->size)
    return 0;

  if (buf->len > buf->size)
    return -1;

  memcpy(buf->p, unknown->data, unknown->length);
  buf->p += unknown->length;
  return 0;
}

struct te_visitor encode_visitor =
{
  .visit_tlv = calc_tlv,
  .visit_tlv_end = calc_tlv_end,
  .visit_subtlv = calc_subtlv,
  .visit_subtlv_end = calc_subtlv_end,
  .visit_encap = calc_encap,
  .visit_ep = calc_ep,
  .visit_color = calc_color,
  .visit_udp_dest_port = calc_udp_dest_port,
  .visit_unknown = calc_unknown,
};

const struct te_visitor *te_get_encode_visitor(void)
{
  return &encode_visitor;
}

int te_init_encode_visitor(struct tlv_buffer *buf, byte *p, uint size)
{
  if (!buf)
    return sizeof(struct te_buffer);

  memset(buf, 0, sizeof(struct te_buffer));
  buf->size = size;
  buf->p = p;

  return 0;
}

uint te_get_encode_length(struct tlv_buffer *buf)
{
  return buf->len;
}

void register_format_tunnel_encap(int type, const char *name, format_tunnel_encap_fn cb)
{
  struct format_fn *fn = malloc(sizeof(struct format_fn));
  memset(fn, 0, sizeof(*fn));
  fn->type = type;
  fn->name = name;
  fn->cb = cb;
  add_head(&format_fns, &fn->n);
}

void unregister_format_tunnel_encap(int type, format_tunnel_encap_fn cb)
{
  struct format_fn *n, *nxt;
  WALK_LIST_DELSAFE(n, nxt, format_fns)
  {
    if (n->type == type && n->cb == cb)
    {
      rem_node(&n->n);
    }
  }
}

void tunnel_encap_init()
{
  init_list(&format_fns);
}