#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <poll.h>

#include "jl_modbus.h"
#include "jl_wifi_ctrl.h"

extern uint32_t custom_tick_get(void);

#ifndef JL_WIFI_CTRL_PATH
#define JL_WIFI_CTRL_PATH "/var/run/wpa_supplicant/wlan0"
#endif

#define JL_WIFI_CTRL_RECONNECT_COOLDOWN_MS 2000
/* Used by the periodic SCAN trigger (JLV_WLAN_SCAN_INTERVAL_MS apart, so a
 * blocking wait here is cheap) and other short one-shot commands. STATUS and
 * SCAN_RESULTS reads no longer go through this: they're polled on a tight UI
 * cycle (every 1-2s) and are handled asynchronously instead (see
 * jl_wifi_ctrl_request() / jl_wifi_ctrl_poll()) so they never block the UI
 * task thread. */
#define JL_WIFI_CTRL_TXN_TIMEOUT_MS 250
/* Used by one-shot config/action commands (SET_NETWORK, ENABLE_NETWORK,
 * RECONNECT, ...) triggered from a button press rather than a poll tick.
 * wpa_supplicant can take a while to reply while it's busy processing a
 * just-finished scan, which is exactly when the user tends to fire these
 * (right after picking a network from the scan list) - 250ms was too
 * tight and made the command silently no-op, leaving the network half
 * configured. */
#define JL_WIFI_CTRL_ACTION_TIMEOUT_MS 1500
/* Safety net for an async STATUS/SCAN_RESULTS request that never gets a
 * reply (e.g. wpa_supplicant wedged) - without this pending_req would stay
 * stuck forever and every future request would silently no-op. */
#define JL_WIFI_CTRL_REQ_TIMEOUT_MS 2000
#define JL_WIFI_CTRL_SSID_RAW_MAX (JLV_WIFI_SSID_BYTES*4+1)
#define JL_WIFI_CTRL_SCAN_CACHE_MAX 32

typedef enum {
	JL_WIFI_CTRL_REQ_NONE = 0,
	JL_WIFI_CTRL_REQ_STATUS,
	JL_WIFI_CTRL_REQ_SCAN_RESULTS,
} jl_wifi_ctrl_req_t;

static struct wpa_ctrl* ctrl_cmd;
static struct wpa_ctrl* ctrl_mon;
static uint32_t last_connect_attempt_ms;
static uint32_t pending_events;

/* STATUS/SCAN_RESULTS are request/response commands on the single ctrl_cmd
 * socket, so only one can be in flight at a time. jl_wifi_ctrl_request()
 * sends without waiting; the reply is picked up later, non-blocking, from
 * jl_wifi_ctrl_poll() (already ticking every 100ms) and cached below. */
static jl_wifi_ctrl_req_t pending_req = JL_WIFI_CTRL_REQ_NONE;
static uint32_t pending_req_sent_ms;

static jl_wifi_ctrl_status_t cached_status;
static bool cached_status_valid;

static jl_wifi_ctrl_scan_result_t cached_scan[JL_WIFI_CTRL_SCAN_CACHE_MAX];
static int cached_scan_count;

static void jl_wifi_ctrl_apply_reply(jl_wifi_ctrl_req_t req, char* reply, size_t reply_len);

static void jl_wifi_ctrl_teardown(void)
{
	if (ctrl_cmd) {
		wpa_ctrl_close(ctrl_cmd);
		ctrl_cmd = NULL;
	}
	if (ctrl_mon) {
		wpa_ctrl_close(ctrl_mon);
		ctrl_mon = NULL;
	}

	/* Any request in flight belonged to the now-closed socket; its reply,
	 * if it ever arrives, is not coming back to us. Drop cached data too,
	 * so callers see "unavailable" instead of stale pre-disconnect state. */
	pending_req = JL_WIFI_CTRL_REQ_NONE;
	cached_status_valid = false;
	cached_scan_count = 0;
}

