// Based on proto/rip/rip.c

#include <stdio.h>
#include <stdlib.h>
#include "nest/protocol.h"
#include "nest/iface.h"
#include "sysdep/linux/wireguard.h"
#include "wireguard.h"
#include "proto/bgp/bgp.h"

static
int get_device(struct wg_proto *p, wg_device **pdev, const char *device_name)
{
  struct wg_config *c = (struct wg_config *) p->p.cf;

  /* if (has_user_space(p)) */
    /* return user_get_device(p, dev, device_name); */
  /* else */
  /*   return wg_get_device(dev, device_name); */

  wg_device *dev = calloc(1, sizeof(wg_device));
  strncpy(dev->name, device_name, sizeof(dev->name));
  if (c->private_key)
    {
      dev->flags |= WGDEVICE_HAS_PRIVATE_KEY | WGDEVICE_HAS_PUBLIC_KEY;
      wg_key_from_base64(dev->private_key, c->private_key);
      wg_generate_public_key(dev->public_key, dev->private_key);
    }
  if (dev->listen_port)
    dev->flags |= WGDEVICE_HAS_LISTEN_PORT;
  dev->listen_port = c->listen_port;
  debug("listen port %d\n", c->listen_port);

  struct peer_config *pc = NULL;
  WALK_LIST(pc,c->peers)
    {
      wg_peer *peer = calloc(1, sizeof(wg_peer));
      if (!dev->first_peer)
	dev->first_peer = peer;
      if (dev->last_peer)
	dev->last_peer->next_peer = peer;
      dev->last_peer = peer;

      if (pc->public_key)
	{
	  peer->flags = WGPEER_HAS_PUBLIC_KEY;
	  wg_key_from_base64(peer->public_key, pc->public_key);
	}
      peer->next_peer = NULL;

      sockaddr_fill((sockaddr*)&peer->endpoint.addr,
		    ipa_is_ip4(pc->endpoint) ? AF_INET : AF_INET6,
		    pc->endpoint, NULL, pc->remote_port);

      wg_allowedip *allowedip = calloc(1, sizeof(wg_allowedip));
      peer->first_allowedip = allowedip;
      peer->last_allowedip = allowedip;

      switch (pc->allowedip.type)
	{
	case NET_IP4:
	  allowedip->family = AF_INET;
	  allowedip->ip4 = ipa_to_in4(net_prefix(&pc->allowedip));
	  allowedip->cidr = net_pxlen(&pc->allowedip);
	  break;
	case NET_IP6:
	  allowedip->family = AF_INET6;
	  allowedip->ip6 = ipa_to_in6(net_prefix(&pc->allowedip));
	  allowedip->cidr = net_pxlen(&pc->allowedip);
	  break;
	default:
	  break;
	}

      allowedip->next_allowedip = NULL;
    }

  *pdev = dev;
  return 0;
}

static void
wg_init_entry(void *e_)
{
  struct wg_entry *e UNUSED = e_;
//  debug("wg_init_entry\n");
}

static void
wg_postconfig(struct proto_config *C UNUSED)
{
  debug("postconfig\n");
}

static void
wg_if_notify(struct proto *P, unsigned flags, struct iface *i)
{
  struct wg_proto *p = (struct wg_proto *) P;
  struct wg_config *c = (struct wg_config *) P->cf;

  debug("WG: if_notify %p %s %d %s\n", i, i->name, flags, c->ifname);
  if (c->ifname && !strcmp(i->name, c->ifname)) {
    debug("WG: found ifname\n");
    p->iface = i;
  }
}

static void
dump(void *ptr, size_t len)
{
  unsigned char *data = ptr;

  for (size_t i=0; i<len; i++) {
    fprintf(stderr, "%02x ", data[i]);
  }
  fprintf(stderr, "\n");
}

#define BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP         1
#define BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR         4
#define BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP     6
#define BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT 8

#define FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP         (1<<BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP)
#define FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR         (1<<BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR)
#define FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP     (1<<BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP)
#define FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT (1<<BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT)

#define BGP_TUNNEL_ENCAP_A_TUNNEL_TYPE_WIREGUARD 65535

