#ifndef INDOOR_WIFISCANNER_LINUX_H
#define INDOOR_WIFISCANNER_LINUX_H

/*
 * scan_access_points.c: Prints all detected access points with wlan0 using NL80211 (netlink).
 *
 * Only works on network interfaces whose drivers are compatible with Netlink. Test this by running `iw list`.
 *
 * Since only privileged users may submit NL80211_CMD_TRIGGER_SCAN, you'll have to run the compiled program as root.
 *
 * Build with: gcc $(pkg-config --cflags --libs libnl-genl-3.0) scan_access_points.c
 *
 * Raspbian prerequisites:
 *      sudo apt-get install libnl-genl-3-dev
 *
 * Resources:
 *      http://git.kernel.org/cgit/linux/kernel/git/jberg/iw.git/tree/scan.c
 *      http://stackoverflow.com/questions/21601521/how-to-use-the-libnl-library-to-trigger-nl80211-commands
 *      http://stackoverflow.com/questions/23760780/how-to-send-single-channel-scan-request-to-libnl-and-receive-single-
 */
#include <errno.h>
#include <netlink/errno.h>
#include <netlink/netlink.h>
#include <netlink/genl/genl.h>
#include <netlink/genl/ctrl.h>
#include <linux/nl80211.h>
#include <net/if.h>
#include <stdio.h>
#include <ctype.h>
#include <string>
#include <thread>

#include "WiFiRAW.h"
#include "WiFiScan.h"

class WiFiScanLinux : public WiFiScan {

	struct TMP {
		Timestamp tsStart;
		WiFiMeasurements res;
	};

	struct trigger_results {
		int done;
		int aborted;
	};

	struct handler_args {  // For family_handler() and nl_get_multicast_id().
		const char *group;
		int id;
	};

	static int error_handler(struct sockaddr_nl* nla, struct nlmsgerr* err, void* arg) {
		(void) nla;
		// Callback for errors.
		printf("error_handler() called.\n");
		int* ret = (int*) arg;
		*ret = err->error;
		return NL_STOP;
	}

	static int finish_handler(struct nl_msg* msg, void* arg) {
		(void) msg;
		// Callback for NL_CB_FINISH.
		int* ret = (int*) arg;
		*ret = 0;
		return NL_SKIP;
	}

	static int ack_handler(struct nl_msg* msg, void* arg) {
		(void) msg;
		// Callback for NL_CB_ACK.
		int* ret = (int*)arg;
		*ret = 0;
		return NL_STOP;
	}

	static int no_seq_check(struct nl_msg* msg, void* arg) {
		(void) msg;
		(void) arg;
		// Callback for NL_CB_SEQ_CHECK.
		return NL_OK;
	}

	static int family_handler(struct nl_msg* msg, void* arg) {
		// Callback for NL_CB_VALID within nl_get_multicast_id(). From http://sourcecodebrowser.com/iw/0.9.14/genl_8c.html.
		struct handler_args* grp = (struct handler_args*) arg;
		struct nlattr *tb[CTRL_ATTR_MAX + 1];
		struct genlmsghdr* gnlh = (struct genlmsghdr*) nlmsg_data(nlmsg_hdr(msg));
		struct nlattr *mcgrp;
		int rem_mcgrp;

		nla_parse(tb, CTRL_ATTR_MAX, genlmsg_attrdata(gnlh, 0), genlmsg_attrlen(gnlh, 0), NULL);

		if (!tb[CTRL_ATTR_MCAST_GROUPS]) return NL_SKIP;

		nla_for_each_nested(mcgrp, tb[CTRL_ATTR_MCAST_GROUPS], rem_mcgrp) {  // This is a loop.
			struct nlattr* tb_mcgrp[CTRL_ATTR_MCAST_GRP_MAX + 1];

			nla_parse(tb_mcgrp, CTRL_ATTR_MCAST_GRP_MAX, (struct nlattr*) nla_data(mcgrp), nla_len(mcgrp), NULL);

			if (!tb_mcgrp[CTRL_ATTR_MCAST_GRP_NAME] || !tb_mcgrp[CTRL_ATTR_MCAST_GRP_ID]) continue;
			if (strncmp((const char*) nla_data(tb_mcgrp[CTRL_ATTR_MCAST_GRP_NAME]), grp->group,
					nla_len(tb_mcgrp[CTRL_ATTR_MCAST_GRP_NAME]))) {
				continue;
					}

			grp->id = nla_get_u32(tb_mcgrp[CTRL_ATTR_MCAST_GRP_ID]);
			break;
		}

		return NL_SKIP;
	}