static bool jl_wifi_ctrl_ensure_connected(void)
{
	if (ctrl_cmd && ctrl_mon) {
		return true;
	}

	uint32_t now = custom_tick_get();
	if (last_connect_attempt_ms != 0 &&
			(uint32_t)(now - last_connect_attempt_ms) < JL_WIFI_CTRL_RECONNECT_COOLDOWN_MS) {
		return false;
	}
	last_connect_attempt_ms = now;

	jl_wifi_ctrl_teardown();

	ctrl_cmd = wpa_ctrl_open(JL_WIFI_CTRL_PATH);
	if (!ctrl_cmd) {
		return false;
	}

	ctrl_mon = wpa_ctrl_open(JL_WIFI_CTRL_PATH);
	if (!ctrl_mon || wpa_ctrl_attach(ctrl_mon) != 0) {
		log_debug("jl_wifi_ctrl: attach failed");
		jl_wifi_ctrl_teardown();
		return false;
	}

	log_info("jl_wifi_ctrl: connected to %s", JL_WIFI_CTRL_PATH);
	return true;
}

/* Bounded (~JL_WIFI_CTRL_TXN_TIMEOUT_MS) request/response, bypassing
 * wpa_ctrl_request()'s internal 10s blocking retry loop - it's fine for the
 * rare ATTACH/DETACH call (done via wpa_ctrl_attach() above), but far too
 * coarse for commands issued from a UI poll tick. */
static bool jl_wifi_ctrl_txn_timeout(struct wpa_ctrl* ctrl, const char* cmd, char* reply, size_t* reply_len, int timeout_ms)
{
	int fd = wpa_ctrl_get_fd(ctrl);
	if (fd < 0) {
		return false;
	}

	if (send(fd, cmd, strlen(cmd), 0) < 0) {
		log_debug("jl_wifi_ctrl: send failed, reconnecting");
		jl_wifi_ctrl_teardown();
		return false;
	}

	struct pollfd pfd = { .fd = fd, .events = POLLIN };
	if (poll(&pfd, 1, timeout_ms) <= 0 || !(pfd.revents & POLLIN)) {
		return false;
	}

	return wpa_ctrl_recv(ctrl, reply, reply_len) == 0;
}

/* Waits (up to timeout_ms) for a STATUS/SCAN_RESULTS request already in
 * flight to be answered, before a synchronous command is allowed to send on
 * the same socket - otherwise the two replies could arrive out of order and
 * get attributed to the wrong request. Rare in practice (async requests are
 * answered within a poll tick or two), so this basically never actually
 * waits. */
static void jl_wifi_ctrl_drain_pending(int timeout_ms)
{
	if (pending_req == JL_WIFI_CTRL_REQ_NONE) {
		return;
	}

	int fd = wpa_ctrl_get_fd(ctrl_cmd);
	struct pollfd pfd = { .fd = fd, .events = POLLIN };
	if (fd >= 0 && poll(&pfd, 1, timeout_ms) > 0 && (pfd.revents & POLLIN)) {
		static char reply[8192];
		size_t reply_len = sizeof(reply) - 1;
		if (wpa_ctrl_recv(ctrl_cmd, reply, &reply_len) == 0) {
			jl_wifi_ctrl_apply_reply(pending_req, reply, reply_len);
		}
	}
	pending_req = JL_WIFI_CTRL_REQ_NONE;
}

static bool jl_wifi_ctrl_cmd_ok_timeout(const char* cmd, int timeout_ms)
{
	if (!jl_wifi_ctrl_ensure_connected()) {
		return false;
	}

	jl_wifi_ctrl_drain_pending(timeout_ms);

	char reply[64];
	size_t reply_len = sizeof(reply) - 1;
	if (!jl_wifi_ctrl_txn_timeout(ctrl_cmd, cmd, reply, &reply_len, timeout_ms)) {
		log_warn("jl_wifi_ctrl: '%s' timed out after %dms", cmd, timeout_ms);
		return false;
	}

	reply[reply_len] = '\0';
	bool ok = strncmp(reply, "OK", 2) == 0;
	if (!ok) {
		log_warn("jl_wifi_ctrl: '%s' -> '%s'", cmd, reply);
	}
	return ok;
}