static
int decode_wireguard(const void *p, size_t sub_tlv_len, wg_key *pubkey, u16 *flags)
{
  if (sub_tlv_len != sizeof(wg_key)) {
    log(L_TRACE "WG: wireguard len error %d", sub_tlv_len);
    return -1;
  }

  memcpy(pubkey, p, sizeof(wg_key));
  *flags |= FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP;
  return 0;
}

static
int decode_color(const void *p, size_t sub_tlv_len, u32 *color, u16 *flags)
{
  if (sub_tlv_len != 8) {
    log(L_TRACE "WG: color len error %d", sub_tlv_len);
    return -1;
  }

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

  *color = get_u32(p+2);
  *flags |= FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR;
  return 0;
}

static
int decode_udp_dest_port(const void *p, size_t sub_tlv_len, u16 *udp_dest_port, u16 *flags)
{
  if (sub_tlv_len != 2) {
    log(L_TRACE "WG: udp dest port len error %d", sub_tlv_len);
    return -1;
  }

  *udp_dest_port = get_u16(p);
  *flags |= FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT;
  return 0;
}

static
int decode_remote_ep(void *p, size_t sub_tlv_len, u32 *as4, ip_addr *remote_ep, u16 *flags)
{
  if (sub_tlv_len < 6) {
    log(L_TRACE "WG: remote ep len error");
    return -1;
  }

  *as4 = get_u32(p);
  u16 af = get_u16(p + 4);
  switch (af) {
  case NET_IP4:
    if (sub_tlv_len != 10) {
      log(L_TRACE "WG: IPv4 len error %d", sub_tlv_len);
      return -1;
    }
    *remote_ep = ipa_from_ip4(get_ip4(p + 6));
    *flags |= FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP;
    return 0;
  case NET_IP6:
    if (sub_tlv_len != 22) {
      log(L_TRACE "WG: IPv6 len error %d", sub_tlv_len);
      return -1;
    }
    *remote_ep = ipa_from_ip6(get_ip6(p + 6));
    *flags |= FLAG_BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP;
    return 0;
  default:
    log(L_TRACE "WG: Address family error %d", af);
    return -1;
  }
}

