#include "apgetconnected.h"

#include <iostream>
#include <sstream>
#include <iomanip>

#include <linux/nl80211.h>
#include <net/if.h>

#include <netlink/genl/genl.h>
#include <netlink/genl/ctrl.h>
#include <netlink/msg.h>


namespace ap
{

static int nl80211Init(struct nl_sock* nl_sock, int* nl80211_id)
{
	int err;

	if (!nl_sock)
	{
		std::cerr<<"Failed to allocate netlink socket.\n";
		return -ENOMEM;
	}

	if (genl_connect(nl_sock))
	{
		std::cerr<<"Failed to connect to generic netlink.\n";
		err = -ENOLINK;
		nl_socket_free(nl_sock);
		return err;
	}

	nl_socket_set_buffer_size(nl_sock, 8192, 8192);

	/* try to set NETLINK_EXT_ACK to 1, ignoring errors */
	err = 1;
	setsockopt(nl_socket_get_fd(nl_sock), SOL_NETLINK, NETLINK_EXT_ACK, &err, sizeof(err));

	*nl80211_id = genl_ctrl_resolve(nl_sock, "nl80211");
	if (*nl80211_id < 0)
	{
		std::cerr<<"nl80211 not found.\n";
		err = -ENOENT;
		nl_socket_free(nl_sock);
		return err;
	}

	return 0;
}

static int errorHandler(struct sockaddr_nl *nla, struct nlmsgerr *err, void *arg)
{
	(void)nla;
	(void)err;
	printf("netlink error\n");
	*reinterpret_cast<int*>(arg) = 0;
	return NL_STOP;
}

static int finishHandler(struct nl_msg *msg, void *arg)
{
	(void)msg;
	*reinterpret_cast<int*>(arg) = 0;
	return NL_SKIP;
}

static int ackHandler(struct nl_msg *msg, void *arg)
{
	(void)msg;
	*reinterpret_cast<int*>(arg) = 0;
	return NL_STOP;
}


static int extractDeviceMac(struct nl_msg *msg, void *arg)
{
	struct nlattr* tb[NL80211_ATTR_MAX + 1];
	struct genlmsghdr* gnlh = (struct genlmsghdr*)nlmsg_data(nlmsg_hdr(msg));

	std::vector<uint64_t>* deviceMacAddrs = reinterpret_cast<std::vector<uint64_t>*>(arg);

	nla_parse(tb, NL80211_ATTR_MAX, genlmsg_attrdata(gnlh, 0), genlmsg_attrlen(gnlh, 0), NULL);
	if (!tb[NL80211_ATTR_STA_INFO])
	{
		std::cerr<<"sta stats missing!\n";
		return -1;
	}

	const unsigned char* macChar = reinterpret_cast<const unsigned char*>(nla_data(tb[NL80211_ATTR_MAC]));
	uint64_t macAddr = 0;
	memcpy(&macAddr, macChar, 6);
	deviceMacAddrs->push_back(macAddr);

	return 0;
}

std::vector<uint64_t> connectedDevices(const std::string& ifDevice, int& error)
{
	struct nl_sock* nl_sock = nl_socket_alloc();
	int nl80211_id = 0;
	error = 0;

	if(nl80211Init(nl_sock, &nl80211_id))
	{
		std::cerr<<"Can not init nl80211\n";
		error = ERR_INIT;
	}

	signed long long devidx = if_nametoindex(ifDevice.c_str());
	if(!devidx)
	{
		std::cerr<<ifDevice<<" is not a valid if device\n";
		error = ERR_NO_DEV;
	}

	struct nl_msg *msg;
	msg = nlmsg_alloc();
	if(!msg) error = ERR_ALLOC;

	struct nl_cb *cb = nl_cb_alloc(NL_CB_DEFAULT);
	struct nl_cb *s_cb = nl_cb_alloc(NL_CB_DEFAULT);
	if (!cb || !s_cb)  error = ERR_ALLOC;

	std::vector<uint64_t> deviceMacAddrs;

	if(error == 0)
	{
		if(!genlmsg_put(msg, 0, 0, nl80211_id, 0, NLM_F_DUMP, NL80211_CMD_GET_STATION, 0))
		{
			std::cerr<<"genlmsg_put() failed\n";
			error = ERR_GENERAL;
		}

		if(error == 0)
		{
			NLA_PUT_U32(msg, NL80211_ATTR_IFINDEX, devidx);

			nl_socket_set_cb(nl_sock, s_cb);
			nl_send_auto_complete(nl_sock, msg);

			{
				int ret = 1;

				nl_cb_err(cb, NL_CB_CUSTOM, errorHandler, &ret);
				nl_cb_set(cb, NL_CB_FINISH, NL_CB_CUSTOM, finishHandler, &ret);
				nl_cb_set(cb, NL_CB_ACK, NL_CB_CUSTOM, ackHandler, &ret);
				nl_cb_set(cb, NL_CB_VALID, NL_CB_CUSTOM, extractDeviceMac, &deviceMacAddrs);

				while (ret > 0) nl_recvmsgs(nl_sock, cb);
			}
		}
	}

nla_put_failure:
	nl_cb_put(cb);
	nl_cb_put(s_cb);
	nlmsg_free(msg);
	nl_socket_free(nl_sock);
	return deviceMacAddrs;
}

std::string macAddrToString(uint64_t macAddr)
{
	unsigned char* macAddrPtr = reinterpret_cast<unsigned char*>(&macAddr);
	std::stringstream ss;
	ss<<std::uppercase<<std::setfill('0')<<std::setw(2)<<std::hex;
	for(unsigned int i = 0; i < 6; ++i) ss<<+macAddrPtr[i]<<':';
	std::string macString = ss.str();
	macString.pop_back();
	return macString;
}

}