//
// Created by nano on 8/23/26.
//

#include "wireless_comms.h"

#include "esp_err.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_now.h"
#include "esp_wifi.h"
#include "esp_crc.h"
#include "esp_mac.h"
#include "esp_random.h"
#include "usb_serial.h"

const char *TAG = "wireless_comms";

#define ESPNOW_MAXDELAY 512

static QueueHandle_t espnow_send_queue = NULL;

static uint8_t s_broadcast_mac[ESP_NOW_ETH_ALEN] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
extern uint8_t node_mac[ESP_NOW_ETH_ALEN];

void wireless_comms::wifi_init(int channel) {
	ESP_ERROR_CHECK(esp_netif_init());
	ESP_ERROR_CHECK(esp_event_loop_create_default());
	wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
	ESP_ERROR_CHECK(esp_wifi_init(&cfg));
	ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));
	ESP_ERROR_CHECK(esp_wifi_set_mode(ESPNOW_WIFI_MODE));
	ESP_ERROR_CHECK(esp_wifi_set_mac(ESPNOW_WIFI_IF,node_mac));
	ESP_ERROR_CHECK(esp_wifi_start());
	ESP_ERROR_CHECK(esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE));


#ifdef CONFIG_ESPNOW_ENABLE_LONG_RANGE
	ESP_ERROR_CHECK(
		esp_wifi_set_protocol(ESPNOW_WIFI_IF, WIFI_PROTOCOL_11B|WIFI_PROTOCOL_11G|WIFI_PROTOCOL_11N|WIFI_PROTOCOL_LR));
#endif
}

void wireless_comms::espnow_send_callback(const esp_now_send_info_t *tx_info, esp_now_send_status_t status) {
	espnow_event_t evt;
	espnow_event_send_cb_t *send_cb = &evt.info.send_cb;

	if (tx_info == NULL) {
		ESP_LOGE(TAG, "Send cb arg error");
		return;
	}

	evt.id = ESPNOW_SEND_CB;
	memcpy(send_cb->mac_addr, tx_info->des_addr, ESP_NOW_ETH_ALEN);
	send_cb->status = status;
	if (xQueueSend(espnow_send_queue, &evt, ESPNOW_MAXDELAY) != pdTRUE) {
		ESP_LOGW(TAG, "Send send queue fail");
	}
}

void wireless_comms::espnow_recv_callback(const esp_now_recv_info_t *recv_info, const uint8_t *data, int len) {
	espnow_event_t evt;
	espnow_event_recv_cb_t *recv_cb = &evt.info.recv_cb;
	uint8_t *mac_addr = recv_info->src_addr;

	if (mac_addr == NULL || data == NULL || len <= 0) {
		ESP_LOGE(TAG, "Receive cb arg error");
		return;
	}

	evt.id = ESPNOW_RECV_CB;
	memcpy(recv_cb->mac_addr, mac_addr, ESP_NOW_ETH_ALEN);
	recv_cb->data = (uint8_t *) malloc(len);
	if (recv_cb->data == NULL) {
		ESP_LOGE(TAG, "Malloc receive data fail");
		return;
	}
	memcpy(recv_cb->data, data, len);
	recv_cb->data_len = len;
	if (xQueueSend(espnow_send_queue, &evt, ESPNOW_MAXDELAY) != pdTRUE) {
		ESP_LOGW(TAG, "Send receive queue fail");
		free(recv_cb->data);
	}
}

/* Parse received ESPNOW data. */
int wireless_comms::espnow_data_parse(uint8_t *data, uint16_t data_len) {
	auto *buf = (espnow_data_t *) data;
	uint16_t crc, crc_cal = 0;

	if (data_len < sizeof(espnow_data_t)) {
		ESP_LOGE(TAG, "Receive ESPNOW data too short, len:%d", data_len);
		return -1;
	}

	crc = buf->crc;
	buf->crc = 0;
	crc_cal = esp_crc16_le(UINT16_MAX, (uint8_t const *) buf, data_len);

	if (crc_cal == crc) {
		return 1;
	}

	return -1;
}