static
int decode_sub_tlv(u8 *p, size_t len, wg_key *pubkey,
		   u32 *remote_ep_as, ip_addr *remote_ep_addr,
		   u32 *color, u16 *udp_dest_port, u16 *flags)
{
  if (len < 3) {
    log(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;

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

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

  int res = 0;

  switch (type) {
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_ENCAP:
    res = decode_wireguard(p, sub_tlv_len, pubkey, flags);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_REMOTE_EP:
    res = decode_remote_ep(p, sub_tlv_len, remote_ep_as, remote_ep_addr, flags);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_COLOR:
    res = decode_color(p, sub_tlv_len, color, flags);
    break;
  case BGP_TUNNEL_ENCAP_A_SUB_TLV_UDP_DEST_PORT:
    res = decode_udp_dest_port(p, sub_tlv_len, udp_dest_port, flags);
    break;
  default:
    /* Skip unsupported sub-TLV. */
    res = 0;
    break;
  }

  if (res < 0)
    return res;

  return p - first + sub_tlv_len;
}

static
int decode_tunnel_encap(const eattr *e, wg_key *pubkey, u32 *as4, ip_addr *remote_ep, u32 *color, u16 *udp_port, u16 *flags)
{
  u8 *p = e->u.ptr->data;
  int len = e->u.ptr->length;

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

  u16 tunnel_type = get_u16(p);

  log(L_DEBUG "WG: tunnel type %d", tunnel_type);

  if (tunnel_type != BGP_TUNNEL_ENCAP_A_TUNNEL_TYPE_WIREGUARD) {
    log(L_TRACE "WG: tunnel type error %d", tunnel_type);
    return -1;
  }

  u16 value_length = get_u16(p + 2);

  log(L_TRACE "WG: tunnel encap value len %d", value_length);

  if (len < value_length + 4) {
    log(L_TRACE "WG: tunnel encap len error %d", value_length);
    return -1;
  }

  for (u8 *cur = p + 4; cur < p + 4 + value_length;) {
    int res = decode_sub_tlv(cur, value_length, pubkey, as4, remote_ep, color, udp_port, flags);

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

    cur += res;
  }

  return 0;
}

static wg_peer *
add_peer(wg_device *dev, wg_key pubkey)
{
  struct wg_peer *peer = malloc(sizeof(struct wg_peer));
  memset(peer, 0, sizeof(struct wg_peer));

  peer->flags = WGPEER_HAS_PUBLIC_KEY;
  memcpy(peer->public_key, pubkey, sizeof(wg_key));

  if (dev->first_peer && dev->last_peer)
    dev->last_peer->next_peer = peer;
  else
    dev->first_peer = peer;
  dev->last_peer = peer;
  return peer;
}

static int
set_peer_remote_ep(wg_peer *peer, ip_addr remote_ep_addr, u16 udp_dest_port)
{
  if (udp_dest_port != 0 && ipa_nonzero(remote_ep_addr) ) {
    if (ipa_is_ip4(remote_ep_addr)) {
      log(L_TRACE "WG:     found ip4 ep");
      peer->endpoint.addr4.sin_family = AF_INET;
      put_ip4(&peer->endpoint.addr4.sin_addr.s_addr, ipa_to_ip4(remote_ep_addr));
      put_u16(&peer->endpoint.addr4.sin_port, udp_dest_port);
    } else {
      log(L_TRACE "WG:     found ip6 ep");
      peer->endpoint.addr6.sin6_family = AF_INET6;
      put_ip6(&peer->endpoint.addr6.sin6_addr, ipa_to_ip6(remote_ep_addr));
      put_u16(&peer->endpoint.addr6.sin6_port, udp_dest_port);
    }
  }

  return 0;
}

static int
add_allowed_ips(u8 net_type, struct network *n, wg_peer *peer)
{
  // Add allowed ip
  struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
  memset(allowedip, 0, sizeof(struct wg_allowedip));

  if (net_type == NET_IP4) {
    allowedip->family = AF_INET;
    allowedip->ip4.s_addr = ip4_to_u32(ip4_hton(net4_prefix(n->n.addr)));
  } else if (net_type == NET_IP6) {
    allowedip->family = AF_INET6;
    ip6_addr addr = ip6_hton(net6_prefix(n->n.addr));
    memcpy(allowedip->ip6.s6_addr, &addr, 16);
  }

  allowedip->cidr = net_pxlen(n->n.addr);
  if (peer->first_allowedip && peer->last_allowedip)
    peer->last_allowedip->next_allowedip = allowedip;
  else
    peer->first_allowedip = allowedip;
  peer->last_allowedip = allowedip;

  return 0;
}

static void
wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
  struct rte *new, struct rte *old UNUSED)
{
  struct wg_proto *p = (struct wg_proto *) P;
  struct wg_config *c = (struct wg_config *) P->cf;
  struct wg_channel *ch = (struct wg_channel *) CH;
  struct wg_entry *en;
  struct iface *iface = NULL;
  const char *ifname = NULL;

//  debug("WG: notify\n");

  if (new) {
    iface = new->attrs->nh.iface;
    ifname = iface ? iface->name : NULL;

    if (iface && iface == p->iface) {
      struct eattr *t;

      debug("WG: found iface\n");
      en = fib_get(&ch->rtable, n->n.addr);

      debug("WG: notify new %d %N\n",
        new->attrs->dest, n->n.addr);

      wg_key pubkey;
      memset(pubkey, 0, sizeof(wg_key));
      u32 remote_ep_as4 = 0;
      ip_addr remote_ep_addr = IPA_NONE;
      u16 udp_dest_port = 0;
      u32 color = 0;
      u16 flags = 0;

      t = ea_find(new->attrs->eattrs, EA_CODE(PROTOCOL_BGP, BA_TUNNEL_ENCAP));
      if (t && t->u.ptr && decode_tunnel_encap(t, &pubkey, &remote_ep_as4, &remote_ep_addr, &color, &udp_dest_port, &flags) == 0) {
	log(L_TRACE "WG:     Attr %x %x %d %04x", t->flags, t->type, t->u.ptr->length, flags);

	struct wg_device *dev = p->dev;

	if (dev != NULL) {
	  bool dirty = false;
	  bool found = false;
	  struct wg_peer *peer = NULL;
	  wg_for_each_peer(dev, peer) {
	    // Look for public key
	    size_t len = 32; // FIXME
	    // MIN(32, t->u.ptr->length)
	    if (memcmp(peer->public_key, pubkey, len) != 0) {
	      log(L_TRACE "WG:     Not found");
	      continue;
	    }

	    log(L_TRACE "WG:     Found");
	    found = true;

	    set_peer_remote_ep(peer, remote_ep_addr, udp_dest_port);
	    add_allowed_ips(ch->c.net_type, n, peer);
	    dirty = true;

	    break;
	  }

	  if (!found) {
	    wg_peer *peer = add_peer(dev, pubkey);
	    set_peer_remote_ep(peer, remote_ep_addr, udp_dest_port);
	    add_allowed_ips(ch->c.net_type, n, peer);
	    dirty = true;
	  }

	  if (dirty) {
	    int res = wg_set_device(dev);
	    log(L_TRACE "WG: wg_set_device %d", res);
	  }
	}

	
/*
	struct sub_tlv_remove_endpoint {
		u32 asn;
		u8 address_family;
		union {
			ip4_addr ip4;
			ip6_addr ip6;
		};
	};

	struct sub_tlv_udp_dest_port {
		u16 port;
	};

	struct sub_tlv_wireguard {
		u8 pubkey[32];
	};

	struct sub_tlv {
		u8 type;
		union {
			struct {
				u8 length;
				struct sub_tlv_value value[];
			} s8;
			struct {
				u16 length;
				struct sub_tlv_value value[];
			} s16;
		} u;
	};

	struct bgp_tunnel_encap_attr {
		u16 type;
		u16 length;
		struct sub_tlv stlv[];
	};
*/
      } else {
        log(L_TRACE "WG:     No Attr");
      }

  //  old_metric = en->valid ? en->metric : -1;

//  en->valid = RIP_ENTRY_VALID;
//  en->metric = rt_metric;
//  en->tag = rt_tag;
//  en->from = (new->attrs->src->proto == P) ? new->u.rip.from : NULL;
//  en->iface = new->attrs->iface;
//  en->next_hop = new->attrs->gw;
      }

  } else {
    debug("WG: notify withdraw %N\n",
	  n->n.addr);

    /* Withdraw */
    en = fib_find(&ch->rtable, n->n.addr);

    if (!en) { // || en->valid != RIP_ENTRY_VALID)
      debug("WG: fib not found\n");
      return;
    }

	struct wg_device *dev = p->dev;

	if (dev != NULL) {
	  bool found = false;
	  struct wg_peer *peer = NULL;
	  wg_for_each_peer(dev, peer) {
	    // Remove from all peers
	    found = true;

	    if (found) {
	      // Add allowed ip
	      struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
	      memset(allowedip, 0, sizeof(struct wg_allowedip));

	      if (ch->c.net_type == NET_IP4) {
		allowedip->family = AF_INET;
		allowedip->ip4.s_addr = ip4_to_u32(ip4_hton(net4_prefix(n->n.addr)));
	      } else if (ch->c.net_type == NET_IP6) {
		allowedip->family = AF_INET6;
		ip6_addr addr = ip6_hton(net6_prefix(n->n.addr));
		memcpy(allowedip->ip6.s6_addr, &addr, 16);
	      }

	      allowedip->cidr = net_pxlen(n->n.addr);

	      struct wg_allowedip *ip = NULL;
	      struct wg_allowedip *previp = NULL;
	      
	      wg_for_each_allowedip(peer, ip) {
		if (allowedip->family != ip->family) {
		  debug("WG: family no match\n");
		  previp = ip;
		  continue;
		}

		if (allowedip->cidr != ip->cidr) {
		  debug("WG: cidr no match\n");
		  previp = ip;
		  continue;
		}
		
		if (memcmp(&allowedip->ip6, &ip->ip6, sizeof(struct in6_addr))) {
		  debug("WG: ip no match\n");
		  dump(&allowedip->ip6, sizeof(struct in6_addr));
		  dump(&ip->ip6, sizeof(struct in6_addr));
		  previp = ip;
		  continue;
		}

		debug("WG: found ip\n");

		if (!previp) {
		  debug("WG: remove first\n");
		  peer->first_allowedip = ip->next_allowedip;
		  if (peer->last_allowedip == ip) {
		    peer->last_allowedip = NULL;
		    peer->first_allowedip = NULL;
		  }
		} else {
		  debug("WG: remove middle\n");
		  // Remove
		  if (peer->last_allowedip == ip)
		      peer->last_allowedip = previp;

		  previp->next_allowedip = ip->next_allowedip;
		}

		free(ip);
		break;
	      }
	      
	      peer->flags |= WGPEER_REPLACE_ALLOWEDIPS;
	      int res = wg_set_device(dev);
	      log(L_TRACE "WG: wg_set_device %d", res);
	    }
	  }

	  /*
    old_metric = en->metric;

    en->valid = RIP_ENTRY_STALE;
    en->metric = p->infinity;
    en->tag = 0;
    en->from = NULL;
    en->iface = NULL;
    en->next_hop = IPA_NONE;
*/
	}
  }

//  TRACE(D_EVENTS, "wg notify %s %s", src_table->name, ifname?ifname:"(null)");
}

static void
wg_reload_routes(struct channel *C)
{
  struct wg_proto *p UNUSED = (struct wg_proto *) C->proto;

  debug("reload routes\n");

// TODO
//  WALK_LIST(c, p->channels)
//    channel_request_feeding(c);
}


static struct proto *
wg_init(struct proto_config *C)
{
  struct wg_config *c UNUSED = (struct wg_config *) C;
  struct proto *P = proto_new(C);
  struct wg_proto *p UNUSED = (struct wg_proto *) P;

  debug("init\n");

  P->if_notify = wg_if_notify;
  P->rt_notify = wg_rt_notify;
  P->reload_routes = wg_reload_routes;
//  P->accept_ra_types = RA_ANY;

  /* Add all channels */
  struct wg_channel_config *cc;
  WALK_LIST(cc, C->channels)
    proto_add_channel(P, &cc->c);

  return P;
}


static int
wg_start(struct proto *P)
{
  struct wg_config *cf UNUSED = (struct wg_config *) P->cf;
  struct wg_proto *p = (struct wg_proto *) P;

  log(L_TRACE "WG: start");

  if (get_device(p, &p->dev, cf->ifname) >= 0)
    {
      int res = wg_set_device(p->dev);
      log(L_TRACE "WG: wg_set_device %d", res);
    }

  struct wg_channel *ch;
  WALK_LIST(ch,p->p.channels) {
    fib_init(&ch->rtable, P->pool, ch->c.net_type, sizeof(struct wg_entry),
	     OFFSETOF(struct wg_entry, n), 0, wg_init_entry);
  }
  return PS_UP;
}

static void
wg_dump(struct proto *P)
{
  struct wg_proto *p = (struct wg_proto *) P;
  int i;

  i = 0;

  struct wg_channel *ch;
  WALK_LIST(ch,p->p.channels) {
    FIB_WALK(&ch->rtable, struct wg_entry, en)
      {
	//    struct wg_entry *en = (struct wg_entry *) e;
	debug("WG: entry #%d:\n",
	      i++);
      }
    FIB_WALK_END;
  }
}

struct peer_config *peer_new(struct wg_config *c)
{
  struct peer_config *pc = cfg_allocz(sizeof(struct peer_config));
  debug("peer_new %p\n", pc);
  add_tail(&c->peers, (node*)pc);
  return pc;
}

struct channel_class channel_wg = {
  .channel_size =       sizeof(struct wg_channel),
  .config_size =        sizeof(struct wg_channel_config),
  /* .init =               wg_channel_init, */
  /* .start =              wg_channel_start, */
  /* .shutdown =           wg_channel_shutdown, */
  /* .cleanup =            wg_channel_cleanup, */
  /* .reconfigure =        wg_channel_reconfigure, */
};

struct protocol proto_wireguard = {
  .name =		"Wireguard",
  .template =		"wg%d",
  .class =              PROTOCOL_WG,
  .channel_mask =       NB_IP,
  .proto_size =         sizeof(struct wg_proto),
  .config_size =	sizeof(struct wg_config),
  .postconfig =		wg_postconfig,
  .init =		wg_init,
  .start =		wg_start,
  .dump = 		wg_dump,
/*  .multitable =		1,
  .preference =		DEF_PREF_PIPE,
  .cleanup =		wg_cleanup,
  .reconfigure =	wg_reconfigure,
  .copy_config = 	wg_copy_config,
  .get_status = 	wg_get_status,
  .show_proto_info = 	wg_show_proto_info*/
};