static bool jl_wifi_ctrl_cmd_ok(const char* cmd)
{
	return jl_wifi_ctrl_cmd_ok_timeout(cmd, JL_WIFI_CTRL_TXN_TIMEOUT_MS);
}

static bool jl_wifi_ctrl_action_ok(const char* cmd)
{
	return jl_wifi_ctrl_cmd_ok_timeout(cmd, JL_WIFI_CTRL_ACTION_TIMEOUT_MS);
}

static const char* jl_wifi_ctrl_find_field(const char* buf, const char* key)
{
	size_t keylen = strlen(key);
	if (strncmp(buf, key, keylen) == 0) {
		return buf + keylen;
	}

	const char* p = buf;
	while ((p = strstr(p, key)) != NULL) {
		if (p[-1] == '\n') {
			return p + keylen;
		}
		p++;
	}

	return NULL;
}

static void jl_wifi_ctrl_copy_field(const char* val, char* out, size_t outsz)
{
	const char* nl = strchr(val, '\n');
	size_t len = nl ? (size_t)(nl - val) : strlen(val);
	if (len >= outsz) {
		len = outsz - 1;
	}
	memcpy(out, val, len);
	out[len] = '\0';
}

static jl_wifi_ctrl_state_t jl_wifi_ctrl_map_state(const char* s)
{
	if (strcmp(s, "COMPLETED") == 0) return JL_WIFI_CTRL_STATE_COMPLETED;
	if (strcmp(s, "DISCONNECTED") == 0) return JL_WIFI_CTRL_STATE_DISCONNECTED;
	if (strcmp(s, "SCANNING") == 0) return JL_WIFI_CTRL_STATE_SCANNING;
	if (strcmp(s, "ASSOCIATED") == 0) return JL_WIFI_CTRL_STATE_ASSOCIATED;
	if (strcmp(s, "4WAY_HANDSHAKE") == 0) return JL_WIFI_CTRL_STATE_4WAY_HANDSHAKE;
	if (strcmp(s, "INACTIVE") == 0) return JL_WIFI_CTRL_STATE_INACTIVE;
	return JL_WIFI_CTRL_STATE_UNKNOWN;
}

static void jl_wifi_ctrl_parse_status(const char* reply, jl_wifi_ctrl_status_t* out)
{
	const char* val;

	if ((val = jl_wifi_ctrl_find_field(reply, "ssid=")) != NULL) {
		char tmp[JL_WIFI_CTRL_SSID_RAW_MAX];
		jl_wifi_ctrl_copy_field(val, tmp, sizeof(tmp));
		jl_wifi_ssid_decode(out->ssid, sizeof(out->ssid), tmp);
	}

	if ((val = jl_wifi_ctrl_find_field(reply, "id=")) != NULL) {
		out->id_bss = (int)strtol(val, NULL, 10);
	}

	if ((val = jl_wifi_ctrl_find_field(reply, "wpa_state=")) != NULL) {
		char tmp[32];
		jl_wifi_ctrl_copy_field(val, tmp, sizeof(tmp));
		out->state = jl_wifi_ctrl_map_state(tmp);
	}

	if ((val = jl_wifi_ctrl_find_field(reply, "ip_address=")) != NULL) {
		jl_wifi_ctrl_copy_field(val, out->ip_address, sizeof(out->ip_address));
	}

	if ((val = jl_wifi_ctrl_find_field(reply, "address=")) != NULL) {
		jl_wifi_ctrl_copy_field(val, out->mac_address, sizeof(out->mac_address));
	}
}