	int nl_get_multicast_id(struct nl_sock *sock, const char *family, const char *group) {

		// From http://sourcecodebrowser.com/iw/0.9.14/genl_8c.html.
		struct nl_msg *msg;
		struct nl_cb *cb;
		int ret, ctrlid;

		//struct handler_args grp = { .group = group, .id = -ENOENT, };
		struct handler_args grp;
		grp.group = group;
		grp.id = -ENOENT;

		msg = nlmsg_alloc();
		if (!msg) return -ENOMEM;

		cb = nl_cb_alloc(NL_CB_DEFAULT);
		if (!cb) {
			ret = -ENOMEM;
			goto out_fail_cb;
		}

		ctrlid = genl_ctrl_resolve(sock, "nlctrl");

		genlmsg_put(msg, 0, 0, ctrlid, 0, 0, CTRL_CMD_GETFAMILY, 0);

		ret = -ENOBUFS;
		NLA_PUT_STRING(msg, CTRL_ATTR_FAMILY_NAME, family);

		ret = nl_send_auto_complete(sock, msg);
		if (ret < 0) goto out;

		ret = 1;

		nl_cb_err(cb, NL_CB_CUSTOM, error_handler, &ret);
		nl_cb_set(cb, NL_CB_ACK, NL_CB_CUSTOM, ack_handler, &ret);
		nl_cb_set(cb, NL_CB_VALID, NL_CB_CUSTOM, family_handler, &grp);

		while (ret > 0) nl_recvmsgs(sock, cb);

		if (ret == 0) ret = grp.id;

		nla_put_failure:
			out:
				nl_cb_put(cb);
			out_fail_cb:
				nlmsg_free(msg);
				return ret;
	}


	static int callback_trigger(struct nl_msg *msg, void *arg) {

		// Called by the kernel when the scan is done or has been aborted.
		struct genlmsghdr* gnlh = (struct genlmsghdr*) nlmsg_data(nlmsg_hdr(msg));
		struct trigger_results* results = (struct trigger_results*) arg;

		//printf("Got something.\n");
		//printf("%d\n", arg);
		//nl_msg_dump(msg, stdout);

		if (gnlh->cmd == NL80211_CMD_SCAN_ABORTED) {
			printf("Got NL80211_CMD_SCAN_ABORTED.\n");
			results->done = 1;
			results->aborted = 1;
		} else if (gnlh->cmd == NL80211_CMD_NEW_SCAN_RESULTS) {
			printf("Got NL80211_CMD_NEW_SCAN_RESULTS.\n");
			results->done = 1;
			results->aborted = 0;
		}  // else probably an uninteresting multicast message.

		return NL_SKIP;
	}



