extern "C" {
#include <netlink/genl/ctrl.h>
#include <netlink/genl/genl.h>
#include <netlink/genl/family.h>
#include <pthread.h>
#include <linux/nl80211.h>
#include <dirent.h>
#include "ieee80211.h"
}

#include "wireless.h"
#include "utils.h"
#include <iostream>
#include <QStringList>
#include <QRegularExpression>

#define BIT(x) (1ULL<<(x))

int IEEE80211Freq[][2] = {
  {1, 2412},
  {2, 2417},
  {3, 2422},
  {4, 2427},
  {5, 2432},
  {6, 2437},
  {7, 2442},
  {8, 2447},
  {9, 2452},
  {10, 2457},
  {11, 2462},
  {12, 2467},
  {13, 2472},
  {14, 2484},
  // We could do the math here, but what about 4ghz nonsense?
  // We'll do table lookups for now.
  {36, 5180},
  {37, 5185},
  {38, 5190},
  {39, 5195},
  {40, 5200},
  {41, 5205},
  {42, 5210},
  {43, 5215},
  {44, 5220},
  {45, 5225},
  {46, 5230},
  {47, 5235},
  {48, 5240},
  {52, 5260},
  {53, 5265},
  {54, 5270},
  {55, 5275},
  {56, 5280},
  {57, 5285},
  {58, 5290},
  {59, 5295},
  {60, 5300},
  {64, 5320},
  {100, 5500},
  {104, 5520},
  {108, 5540},
  {112, 5560},
  {116, 5580},
  {120, 5600},
  {124, 5620},
  {128, 5640},
  {132, 5660},
  {136, 5680},
  {140, 5700},
  {149, 5745},
  {150, 5750},
  {152, 5760},
  {153, 5765},
  {157, 5785},
  {160, 5800},
  {161, 5805},
  {165, 5825},
  {0, 0}
};

struct nl_callback {
  char *phyname;
  void *bands;
};

Wireless::Wireless(QString interface) :
QObject(0),
m_interface(interface),
m_bands(),
m_isValid(false),
m_isNL80211(false)
{
  void *handle = NULL;
  struct genl_family *family = NULL;
  struct nl_cache *cache = NULL;
  struct nl_msg *msg = nlmsg_alloc();
  struct nl_cb *cb = nl_cb_alloc(NL_CB_DEFAULT);
  QByteArray interfaceArray = m_interface.toLatin1();
  const char *interfaceStr = interfaceArray.constData();
  int err = 1;

  char *phyname = nl80211_find_parent(interfaceStr);

  if(phyname == NULL) {
    if(!(Utils::interfaceIndex(interfaceStr))) {
      fprintf(stderr, "Interface %s doesn't exist\n", interfaceStr);
    }else{
      fprintf(stderr, "could not find a parent phy device for interface %s, it isn't nl80211?\n", interfaceStr);
    }

    return;
  }else{
    m_isNL80211 = true;
  }

  if(nl80211_connect(interfaceStr, &handle, &cache, &family) < 0) {
    std::cerr << "could not connect to nl80211" << std::endl;
    return;
  }

  struct nl_callback callback = { .phyname = phyname, .bands = &m_bands };

  nl_cb_set(cb, NL_CB_VALID, NL_CB_CUSTOM, nl80211_freqlist_cb, &callback);
  nl_cb_set(cb, NL_CB_FINISH, NL_CB_CUSTOM, nl80211_finish_cb, &err);
  nl_cb_err(cb, NL_CB_CUSTOM, nl80211_error_cb, &err);

  genlmsg_put(msg, 0, 0, genl_family_get_id((struct genl_family *)family), 0, NLM_F_DUMP, NL80211_CMD_GET_WIPHY, 0);

  if(nl_send_auto_complete((struct nl_sock *)handle, msg) < 0) {
    fprintf(stderr, "%s: Failed to write nl80211 message\n",
        __FUNCTION__);
    nl80211_disconnect(handle, cache, family);
    return;
  }

  while(err)
    nl_recvmsgs((struct nl_sock *)handle, cb);

  nl80211_disconnect(handle, cache, family);

  m_isValid = true;
}

Wireless::~Wireless() {
  std::cerr << "del wireless" << std::endl;

  m_bands.clear();
}

QList<int> Wireless::allChannels() const {
  QList<int> channels;

  foreach(BandInfo info, m_bands) {
    foreach(ChannelInfo chan, info.channels) {
      channels.append(chan.chan);
    }
  }

  return channels;
}

QList<int> Wireless::allowedChannels(Direction dir) const {
  QList<int> channels;

  foreach(BandInfo info, m_bands) {
    foreach(ChannelInfo chan, info.channels) {
      if(!chan.disabled) {
        if(((dir & Direction_Both) == Direction_Both) && (!chan.passive)) {
          channels.append(chan.chan);
        }else if(((dir & Direction_RX) == Direction_RX)) {
          channels.append(chan.chan);
        }
      }
    }
  }

  return channels;
}

QList<int> Wireless::allFrequencies() const {
  QList<int> freqs;

  foreach(BandInfo band, m_bands) {
    foreach(ChannelInfo info, band.channels) {
      freqs.append(info.freq);
    }
  }

  return freqs;
}