/* Prepare ESPNOW data to be sent. */
void wireless_comms::espnow_send_data_prepare(espnow_send_param_t *send_param) {
	espnow_data_t *buf = (espnow_data_t *) send_param->buffer;
	assert(send_param->len >= sizeof(espnow_data_t));
	buf->crc = 0;
	buf->crc = esp_crc16_le(UINT16_MAX, (uint8_t const *) buf, send_param->len);
	ESP_LOGI("espnow send_data_prepare", "%d %s", send_param->len - sizeof(espnow_data_t), buf->payload);
}

esp_err_t wireless_comms::espnow_init(int wifi_channel) {
	// Create send queue
	espnow_send_queue = xQueueCreate(ESPNOW_QUEUE_SIZE, sizeof(espnow_event_t));
	if (espnow_send_queue == NULL) {
		ESP_LOGE(TAG, "Create queue fail");
		return ESP_FAIL;
	}

	/* Initialize ESPNOW and register sending and receiving callback function. */
	ESP_ERROR_CHECK(esp_now_init());
	ESP_ERROR_CHECK(esp_now_register_send_cb(espnow_send_callback));
	ESP_ERROR_CHECK(esp_now_register_recv_cb(espnow_recv_callback));
#if CONFIG_ESPNOW_ENABLE_POWER_SAVE
	ESP_ERROR_CHECK(esp_now_set_wake_window(CONFIG_ESPNOW_WAKE_WINDOW));
	ESP_ERROR_CHECK(esp_wifi_connectionless_module_set_wake_interval(CONFIG_ESPNOW_WAKE_INTERVAL));
#endif
	/* Set primary master key. */
	ESP_ERROR_CHECK(esp_now_set_pmk((uint8_t *)ESPNOW_PMK));

	/* Add broadcast peer information to peer list. */
	auto *peer = static_cast<esp_now_peer_info_t *>(malloc(sizeof(esp_now_peer_info_t)));
	if (peer == NULL) {
		ESP_LOGE(TAG, "Malloc peer information fail");
		vQueueDelete(espnow_send_queue);
		espnow_send_queue = NULL;
		esp_now_deinit();
		return ESP_FAIL;
	}
	memset(peer, 0, sizeof(esp_now_peer_info_t));
	peer->channel = wifi_channel;
	peer->ifidx = ESPNOW_WIFI_IF;
	peer->encrypt = false;
	memcpy(peer->peer_addr, s_broadcast_mac, ESP_NOW_ETH_ALEN);
	ESP_ERROR_CHECK(esp_now_add_peer(peer));
	free(peer);

	/* Initialize sending parameters. */
	send_param = (espnow_send_param_t *) malloc(sizeof(espnow_send_param_t));
	if (send_param == NULL) {
		ESP_LOGE(TAG, "Malloc send parameter fail");
		vQueueDelete(espnow_send_queue);
		espnow_send_queue = NULL;
		esp_now_deinit();
		return ESP_FAIL;
	}
	memset(send_param, 0, sizeof(espnow_send_param_t));
	send_param->buffer = (uint8_t *) malloc(ESPNOW_SEND_LEN);
	send_param->len = ESPNOW_SEND_LEN;
	if (send_param->buffer == NULL) {
		ESP_LOGE(TAG, "Malloc send buffer fail");
		free(send_param);
		vQueueDelete(espnow_send_queue);
		espnow_send_queue = NULL;
		esp_now_deinit();
		return ESP_FAIL;
	}
	memcpy(send_param->dest_mac, s_broadcast_mac, ESP_NOW_ETH_ALEN);
	espnow_send_data_prepare(send_param);

	return ESP_OK;
}

void wireless_comms::espnow_deinit(espnow_send_param_t *send_param) {
	free(send_param->buffer);
	free(send_param);
	vQueueDelete(espnow_send_queue);
	espnow_send_queue = NULL;
	esp_now_deinit();
}