	static int addResult(struct nl_msg* msg, void* arg) {

		TMP* tmp = (TMP*) arg;

		// Called by the kernel with a dump of the successful scan's data. Called for each SSID.
		struct genlmsghdr* gnlh = (struct genlmsghdr*) nlmsg_data(nlmsg_hdr(msg));
	//	char mac_addr[20];
		struct nlattr *tb[NL80211_ATTR_MAX + 1];
		struct nlattr *bss[NL80211_BSS_MAX + 1];

		static struct nla_policy bss_policy[NL80211_BSS_MAX + 1];
		memset(&bss_policy, 0, sizeof(bss_policy));
			bss_policy[NL80211_BSS_TSF].type = NLA_U64;
			bss_policy[NL80211_BSS_FREQUENCY].type = NLA_U32;
//			bss_policy[NL80211_BSS_BSSID] = { };
			bss_policy[NL80211_BSS_BEACON_INTERVAL].type = NLA_U16;
			bss_policy[NL80211_BSS_CAPABILITY].type = NLA_U16;
//			bss_policy[NL80211_BSS_INFORMATION_ELEMENTS] = { };
			bss_policy[NL80211_BSS_SIGNAL_MBM].type = NLA_U32;
			bss_policy[NL80211_BSS_SIGNAL_UNSPEC].type = NLA_U8;
			bss_policy[NL80211_BSS_STATUS].type = NLA_U32;
			bss_policy[NL80211_BSS_SEEN_MS_AGO].type = NLA_U32;
//			bss_policy[NL80211_BSS_BEACON_IES] = { };


		// Parse and error check.
		nla_parse(tb, NL80211_ATTR_MAX, genlmsg_attrdata(gnlh, 0), genlmsg_attrlen(gnlh, 0), NULL);
		if (!tb[NL80211_ATTR_BSS]) {
			printf("bss info missing!\n");
			return NL_SKIP;
		}
		if (nla_parse_nested(bss, NL80211_BSS_MAX, tb[NL80211_ATTR_BSS], bss_policy)) {
			printf("failed to parse nested attributes!\n");
			return NL_SKIP;
		}
		if (!bss[NL80211_BSS_BSSID]) return NL_SKIP;
		if (!bss[NL80211_BSS_INFORMATION_ELEMENTS]) return NL_SKIP;

		const uint64_t seen_ago_ms = nla_get_u32(bss[NL80211_BSS_SEEN_MS_AGO]);
		const int rssi = (nla_get_s32(bss[NL80211_BSS_SIGNAL_MBM])) / 100.0f;

//		// Start printing.
//		mac_addr_n2a(mac_addr, (unsigned char*) nla_data(bss[NL80211_BSS_BSSID]));
//		printf("%s, ", mac_addr);
//		printf("%d MHz, ", nla_get_u32(bss[NL80211_BSS_FREQUENCY]));
//		//print_ssid((unsigned char*) nla_data(bss[NL80211_BSS_INFORMATION_ELEMENTS]), nla_len(bss[NL80211_BSS_INFORMATION_ELEMENTS]));
//		printf(" %d ms", seen_ago_ms);
//		printf(" %d dBm", rssi);
//		printf("\n");

		WiFiRAW::TaggedParams params = WiFiRAW::parseTaggedParams(
			(const uint8_t*) nla_data(bss[NL80211_BSS_INFORMATION_ELEMENTS]),
			nla_len(bss[NL80211_BSS_INFORMATION_ELEMENTS])
		);

		const Timestamp ts = Timestamp::fromUnixTime() - Timestamp::fromMS(seen_ago_ms);
		const int freq_MHz = nla_get_u32(bss[NL80211_BSS_FREQUENCY]);
		const uint8_t* macPtr = (const uint8_t*) nla_data(bss[NL80211_BSS_BSSID]);
		const uint64_t macLng =  ((uint64_t)macPtr[5]<<40)|((uint64_t)macPtr[4]<<32)|((uint64_t)macPtr[3]<<24)|((uint64_t)macPtr[2]<<16)|((uint64_t)macPtr[1]<<8)|((uint64_t)macPtr[0]<<0);
		const MACAddress mac(macLng);
		const AccessPoint ap(mac, params.ssid);
		const WiFiMeasurement mes(ap, rssi, freq_MHz, ts);

		// by default, linux also lists older scan results
		// remove them here!
		if (ts > tmp->tsStart) {
			//std::cout << seen_ago_ms << std::endl;
			tmp->res.entries.push_back(mes);
			std::cout << mes.asString() << std::endl;
		}



		return NL_SKIP;

	}


//	int do_scan_trigger(struct nl_sock *socket, int if_index, int driver_id) {

//		// Starts the scan and waits for it to finish. Does not return until the scan is done or has been aborted.
//		struct trigger_results results;
//		results.done = 0;
//		results.aborted = 0;

//		struct nl_msg *msg;
//		struct nl_cb *cb;
//		struct nl_msg *ssids_to_scan;
//		int err;
//		int ret;
//		int mcid = nl_get_multicast_id(socket, "nl80211", "scan");
//		nl_socket_add_membership(socket, mcid);  // Without this, callback_trigger() won't be called.

//		// Allocate the messages and callback handler.
//		msg = nlmsg_alloc();
//		if (!msg) {throw Exception("Failed to allocate netlink message for msg.");}

//		ssids_to_scan = nlmsg_alloc();
//		if (!ssids_to_scan) {
//			nlmsg_free(msg);
//			throw Exception("Failed to allocate netlink message for ssids_to_scan.");
//		}

//		cb = nl_cb_alloc(NL_CB_DEFAULT);
//		if (!cb) {
//			nlmsg_free(msg);
//			nlmsg_free(ssids_to_scan);
//			throw Exception("Failed to allocate netlink callbacks.");
//		}

//		// Setup the messages and callback handler.
//		genlmsg_put(msg, 0, 0, driver_id, 0, 0, NL80211_CMD_TRIGGER_SCAN, 0);  // Setup which command to run.
//		nla_put_u32(msg, NL80211_ATTR_IFINDEX, if_index);  // Add message attribute, which interface to use.
//		nla_put(ssids_to_scan, 1, 0, "");  // Scan all SSIDs.
//		nla_put_nested(msg, NL80211_ATTR_SCAN_SSIDS, ssids_to_scan);  // Add message attribute, which SSIDs to scan for.
//		nlmsg_free(ssids_to_scan);  // Copied to `msg` above, no longer need this.
//		nl_cb_set(cb, NL_CB_VALID, NL_CB_CUSTOM, callback_trigger, &results);  // Add the callback.
//		nl_cb_err(cb, NL_CB_CUSTOM, error_handler, &err);
//		nl_cb_set(cb, NL_CB_FINISH, NL_CB_CUSTOM, finish_handler, &err);
//		nl_cb_set(cb, NL_CB_ACK, NL_CB_CUSTOM, ack_handler, &err);
//		nl_cb_set(cb, NL_CB_SEQ_CHECK, NL_CB_CUSTOM, no_seq_check, NULL);  // No sequence checking for multicast messages.

//		// Send NL80211_CMD_TRIGGER_SCAN to start the scan. The kernel may reply with NL80211_CMD_NEW_SCAN_RESULTS on
//		// success or NL80211_CMD_SCAN_ABORTED if another scan was started by another process.
//		err = 1;
//		ret = nl_send_auto(socket, msg);  // Send the message.
//		printf("NL80211_CMD_TRIGGER_SCAN sent %d bytes to the kernel.\n", ret);
//		printf("Waiting for scan to complete...\n");
//		while (err > 0) ret = nl_recvmsgs(socket, cb);  // First wait for ack_handler(). This helps with basic errors.
//		if (err < 0) {
//			printf("WARNING: err has a value of %d.\n", err);
//		}
//		if (ret < 0) {
//			printf("ERROR: nl_recvmsgs() returned %d (%s).\n", ret, nl_geterror(-ret));
//			return ret;
//		}
//		while (!results.done) nl_recvmsgs(socket, cb);  // Now wait until the scan is done or aborted.
//		if (results.aborted) {
//			printf("ERROR: Kernel aborted scan.\n");
//			return 1;
//		}
//		printf("Scan is done.\n");

//		// Cleanup.
//		nlmsg_free(msg);
//		nl_cb_put(cb);
//		nl_socket_drop_membership(socket, mcid);  // No longer need this.
//		return 0;
//	}