static int jl_wifi_ctrl_parse_scan_results(char* reply, jl_wifi_ctrl_scan_result_t* out, int max)
{
	int count = 0;
	bool header_skipped = false;
	char* line_save = NULL;
	char* line = strtok_r(reply, "\n", &line_save);

	while (line && count < max) {
		if (!header_skipped) {
			header_skipped = true;
			line = strtok_r(NULL, "\n", &line_save);
			continue;
		}

		jl_wifi_ctrl_scan_result_t* ap = &out[count];
		memset(ap, 0, sizeof(*ap));

		char* field_save = NULL;
		char* field = strtok_r(line, "\t", &field_save);
		int step = 0;
		while (field) {
			switch (step) {
			case 0:
				snprintf(ap->bssid, sizeof(ap->bssid), "%s", field);
				break;
			case 1:
				ap->freq = atoi(field);
				break;
			case 2:
				ap->level = strtol(field, NULL, 10);
				break;
			case 3:
				ap->wpa2 = true;
				break;
			case 4:
				jl_wifi_ssid_decode(ap->ssid, sizeof(ap->ssid), field);
				break;
			}
			field = strtok_r(NULL, "\t", &field_save);
			step++;
		}

		if (step == 5) {
			count++;
		}

		line = strtok_r(NULL, "\n", &line_save);
	}

	return count;
}

/* Parses a completed async reply into the cache matching the request that
 * produced it. Shared by jl_wifi_ctrl_poll() (normal, non-blocking path) and
 * jl_wifi_ctrl_drain_pending() (the rare fallback used by synchronous action
 * commands). */
static void jl_wifi_ctrl_apply_reply(jl_wifi_ctrl_req_t req, char* reply, size_t reply_len)
{
	reply[reply_len] = '\0';

	if (req == JL_WIFI_CTRL_REQ_STATUS) {
		jl_wifi_ctrl_status_t st = {0};
		st.state = JL_WIFI_CTRL_STATE_UNKNOWN;
		jl_wifi_ctrl_parse_status(reply, &st);
		st.valid = true;
		cached_status = st;
		cached_status_valid = true;
	} else if (req == JL_WIFI_CTRL_REQ_SCAN_RESULTS) {
		cached_scan_count = jl_wifi_ctrl_parse_scan_results(reply, cached_scan, JL_WIFI_CTRL_SCAN_CACHE_MAX);
	}
}

/* Fires a STATUS/SCAN_RESULTS request without waiting for the reply. A
 * no-op if one is already in flight - the caller just keeps seeing the last
 * cached result until it resolves. */
static void jl_wifi_ctrl_request(jl_wifi_ctrl_req_t req, const char* cmd)
{
	if (pending_req != JL_WIFI_CTRL_REQ_NONE) {
		return;
	}

	int fd = wpa_ctrl_get_fd(ctrl_cmd);
	if (fd < 0) {
		return;
	}

	if (send(fd, cmd, strlen(cmd), 0) < 0) {
		log_debug("jl_wifi_ctrl: send failed, reconnecting");
		jl_wifi_ctrl_teardown();
		return;
	}

	pending_req = req;
	pending_req_sent_ms = custom_tick_get();
}

static void jl_wifi_ctrl_handle_event(const char* msg)
{
	const char* p = msg;
	if (*p == '<') {
		const char* gt = strchr(p, '>');
		if (gt) {
			p = gt + 1;
		}
	}

	if (strncmp(p, WPA_EVENT_SCAN_RESULTS, strlen(WPA_EVENT_SCAN_RESULTS)) == 0) {
		pending_events |= JL_WIFI_EVT_SCAN_RESULTS;
	} else if (strncmp(p, WPA_EVENT_CONNECTED, strlen(WPA_EVENT_CONNECTED)) == 0) {
		pending_events |= JL_WIFI_EVT_CONNECTED;
	} else if (strncmp(p, WPA_EVENT_DISCONNECTED, strlen(WPA_EVENT_DISCONNECTED)) == 0) {
		pending_events |= JL_WIFI_EVT_DISCONNECTED;
	} else if (strncmp(p, WPA_EVENT_TERMINATING, strlen(WPA_EVENT_TERMINATING)) == 0) {
		pending_events |= JL_WIFI_EVT_TERMINATING;
	} else if (strncmp(p, WPA_EVENT_TEMP_DISABLED, strlen(WPA_EVENT_TEMP_DISABLED)) == 0) {
		pending_events |= JL_WIFI_EVT_AUTH_FAILED;
	}
}