void wireless_comms::espnow_task(void *pvParameter) {
	espnow_event_t evt;
	int ret;
	auto task_params = (espnow_task_params_t *) pvParameter;
	espnow_send_param_t send_param = *task_params->send_param;
	uart_serial uart = task_params->uart;
	vTaskDelay(5000 / portTICK_PERIOD_MS);

	while (xQueueReceive(espnow_send_queue, &evt, portMAX_DELAY) == pdPASS) {
		switch (evt.id) {
			case ESPNOW_SEND_CB: {
				espnow_event_send_cb_t *send_cb = &evt.info.send_cb;
				/* Delay a while before sending the next data. */
				ESP_LOGI(TAG, "send data to " MACSTR "", MAC2STR(send_cb->mac_addr));
				break;
			}
			case ESPNOW_RECV_CB: {
				espnow_event_recv_cb_t *recv_cb = &evt.info.recv_cb;

				ret = espnow_data_parse(recv_cb->data, recv_cb->data_len);
				if (ret == 1) {
					auto *rx = (espnow_data_t *) recv_cb->data;
					ESP_LOGI(TAG, "Receive data from: " MACSTR ", len: %d, byte0: 0x%02x",
							 MAC2STR(recv_cb->mac_addr), recv_cb->data_len, rx->payload[0]);
					usb_serial::write(reinterpret_cast<const char *>(&rx->payload), PAYLOAD_SIZE);
					uart.write(reinterpret_cast<const char *>(&rx->payload), PAYLOAD_SIZE);
				} else {
					ESP_LOGI(TAG, "Receive error (bad CRC) from: " MACSTR "", MAC2STR(recv_cb->mac_addr));
				}
				free(recv_cb->data);
				break;
			}
			default:
				ESP_LOGE(TAG, "Callback type error: %d", evt.id);
				break;
		}
	}
	ESP_LOGI(TAG,"ENDED TASK WIRELESS");
}

void wireless_comms::set_dest_mac(const uint8_t dest_mac[ESP_NOW_ETH_ALEN]) {
	ESP_ERROR_CHECK(add_peer(dest_mac));
	memcpy(send_param->dest_mac, dest_mac, ESP_NOW_ETH_ALEN);
	espnow_send_data_prepare(send_param);
}

esp_err_t wireless_comms::add_peer(const uint8_t dest_mac[ESP_NOW_ETH_ALEN]) {
	auto *peer = static_cast<esp_now_peer_info_t *>(malloc(sizeof(esp_now_peer_info_t)));
	if (peer == NULL) {
		ESP_LOGE(TAG, "Malloc peer information fail");
		vQueueDelete(espnow_send_queue);
		espnow_send_queue = NULL;
		esp_now_deinit();
		return ESP_FAIL;
	}
	memset(peer, 0, sizeof(esp_now_peer_info_t));
	peer->channel = ESPNOW_CHANNEL;
	peer->ifidx = ESPNOW_WIFI_IF;
	peer->encrypt = false;
	memcpy(peer->peer_addr, dest_mac, ESP_NOW_ETH_ALEN);
	ESP_ERROR_CHECK(esp_now_add_peer(peer));
	free(peer);
	ESP_LOGI(TAG, "Peer " MACSTR " added", MAC2STR(dest_mac));
	return ESP_OK;
}

void wireless_comms::delete_peer(const uint8_t dest_mac[ESP_NOW_ETH_ALEN]) {
	ESP_ERROR_CHECK(esp_now_del_peer(dest_mac));
}

esp_err_t wireless_comms::send_byte(uint8_t data) {
	auto send_parameters = (espnow_send_param_t *) malloc(sizeof(espnow_send_param_t));
	memset(send_parameters, 0, sizeof(espnow_send_param_t));
	send_parameters->len = sizeof(espnow_data_t);
	send_parameters->buffer = (uint8_t *) malloc(sizeof(espnow_data_t));
	if (send_parameters->buffer == NULL) {
		ESP_LOGE(TAG, "Malloc send buffer fail");
		free(send_parameters);
		vQueueDelete(espnow_send_queue);
		espnow_send_queue = NULL;
		esp_now_deinit();
		return ESP_FAIL;
	}
	((espnow_data_t *) (send_parameters->buffer))->payload[0] = data; // i am disgusted at myself
	memset(((espnow_data_t *) (send_parameters->buffer))->payload + 1, 0,PAYLOAD_SIZE - 1);
	// fill the rest of the payload with zeroes

	memcpy(send_parameters->dest_mac, s_broadcast_mac, ESP_NOW_ETH_ALEN); // TODO: THIS IS IGNORED FOR NOW
	espnow_send_data_prepare(send_parameters);
	if (esp_now_send(NULL, send_parameters->buffer, send_parameters->len) != ESP_OK) {
		// send to all peers
		ESP_LOGE(TAG, "Send error");
		espnow_deinit(send_parameters);
	}
	free(send_parameters);
	ESP_LOGI(TAG, "SENT DATA");
	return ESP_OK;
}