/*
 * SPDX-License-Identifier: AGPL-3.0-only
 * Copyright (c) Camden Dixie O'Brien
 */

#include "wifi.h"

#include "console.h"
#include "settings.h"

#include "esp_event.h"
#include "esp_log.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "freertos/event_groups.h"
#include <string.h>

#define TAG "Wifi"

static wifi_config_t config;
static WifiStatus status;
static unsigned retries;
static esp_netif_ip_info_t ip_info;
static esp_netif_t *net_if;

static void wifi_event_handler(
    void *arg, esp_event_base_t event_base, int32_t event_id,
    void *event_data)
{
	(void)arg;
	(void)event_base;

	esp_err_t error;

	switch (event_id) {
	case WIFI_EVENT_STA_START:
		error = esp_wifi_connect();
		if (error != ESP_OK)
			ESP_LOGE(TAG, "Error connecting to WiFi: %04x", error);
		break;

	case WIFI_EVENT_STA_DISCONNECTED:
		if (retries < CONFIG_WIFI_MAX_RETRIES) {
			ESP_LOGI(TAG, "Retrying connection");
			error = esp_wifi_connect();
			if (error != ESP_OK)
				ESP_LOGE(TAG, "Error connecting to WiFi: %04x", error);
			++retries;
		} else {
			ESP_LOGW(
			    TAG, "Failed to connect to network %s", config.sta.ssid);
		}
		break;

	default:
		ESP_LOGI(TAG, "Unhandled WiFi event: %ld", event_id);
		break;
	}
}

static void got_ip_event_handler(
    void *arg, esp_event_base_t event_base, int32_t event_id,
    void *event_data)
{
	(void)arg;
	(void)event_base;

	if (event_id != IP_EVENT_STA_GOT_IP) {
		ESP_LOGW(TAG, "Unexpected IP event received: %ld", event_id);
		return;
	}

	ip_event_got_ip_t *got_ip_data = (ip_event_got_ip_t *)event_data;
	memcpy(&ip_info, &got_ip_data->ip_info, sizeof(ip_info));

	ESP_LOGI(TAG, "Got IP " IPSTR, IP2STR(&ip_info.ip));
	status = WIFI_STATUS_CONNECTED;
}

static void handle_hostname_update(const char *hostname)
{
	const esp_err_t error = esp_netif_set_hostname(net_if, hostname);
	if (error != ESP_OK)
		ESP_LOGE(TAG, "Failed to set hostname: %04x", error);

	wifi_reconnect();
}

static void handle_ssid_update(const char *ssid)
{
	ESP_LOGI(TAG, "SSID updated");

	if (strlen(ssid) > sizeof(config.sta.ssid) - 1)
		ESP_LOGW(TAG, "SSID too long (truncated)");
	strncpy((char *)config.sta.ssid, ssid, sizeof(config.sta.ssid));
	config.sta.ssid[sizeof(config.sta.ssid) - 1] = '\0';

	const esp_err_t error = esp_wifi_set_config(WIFI_IF_STA, &config);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error setting config: %04x", error);
		return;
	}

	wifi_reconnect();
}

static void handle_psk_update(const char *psk)
{
	ESP_LOGI(TAG, "PSK updated");

	if (strlen(psk) > sizeof(config.sta.password) - 1)
		ESP_LOGW(TAG, "PSK too long (truncated)");
	strncpy((char *)config.sta.password, psk, sizeof(config.sta.password));
	config.sta.password[sizeof(config.sta.password) - 1] = '\0';

	const esp_err_t error = esp_wifi_set_config(WIFI_IF_STA, &config);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error setting config: %04x", error);
		return;
	}

	wifi_reconnect();
}

static int command_func(int argc, char **argv)
{
	if (argc != 2) {
		printf("Invalid number of arguments\n");
		return 1;
	} else if (strcmp(argv[1], "status") == 0) {
		printf(
		    "%s\n",
		    status == WIFI_STATUS_CONNECTED ? "Connected" : "Disconnected");
		return 0;
	} else if (strcmp(argv[1], "reconnect") == 0) {
		wifi_reconnect();
		return 0;
	} else if (strcmp(argv[1], "ip") == 0) {
		printf(IPSTR "\n", IP2STR(&ip_info.ip));
		return 0;
	} else {
		printf("Subcommand not recognised\n");
		return 1;
	}
}

void wifi_init(void)
{
	esp_err_t error;

	status = WIFI_STATUS_DISCONNECTED;

	error = esp_netif_init();
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error initializing network stack: %04x", error);
		return;
	}

	error = esp_event_loop_create_default();
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error creating event loop: %04x", error);
		return;
	}
	net_if = esp_netif_create_default_wifi_sta();

	wifi_init_config_t init_config = WIFI_INIT_CONFIG_DEFAULT();
	error = esp_wifi_init(&init_config);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error initializing WiFi stack: %04x", error);
		return;
	}

	esp_event_handler_instance_t instance;
	error = esp_event_handler_instance_register(
	    WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, &instance);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error registering WiFi event handler: %04x", error);
		return;
	}
	error = esp_event_handler_instance_register(
	    IP_EVENT, IP_EVENT_STA_GOT_IP, &got_ip_event_handler, NULL,
	    &instance);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error registering IP event handler: %04x", error);
		return;
	}

	error = esp_wifi_set_mode(WIFI_MODE_STA);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error setting mode: %04x", error);
		return;
	}

	char hostname[SETTINGS_MAX_VALUE_SIZE];
	(void)settings_get_hostname(hostname, SETTINGS_MAX_VALUE_SIZE);
	error = esp_netif_set_hostname(net_if, hostname);
	if (error != ESP_OK)
		ESP_LOGE(TAG, "Failed to set hostname: %04x", error);

	memset(&config, 0, sizeof(config));
	if (settings_get_ssid((char *)config.sta.ssid, sizeof(config.sta.ssid))
	    > sizeof(config.sta.ssid))
		ESP_LOGW(TAG, "SSID too long (truncated)");
	if (settings_get_psk(
	        (char *)config.sta.password, sizeof(config.sta.password))
	    > sizeof(config.sta.password))
		ESP_LOGW(TAG, "PSK too long (truncated)");
	error = esp_wifi_set_config(WIFI_IF_STA, &config);
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error setting config: %04x", error);
		return;
	}

	error = esp_wifi_start();
	if (error != ESP_OK) {
		ESP_LOGE(TAG, "Error starting WiFi: %04x", error);
		return;
	}

	settings_add_hostname_callback(&handle_hostname_update);
	settings_add_ssid_callback(&handle_ssid_update);
	settings_add_psk_callback(&handle_psk_update);

	console_register(
	    "wifi", "Manage WiFi connection", "wifi <status|reconnect|ip>",
	    command_func);
}

void wifi_reconnect(void)
{
	esp_err_t error;

	error = esp_wifi_disconnect();
	if (error != ESP_OK)
		ESP_LOGE(TAG, "Error disconnecting from WiFi: %04x", error);
	error = esp_wifi_connect();
	if (error != ESP_OK)
		ESP_LOGE(TAG, "Error connecting to WiFi: %04x", error);

	retries = 0;
}

WifiStatus get_wifi_status(void)
{
	return status;
}