void jl_wifi_ctrl_poll(void)
{
	if (!jl_wifi_ctrl_ensure_connected()) {
		return;
	}

	if (pending_req != JL_WIFI_CTRL_REQ_NONE) {
		int cmd_fd = wpa_ctrl_get_fd(ctrl_cmd);
		struct pollfd pfd_cmd = { .fd = cmd_fd, .events = POLLIN };
		if (poll(&pfd_cmd, 1, 0) > 0 && (pfd_cmd.revents & POLLIN)) {
			static char reply[8192];
			size_t reply_len = sizeof(reply) - 1;
			if (wpa_ctrl_recv(ctrl_cmd, reply, &reply_len) == 0) {
				jl_wifi_ctrl_apply_reply(pending_req, reply, reply_len);
			}
			pending_req = JL_WIFI_CTRL_REQ_NONE;
		} else if ((uint32_t)(custom_tick_get() - pending_req_sent_ms) > JL_WIFI_CTRL_REQ_TIMEOUT_MS) {
			log_warn("jl_wifi_ctrl: async request timed out after %dms", JL_WIFI_CTRL_REQ_TIMEOUT_MS);
			pending_req = JL_WIFI_CTRL_REQ_NONE;
		}
	}

	struct pollfd pfd_mon = { .fd = wpa_ctrl_get_fd(ctrl_mon), .events = POLLIN };
	while (poll(&pfd_mon, 1, 0) > 0 && (pfd_mon.revents & POLLIN)) {
		static char buf[4096];
		size_t len = sizeof(buf) - 1;

		if (wpa_ctrl_recv(ctrl_mon, buf, &len) < 0) {
			log_debug("jl_wifi_ctrl: monitor recv failed, reconnecting");
			jl_wifi_ctrl_teardown();
			return;
		}

		buf[len] = '\0';
		jl_wifi_ctrl_handle_event(buf);
	}
}

uint32_t jl_wifi_ctrl_events_get_and_clear(void)
{
	uint32_t events = pending_events;
	pending_events = 0;
	return events;
}

bool jl_wifi_ctrl_status_get(jl_wifi_ctrl_status_t* out)
{
	if (!jl_wifi_ctrl_ensure_connected()) {
		memset(out, 0, sizeof(*out));
		out->state = JL_WIFI_CTRL_STATE_UNKNOWN;
		return false;
	}

	jl_wifi_ctrl_request(JL_WIFI_CTRL_REQ_STATUS, "STATUS");

	if (!cached_status_valid) {
		memset(out, 0, sizeof(*out));
		out->state = JL_WIFI_CTRL_STATE_UNKNOWN;
		return false;
	}

	*out = cached_status;
	return true;
}

bool jl_wifi_ctrl_enable_network(int id)
{
	char cmd[32];
	snprintf(cmd, sizeof(cmd), "ENABLE_NETWORK %d", id);
	return jl_wifi_ctrl_action_ok(cmd);
}

bool jl_wifi_ctrl_disable_network(int id)
{
	char cmd[32];
	snprintf(cmd, sizeof(cmd), "DISABLE_NETWORK %d", id);
	return jl_wifi_ctrl_action_ok(cmd);
}

bool jl_wifi_ctrl_reconnect(void)
{
	return jl_wifi_ctrl_action_ok("RECONNECT");
}

bool jl_wifi_ctrl_disconnect(void)
{
	return jl_wifi_ctrl_action_ok("DISCONNECT");
}

static bool jl_wifi_ctrl_txn_reply(const char* cmd, char* reply, size_t reply_sz, int timeout_ms)
{
	if (!jl_wifi_ctrl_ensure_connected()) {
		return false;
	}

	jl_wifi_ctrl_drain_pending(timeout_ms);

	size_t reply_len = reply_sz - 1;
	if (!jl_wifi_ctrl_txn_timeout(ctrl_cmd, cmd, reply, &reply_len, timeout_ms)) {
		log_warn("jl_wifi_ctrl: '%s' timed out after %dms", cmd, timeout_ms);
		return false;
	}
	reply[reply_len] = '\0';
	return true;
}

