#include "./wgcpp.hpp"
#include "wireguard/wireguard.h"
#include <algorithm>
#include <cstring>
#include <netinet/in.h>
#include <sys/socket.h>
#include <vector>

std::map<std::string_view, std::vector<User>> Users::acl_aggr() {
  std::map<std::string_view, std::vector<User>> acl_list;
  for (const User &usr : this->users) {
    for (const std::string_view member : usr.acl_member) {
      auto x = acl_list.find(member);
      if (x != acl_list.end()) {
        acl_list.at(member).push_back(usr);
      } else {
        std::vector<User> acl_members;
        acl_members.push_back(usr);
        acl_list[member] = acl_members;
      }
    }
  }
  return acl_list;
}

void Users::sort_with_ip() {
  std::sort(this->users.begin(), this->users.end(),
            [](User &a, User &b) { return a.ipv4 < b.ipv4; });
}

WG::WG(Users &wg_users, const char name[16], std::uint16_t port,
       uint32_t network, uint8_t cidr)
    : users(wg_users), network_ipv4(network), cidr(cidr) {
  wg_generate_private_key(this->wg_dev.private_key);
  wg_key_to_base64(this->b64_private_key, this->wg_dev.private_key);
  wg_generate_public_key(this->wg_dev.public_key, this->wg_dev.private_key);
  wg_key_to_base64(this->b64_public_key, this->wg_dev.public_key);
  this->wg_dev.first_peer = nullptr;
  this->wg_dev.last_peer = nullptr;

  std::strncpy(this->wg_dev.name, name, 16);
  this->wg_dev.listen_port = port;
  this->wg_dev.flags = (enum wg_device_flags)(WGDEVICE_HAS_LISTEN_PORT |
                                              WGDEVICE_HAS_PRIVATE_KEY |
                                              WGDEVICE_HAS_PUBLIC_KEY);

  uint32_t wildcard = 0xffffffff;

  if (this->cidr < 32) {
    wildcard = ((1 << (32 - this->cidr)) - 1);
  }
  for (uint32_t i = 2; i < wildcard; i++) {
    if ((i ^ 255) != 0) {
      this->ipv4_pool.insert(this->network_ipv4 + i);
    }
  }
  for (const auto &user : wg_users.users) {
    this->ipv4_pool.erase(user.ipv4);
  }
}

WG::WG(Users &wg_users, const char name[16], std::uint16_t port,
       uint32_t network, uint8_t cidr, std::string_view srv_priv_key,
       std::string_view srv_pub_key)
    : users(wg_users), network_ipv4(network), cidr(cidr) {
  std::strncpy(this->b64_private_key, srv_priv_key.data(), 45);
  std::strncpy(this->b64_public_key, srv_pub_key.data(), 45);
  wg_key_from_base64(this->wg_dev.private_key, this->b64_private_key);
  wg_key_from_base64(this->wg_dev.public_key, this->b64_public_key);

  this->wg_dev.first_peer = nullptr;
  this->wg_dev.last_peer = nullptr;
  std::strncpy(this->wg_dev.name, name, 16);
  this->wg_dev.listen_port = port;
  this->wg_dev.flags = (enum wg_device_flags)(WGDEVICE_HAS_LISTEN_PORT |
                                              WGDEVICE_HAS_PRIVATE_KEY |
                                              WGDEVICE_HAS_PUBLIC_KEY);

  uint32_t wildcard = 0xffffffff;

  if (this->cidr < 32) {
    wildcard = ((1 << (32 - this->cidr)) - 1);
  }
  for (uint32_t i = 2; i < wildcard; i++) {
    if ((i ^ 255) != 0) {
      this->ipv4_pool.insert(this->network_ipv4 + i);
    }
  }
  for (const auto &user : wg_users.users) {
    this->ipv4_pool.erase(user.ipv4);
  }
}

void WG::merge() {
  User *previous = nullptr;
  for (User &user : this->users.users) {
    if (previous) {
      previous->node.next_peer = &user.node;
    } else {
      previous = &user;
    }
  }
}

User &WG::add_cfg(strv name) {
  uint32_t ipv4 = this->ipv4_alloc();
  return this->add_cfg(name, ipv4);
}

User &WG::add_cfg(strv name, uint32_t ipv4) {
  wg_key pub, priv;
  wg_generate_private_key(priv);
  wg_generate_public_key(pub, priv);
  User new_user(name, ipv4, pub, priv);
  this->users.users.push_back(new_user);
  return this->users.users.back();
}
std::vector<WGPeerStat> *WG::peer_state() {
  std::vector<WGPeerStat> *x = new std::vector<WGPeerStat>();
  for (User &user : this->users.users) {

    WGPeerStat stat;
    stat.active = !(user.node.flags & WGPEER_REMOVE_ME);
    stat.ipv4 = user.node.first_allowedip->ip4.s_addr;
    wg_key_to_base64(stat.pubkey, user.node.public_key);
    x->push_back(stat);
  }
  return x;
}

User::User(strv name, uint32_t ipv4, wg_key pubkey, wg_key privkey)
    : username(name), ipv4(ipv4) {
  memset(&this->node, 0, sizeof(wg_peer));
  this->node.flags = (enum wg_peer_flags)(WGPEER_HAS_PUBLIC_KEY | WGPEER_REPLACE_ALLOWEDIPS);
  wg_allowedip *allowed_ip = new wg_allowedip;
  allowed_ip->cidr = 32;
  allowed_ip->family = AF_INET;
  allowed_ip->ip4.s_addr = htonl(ipv4);
  this->node.first_allowedip = allowed_ip;
  this->node.last_allowedip = allowed_ip;
  memcpy(this->node.public_key, pubkey, sizeof(uint8_t) * 32);
  memcpy(this->privkey, privkey, sizeof(uint8_t) * 32);
}