#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <errno.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>

#include "rtnetlink.h"

static inline int align4(int len)
{
  return (len + 3) & (~3);
}

struct newaddr_msg {
  struct nlmsghdr nl;
  struct ifaddrmsg ifaddr;
};

struct deladdr_msg {
  struct nlmsghdr nl;
  struct ifaddrmsg ifaddr;
  struct {
    struct rtattr rtattr;
    uint32_t value;
  } attr[2];
};

namespace rtnetlink {

  bool parse_newaddr(uint8_t const * buf, ssize_t length,
                     uint32_t ifa_index,
                     newaddr * out, int out_length,
                     int & index)
  {
    uint8_t const * bufi = buf;

    while (index < out_length) {
      struct nlmsghdr * nl = (struct nlmsghdr *)&bufi[0];
      struct ifaddrmsg * ifaddr = (struct ifaddrmsg *)&bufi[(sizeof (struct nlmsghdr))];
      size_t rtattr_offset = (sizeof (struct nlmsghdr)) + (sizeof (struct ifaddrmsg));

      assert((nl->nlmsg_flags & NLM_F_MULTI) == NLM_F_MULTI);

      if (nl->nlmsg_type == NLMSG_DONE) {
        return true;
      }

      assert(nl->nlmsg_type == RTM_NEWADDR);

      out[index].ifa_prefixlen = ifaddr->ifa_prefixlen;

      while (rtattr_offset < nl->nlmsg_len) {
        struct rtattr * rtattr = (struct rtattr *)&bufi[rtattr_offset];
        rtattr_offset += align4(rtattr->rta_len);

        void * value = (void *)(((ptrdiff_t)rtattr) + (sizeof (struct rtattr)));

        switch (rtattr->rta_type) {
        case IFA_ADDRESS:
          out[index].address = *((uint32_t *)value);
          break;
        case IFA_LOCAL:
          out[index].local = *((uint32_t *)value);
          break;
        }
      }

      if (ifaddr->ifa_index == ifa_index) {
        index += 1;
      }

      bufi = &bufi[nl->nlmsg_len];
      if (bufi - buf == length) {
        break;
      }
    }

    return false;
  }

  int handle_getaddr(int fd, uint32_t ifa_index, newaddr * out, int out_length)
  {
    bool done = false;
    int index = 0;

    while (!done) {
      struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
      struct iovec iov[1] = {};
      struct msghdr msg = {
        .msg_name = (void *)&addr,
        .msg_namelen = (sizeof (sockaddr_nl)),
        .msg_iov = iov,
        .msg_iovlen = 1,
        .msg_flags = MSG_TRUNC,
      };
      ssize_t peek_length = recvmsg(fd, &msg, MSG_PEEK | MSG_TRUNC);
      if (peek_length < 0) {
        perror("recvmsg(MSG_PEEK|MSG_TRUNC)");
        return -1;
      }

      uint8_t * buf = (uint8_t *)malloc(peek_length);

      iov[0].iov_base = buf;
      iov[0].iov_len = peek_length;
      msg.msg_flags = 0;
      ssize_t length = recvmsg(fd, &msg, 0);
      if (length < 0) {
        perror("recvmsg");
        return -1;
      }
      assert(peek_length == length);

      done = parse_newaddr(buf, length, ifa_index, out, out_length, index);
      free(buf);
    }

    return index;
  }

  int send_getaddr(int fd, uint32_t ifa_index)
  {
    const size_t message_length = (sizeof (struct newaddr_msg));

    struct newaddr_msg rtmmsg = {
      .nl = {
        .nlmsg_len = message_length,
        .nlmsg_type = RTM_GETADDR,
        .nlmsg_flags = NLM_F_REQUEST | NLM_F_DUMP,
        .nlmsg_seq = 1234,
      },
      .ifaddr = {
        .ifa_family = AF_INET,
        .ifa_scope = RT_SCOPE_UNIVERSE,
        .ifa_index = ifa_index,
      }
    };

    struct iovec iov[1] = {
      {
        .iov_base = &rtmmsg,
        .iov_len = message_length,
      }
    };
    struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
    struct msghdr msg = {
      .msg_name = (void *)&addr,
      .msg_namelen = (sizeof (sockaddr_nl)),
      .msg_iov = iov,
      .msg_iovlen = 1,
    };

    ssize_t sret = sendmsg(fd, &msg, 0);
    if (sret < 0) {
      perror("sendmsg");
      return -1;
    }

    return 0;
  }

  int send_typeaddr(int fd, uint16_t nlmsg_type, uint32_t ifa_index, uint8_t ifa_prefixlen, uint32_t local, uint32_t address)
  {
    const size_t message_length = (sizeof (struct deladdr_msg));

    struct deladdr_msg rtmmsg = {
      .nl = {
        .nlmsg_len = message_length,
        .nlmsg_type = nlmsg_type,
        .nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK,
        .nlmsg_seq = 12345,
      },
      .ifaddr = {
        .ifa_family = AF_INET,
        .ifa_prefixlen = ifa_prefixlen,
        .ifa_scope = RT_SCOPE_UNIVERSE,
        .ifa_index = ifa_index,
      },
      .attr = {
        {
          .rtattr = { .rta_len = 8, .rta_type = IFA_LOCAL },
          .value = local,
        },
        {
          .rtattr = { .rta_len = 8, .rta_type = IFA_ADDRESS },
          .value = address,
        }
      }
    };

    struct iovec iov[1] = {
      {
        .iov_base = &rtmmsg,
        .iov_len = message_length,
      }
    };
    struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
    struct msghdr msg = {
      .msg_name = (void *)&addr,
      .msg_namelen = (sizeof (sockaddr_nl)),
      .msg_iov = iov,
      .msg_iovlen = 1,
    };

    ssize_t sret = sendmsg(fd, &msg, 0);
    if (sret < 0) {
      perror("sendmsg");
      return -1;
    }

    return 0;
  }

  int handle_error(int fd)
  {
    struct sockaddr_nl addr = { .nl_family = AF_NETLINK };
    struct iovec iov[1] = {};
    struct msghdr msg = {
      .msg_name = (void *)&addr,
      .msg_namelen = (sizeof (sockaddr_nl)),
      .msg_iov = iov,
      .msg_iovlen = 1,
      .msg_flags = MSG_TRUNC,
    };
    ssize_t peek_length = recvmsg(fd, &msg, MSG_PEEK | MSG_TRUNC);
    if (peek_length < 0) {
      perror("recvmsg(MSG_PEEK|MSG_TRUNC)");
      return -1;
    }

    uint8_t * buf = (uint8_t *)malloc(peek_length);

    iov[0].iov_base = buf;
    iov[0].iov_len = peek_length;
    msg.msg_flags = 0;
    ssize_t length = recvmsg(fd, &msg, 0);
    if (length < 0) {
      perror("recvmsg");
      return -1;
    }
    assert(peek_length == length);

    struct nlmsghdr * nl = (struct nlmsghdr *)&buf[0];
    struct nlmsgerr * msgerr = (struct nlmsgerr *)&buf[(sizeof (struct nlmsghdr))];

    assert(nl->nlmsg_type == NLMSG_ERROR);
    if (msgerr->error < 0) {
      printf("msgerr %d %s\n", msgerr->error, strerror(-msgerr->error));
      return -1;
    }

    free(buf);

    return 0;
  }
}