	int if_index;
	struct nl_sock* socket;
	int driver_id;


	struct nl_cb *cb;
	int mcid;

	int err;

	struct trigger_results results;

	/** configure all needed callback (from netlink to code) once */
	void setupOnce() {

		mcid = nl_get_multicast_id(socket, "nl80211", "scan");
		nl_socket_add_membership(socket, mcid);  // Without this, callback_trigger() won't be called.

		cb = nl_cb_alloc(NL_CB_DEFAULT);
		if (!cb) {throw Exception("Failed to allocate netlink callbacks.");}

		// Setup the messages and callback handler.
		nl_cb_set(cb, NL_CB_VALID, NL_CB_CUSTOM, callback_trigger, &results);  // Add the callback.
		nl_cb_err(cb, NL_CB_CUSTOM, error_handler, &err);
		nl_cb_set(cb, NL_CB_FINISH, NL_CB_CUSTOM, finish_handler, &err);
		nl_cb_set(cb, NL_CB_ACK, NL_CB_CUSTOM, ack_handler, &err);
		nl_cb_set(cb, NL_CB_SEQ_CHECK, NL_CB_CUSTOM, no_seq_check, NULL);  // No sequence checking for multicast messages.

	}

	/** triggers a new scan within the wifi hardware */
	void triggerNewScan() {

		std::cout << "triggerNewScan()" << std::endl;

		struct nl_msg *ssids_to_scan;
		ssids_to_scan = nlmsg_alloc();
		if (!ssids_to_scan) {throw Exception("Failed to allocate netlink message for ssids_to_scan.");}
		nla_put(ssids_to_scan, 1, 0, "");  // Scan all SSIDs.

		// construct message
		struct nl_msg* msg = nlmsg_alloc();
		if (!msg) {throw Exception("Failed to allocate netlink message for msg.");}

		genlmsg_put(msg, 0, 0, driver_id, 0, 0, NL80211_CMD_TRIGGER_SCAN, 0);  // Setup which command to run.
		nla_put_u32(msg, NL80211_ATTR_IFINDEX, if_index);  // Add message attribute, which interface to use.
		nla_put_nested(msg, NL80211_ATTR_SCAN_SSIDS, ssids_to_scan);  // Add message attribute, which SSIDs to scan for.
		nlmsg_free(ssids_to_scan);  // Copied to `msg` above, no longer need this.

		results.done = 0;
		results.aborted = 0;

		// trigger scan by sending the constructed message
		const int ret = nl_send_auto(socket, msg);  // Send the message.
		printf("NL80211_CMD_TRIGGER_SCAN sent %d bytes to the kernel.\n", ret);
		printf("Waiting for scan to complete...\n");
		nlmsg_free(msg);

	}