QList<int> Wireless::allowedFrequencies(Direction dir) const {
  QList<int> freqs;

  foreach(BandInfo band, m_bands) {
    foreach(ChannelInfo info, band.channels) {
      if(!info.disabled) {
        if(((dir & Direction_Both) == Direction_Both) && (!info.passive)) {
          freqs.append(info.freq);
        }else if(((dir & Direction_RX) == Direction_RX)) {
          freqs.append(info.freq);
        }
      }
    }
  }

  return freqs;
}

bool Wireless::channelSupported(int chan) const {
  bool found = false;

  foreach(BandInfo band, m_bands) {
    foreach(ChannelInfo info, band.channels) {
      if(info.chan == chan) {
        found = true;
        break;
      }
    }
  }

  return found;
}

bool Wireless::frequencySupported(int freq) const {
  bool found = false;

  foreach(BandInfo band, m_bands) {
    foreach(ChannelInfo info, band.channels) {
      if(info.freq == freq) {
        found = true;
        break;
      }
    }
  }

  return found;
}

bool Wireless::isValid() const {
  return m_isValid;
}

bool Wireless::isNL80211() const {
  return m_isNL80211;
}

const QString& Wireless::name() const {
  return m_interface;
}

int Wireless::ChanToFreq(int in_chan) {
  int x = 0;

  while(IEEE80211Freq[x][0] != 0) {
    if(IEEE80211Freq[x][0] == in_chan) {
      return IEEE80211Freq[x][1];
    }

    ++x;
  }

  return in_chan;
}

QString Wireless::ChanToFreq(QString in_chan, QString outputFormat) {
  QString freq;
  QRegularExpression re("[^\\d]");
  in_chan = in_chan.replace(re, "");

  if(outputFormat.isEmpty()) {
    bool ok = false;
    int chanInt = in_chan.toInt(&ok);

    if(ok) {
      freq = QString::number(ChanToFreq(chanInt));
    }
  }

  return freq;
}

int Wireless::FreqToChan(int in_freq) {
  int x = 0;

  while(IEEE80211Freq[x][1] != 0) {
    if(IEEE80211Freq[x][1] == in_freq) {
      return IEEE80211Freq[x][0];
    }

    ++x;
  }

  return in_freq;
}

QString Wireless::FreqToChan(QString in_freq, QString outputFormat) {
  QString chan;
  QRegularExpression re("[^\\d]");
  in_freq = in_freq.replace(re, "");

  if(in_freq.length() < 4) {
    for(int i = in_freq.length() + 1; i <= 4; ++i) {
      in_freq.append("0");
    }
  }

  if(outputFormat.isEmpty()) {
    bool ok = false;
    int freqInt = in_freq.toInt(&ok);

    if(ok) {
      chan = QString::number(FreqToChan(freqInt));
    }
  }

  return chan;
}

int Wireless::nl80211_connect(const char * /*interface*/, void **handle, struct nl_cache **cache, struct genl_family **family) const {
  struct nl_sock *nl_handle;
  struct nl_cache *nl_cache;
  struct genl_family *nl80211;

  if((nl_handle = nl_socket_alloc()) == NULL) {
    fprintf(stderr, "%s failed to allocate nlhandle\n",
         __FUNCTION__);
    return -1;
  }

  if(genl_connect(nl_handle)) {
    fprintf(stderr, "%s failed to connect to generic netlink\n",
         __FUNCTION__);
    nl_socket_free(nl_handle);
    return -1;
  }

  if(genl_ctrl_alloc_cache(nl_handle, &nl_cache) != 0) {
    fprintf(stderr, "%s failed to allocate "
         "generic netlink cache\n", __FUNCTION__);
    nl_socket_free(nl_handle);
    return -1;
  }

  if ((nl80211 = genl_ctrl_search_by_name(nl_cache, "nl80211")) == NULL) {
    fprintf(stderr, "%s failed to find "
         "nl80211 controls, kernel may be too old\n", __FUNCTION__);
    nl_socket_free(nl_handle);
    return -1;
  }

  (*handle) = (void *) nl_handle;
  (*cache) = nl_cache;
  (*family) = nl80211;

  return 1;
}

void Wireless::nl80211_disconnect(void *handle, struct nl_cache *cache, struct genl_family *family) const {
  genl_family_put(family);
  nl_cache_free(cache);
  nl_socket_free((struct nl_sock *)handle);
}