int jl_wifi_ctrl_add_network(void)
{
	char reply[32];
	if (!jl_wifi_ctrl_txn_reply("ADD_NETWORK", reply, sizeof(reply),
				   JL_WIFI_CTRL_ACTION_TIMEOUT_MS)) {
		return -1;
	}
	if (strncmp(reply, "FAIL", 4) == 0) {
		log_warn("jl_wifi_ctrl: ADD_NETWORK -> '%s'", reply);
		return -1;
	}
	int id = (int)strtol(reply, NULL, 10);
	log_info("jl_wifi_ctrl: ADD_NETWORK -> %d", id);
	return id;
}

static bool jl_wifi_ctrl_list_has_id0(void)
{
	char reply[512];
	if (!jl_wifi_ctrl_txn_reply("LIST_NETWORKS", reply, sizeof(reply),
				   JL_WIFI_CTRL_ACTION_TIMEOUT_MS)) {
		return false;
	}

	char* line = strchr(reply, '\n');
	if (!line) {
		return false;
	}
	for (line++; line && *line; ) {
		if (line[0] == '0' && (line[1] == '\t' || line[1] == ' ' ||
				       line[1] == '\n' || line[1] == '\0')) {
			return true;
		}
		char* nl = strchr(line, '\n');
		line = nl ? nl + 1 : NULL;
	}
	return false;
}

bool jl_wifi_ctrl_ensure_network0(void)
{
	if (jl_wifi_ctrl_list_has_id0()) {
		return true;
	}
	if (jl_wifi_ctrl_add_network() != 0) {
		return false;
	}
	jl_wifi_ctrl_save_config();
	return true;
}

bool jl_wifi_ctrl_set_network_ssid_hex(int id, const char* ssid_hex)
{
	char cmd[JL_WIFI_CTRL_SSID_RAW_MAX + 32];
	int n = snprintf(cmd, sizeof(cmd), "SET_NETWORK %d ssid %s", id, ssid_hex);
	if (n < 0 || (size_t)n >= sizeof(cmd)) {
		return false;
	}
	if (jl_wifi_ctrl_action_ok(cmd)) {
		return true;
	}
	if (id != 0) {
		return false;
	}
	/* Empty factory network={} can occupy id 0 and still reject SET. */
	jl_wifi_ctrl_action_ok("REMOVE_NETWORK 0");
	if (jl_wifi_ctrl_add_network() != 0) {
		return false;
	}
	return jl_wifi_ctrl_action_ok(cmd);
}

bool jl_wifi_ctrl_set_network_psk(int id, const char* psk_plaintext)
{
	char cmd[128];
	int n = snprintf(cmd, sizeof(cmd), "SET_NETWORK %d psk \"%s\"", id, psk_plaintext);
	if (n < 0 || (size_t)n >= sizeof(cmd)) {
		return false;
	}
	return jl_wifi_ctrl_action_ok(cmd);
}

bool jl_wifi_ctrl_set_network_scan_ssid(int id, bool enable)
{
	char cmd[32];
	snprintf(cmd, sizeof(cmd), "SET_NETWORK %d scan_ssid %d", id, enable ? 1 : 0);
	return jl_wifi_ctrl_action_ok(cmd);
}

bool jl_wifi_ctrl_save_config(void)
{
	return jl_wifi_ctrl_action_ok("SAVE_CONFIG");
}

bool jl_wifi_ctrl_scan(void)
{
	return jl_wifi_ctrl_cmd_ok("SCAN");
}

int jl_wifi_ctrl_scan_results_get(jl_wifi_ctrl_scan_result_t* out, int max)
{
	if (!jl_wifi_ctrl_ensure_connected()) {
		return 0;
	}

	jl_wifi_ctrl_request(JL_WIFI_CTRL_REQ_SCAN_RESULTS, "SCAN_RESULTS");

	int n = cached_scan_count < max ? cached_scan_count : max;
	if (n > 0) {
		memcpy(out, cached_scan, n * sizeof(*out));
	}
	return n;
}