	/** blocks until the scan-result is available. true if OK, false otherwise */
	bool waitForScanResult() {

		// Send NL80211_CMD_TRIGGER_SCAN to start the scan. The kernel may reply with NL80211_CMD_NEW_SCAN_RESULTS on
		// success or NL80211_CMD_SCAN_ABORTED if another scan was started by another process.
		err = 0;
//		ret = nl_send_auto(socket, msg);  // Send the message.
//		printf("NL80211_CMD_TRIGGER_SCAN sent %d bytes to the kernel.\n", ret);
//		printf("Waiting for scan to complete...\n");
//		while (err > 0) ret = nl_recvmsgs(socket, cb);  // First wait for ack_handler(). This helps with basic errors.
//		if (err < 0) {
//			printf("WARNING: err has a value of %d.\n", err);
//		}

		while(true) {
			const int ret = nl_recvmsgs(socket, cb);
			printf("-- ret: %d err: %d \n", ret, err);
			if (results.done) {
				return true;
			}
			if (ret < 0 || err < 0) {
				nl_recvmsgs(socket, cb);	// seems to fix issues when device is busy?!
				printf("ERROR: nl_recvmsgs() returned %d (%s).\n", ret, nl_geterror(-ret));
				return false;
			}
		}

	}

	void scanResult(TMP* res) {

		// Now get info for all SSIDs detected.
		struct nl_msg *msg = nlmsg_alloc();  // Allocate a message.
		genlmsg_put(msg, 0, 0, driver_id, 0, NLM_F_DUMP, NL80211_CMD_GET_SCAN, 0);  // Setup which command to run.
		nla_put_u32(msg, NL80211_ATTR_IFINDEX, if_index);  // Add message attribute, which interface to use.
		nl_socket_modify_cb(socket, NL_CB_VALID, NL_CB_CUSTOM, addResult, res);  // Add the callback and the measurements to fill
		int ret = nl_send_auto(socket, msg);  // Send the message.
		printf("NL80211_CMD_GET_SCAN sent %d bytes to the kernel.\n", ret);
		ret = nl_recvmsgs_default(socket);  // Retrieve the kernel's answer. callback_dump() prints SSIDs to stdout.
		nlmsg_free(msg);
		if (ret < 0) {
			printf("ERROR: nl_recvmsgs_default() returned %d (%s).\n", ret, nl_geterror(-ret));
			throw "error";
		}

	}

	void scanCleanup() {

		// Cleanup.
		//nlmsg_free(msg);
		nl_cb_put(cb);
		nl_socket_drop_membership(socket, mcid);  // No longer need this.
		//return 0;

	}





public:

	WiFiScanLinux(const std::string& devName) {

		// convert interface-name to interface-index
		if_index = if_nametoindex(devName.c_str());

		// Open socket to kernel.
		socket = nl_socket_alloc();							// Allocate new netlink socket in memory.
		genl_connect(socket);								// Create file descriptor and bind socket.
		driver_id = genl_ctrl_resolve(socket, "nl80211");	// Find the nl80211 driver ID.

		setupOnce();

	}

	~WiFiScanLinux() {

		// cleanup
		nl_socket_free(socket);

	}



	/** triger WiFiScan and fetch the result */
	WiFiMeasurements scan() {


		TMP res;


		// Issue NL80211_CMD_TRIGGER_SCAN to the kernel and wait for it to finish.
//		while(true) {

//			// use the current timestamp to suppress older scan results
//			// which are cached by linux by default
//			res.tsStart = Timestamp::fromUnixTime();

//			// trigger a scan
//			//int err = do_scan_trigger(socket, if_index, driver_id);
//			int err = scanTrigger();
////			if (err == -25)	{std::this_thread::sleep_for(std::chrono::milliseconds(100)); continue;}				// currently busy. try again
////			if (err != 0)	{throw Exception("do_scan_trigger() failed with code: " + std::to_string(err));}
////			break;

//		}

		again:;

		triggerNewScan();

		std::cout << "scan triggered" << std::endl;

		if (waitForScanResult()) {

			std::cout << "scan done" << std::endl;
			scanResult(&res);

			// return constructed result
			return res.res;


		} else {

			std::this_thread::sleep_for(std::chrono::milliseconds(100));
			goto again;

		}




	}



};


// gcc -I /usr/include/libnl3 main.c -lnl-genl-3 -lnl-3



#endif // INDOOR_WIFISCANNER_LINUX_H