int Wireless::nl80211_freqlist_cb(struct nl_msg *msg, void *arg) {
  struct nlattr *tb_msg[NL80211_ATTR_MAX + 1];
  struct genlmsghdr *gnlh = (struct genlmsghdr *) nlmsg_data(nlmsg_hdr(msg));
  struct nlattr *tb_band[NL80211_BAND_ATTR_MAX + 1];
  struct nlattr *tb_freq[NL80211_FREQUENCY_ATTR_MAX + 1];
  struct nlattr *nl_band, *nl_freq;
  int rem_band, rem_freq = 0;
  struct nl_callback *callback = (struct nl_callback*)arg;

  nla_parse(tb_msg, NL80211_ATTR_MAX, genlmsg_attrdata(gnlh, 0),
        genlmsg_attrlen(gnlh, 0), NULL);

  if (!tb_msg[NL80211_ATTR_WIPHY_BANDS]) {
    return NL_SKIP;
  }

  if (tb_msg[NL80211_ATTR_WIPHY_NAME]) {
    if (strcmp(nla_get_string(tb_msg[NL80211_ATTR_WIPHY_NAME]),
           callback->phyname) != 0) {
      return NL_SKIP;
    }
  }

  // Count the number of channels
  for (nl_band = (struct nlattr *) nla_data(tb_msg[NL80211_ATTR_WIPHY_BANDS]),
     rem_band = nla_len(tb_msg[NL80211_ATTR_WIPHY_BANDS]);
     nla_ok(nl_band, rem_band);
     nl_band = (struct nlattr *) nla_next(nl_band, &rem_band)) {

    nla_parse(tb_band, NL80211_BAND_ATTR_MAX, (struct nlattr *) nla_data(nl_band),
          nla_len(nl_band), NULL);

    QList<BandInfo> *bands = static_cast<QList<BandInfo>*>(callback->bands);
    BandInfo band;

    if(tb_band[NL80211_BAND_ATTR_HT_CAPA]) {
      __u16 cap = nla_get_u16(tb_band[NL80211_BAND_ATTR_HT_CAPA]);
      band.protocols = static_cast<Protocols>(band.protocols | PROTOCOL_80211_N);
      band.width |= WIDTH_20MHZ;

      if((cap & IEEE80211_HT_CAP_SUP_WIDTH_20_40)) { // 40-mhz 802.11n
        band.width |= WIDTH_40MHZ;
      }
    }

    for (nl_freq = (struct nlattr *) nla_data(tb_band[NL80211_BAND_ATTR_FREQS]),
       rem_freq = nla_len(tb_band[NL80211_BAND_ATTR_FREQS]);
       nla_ok(nl_freq, rem_freq);
       nl_freq = (struct nlattr *) nla_next(nl_freq, &rem_freq)) {

      nla_parse(tb_freq, NL80211_FREQUENCY_ATTR_MAX,
            (struct nlattr *) nla_data(nl_freq),
            nla_len(nl_freq), NULL);

      if (!tb_freq[NL80211_FREQUENCY_ATTR_FREQ])
        continue;

      uint32_t freq = nla_get_u32(tb_freq[NL80211_FREQUENCY_ATTR_FREQ]);

      ChannelInfo info;
      info.chan = FreqToChan(freq);
      info.freq = freq;
      info.disabled = (tb_freq[NL80211_FREQUENCY_ATTR_DISABLED]) ? 1 : 0;
      info.passive = (tb_freq[NL80211_FREQUENCY_ATTR_PASSIVE_SCAN]) ? 1 : 0;
      info.radar = (tb_freq[NL80211_FREQUENCY_ATTR_RADAR]) ? 1 : 0;
      info.max_txpower = nla_get_u32(tb_freq[NL80211_FREQUENCY_ATTR_MAX_TX_POWER]);

      // check supported channels to see what bands we support (only makes sense for A and B)
      // is there a better way to do this?
      if(!((band.protocols & PROTOCOL_80211_A) == PROTOCOL_80211_A)) {
        if(info.chan >= 36) { // first US 5GHz channel
          band.protocols = static_cast<Protocols>(band.protocols | PROTOCOL_80211_A);
          band.band = static_cast<Bands>(band.band | BAND_5GHZ);
        }
      }

      if(!((band.protocols & PROTOCOL_80211_B) == PROTOCOL_80211_B)) {
        if(info.chan >= 1) { // first 2.4GHz channel
          band.protocols = static_cast<Protocols>(band.protocols | PROTOCOL_80211_B);
          band.band = static_cast<Bands>(band.band | BAND_2GHZ);
        }
      }

      band.channels.append(info);
    }

    bands->append(band);
  }

  return NL_SKIP;
}

int Wireless::nl80211_error_cb(struct sockaddr_nl * /*nla*/, struct nlmsgerr *err, void *arg) {
  int *ret = (int *) arg;
  *ret = err->error;
  return NL_STOP;
}

int Wireless::nl80211_finish_cb(struct nl_msg * /*msg*/, void *arg) {
  int *ret = (int *) arg;
  *ret = 0;
  return NL_SKIP;
}

char* Wireless::nl80211_find_parent(const char *interface) const {
  DIR *devdir;
  struct dirent *devfile;
  char dirpath[1024];
  char *ret;

  snprintf(dirpath, 1024, "/sys/class/net/%s/phy80211/device/ieee80211", interface);

  if ((devdir = opendir(dirpath)) == NULL)
    return NULL;

  while ((devfile = readdir(devdir)) != NULL) {
    if(strncmp("phy", devfile->d_name, 3) == 0) {
      ret = strdup(devfile->d_name);
      closedir(devdir);
      return ret;
    }
  }

  closedir(devdir);

  return NULL;
}

const QList<BandInfo>& Wireless::bandMap() const {
  return m_bands;
}