#include "esp_netif_ip_addr.h"
#include "esp_wifi_types_generic.h"
#include "freertos/idf_additions.h"
#include "ssd1306.h"

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "portmacro.h"
#include "driver/gpio.h"
#include "hal/gpio_types.h"
#include "esp_wifi.h"
#include "esp_log.h"
#include "nvs_flash.h"

#include <string.h>
#include <stdio.h>

#define TICKS_MS(x) ((x) / portTICK_PERIOD_MS)
#define LOG(...) ESP_LOGI("desk-andon", __VA_ARGS__)

/* The event group allows multiple bits for each event, but we only care about two events:
 * - we are connected to the AP with an IP
 * - we failed to connect after the maximum amount of retries */
#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT      BIT1

#define WIFI_ATTEMPT_MAX 5

#define WIFI_SSID "andon-net"
#define WIFI_PASS "uzw00ystu1"

enum Lane {
	RED,
	ORANGE,
	GREEN,
	BLUE,
	WHITE,
	SPEAKER,
};

static const uint8_t bitmap_frame[] = {
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xc0, 0x00, 0x00, 0x20, 0x00, 0x00, 0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x04, 0x00, 0x00, 0x03, 
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
};

static const uint8_t bitmap_light1[] = {
0x07, 0x00, 0x0e, 0x0e, 0x00, 0x1c, 0x1c, 0x00, 0x38, 0x38, 0x00, 0x70, 0x70, 0x00, 0xe0, 0xe0, 
0x01, 0xc0, 0xc0, 0x03, 0x80, 0x80, 0x07, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x38, 
0x00, 0x00, 0x70, 0x00, 0x00, 0xe0, 0x01, 0x01, 0xc0, 0x03, 0x03, 0x80, 0x07, 0x07, 0x00, 0x0e, 
0x0e, 0x00, 0x1c, 0x1c, 0x00, 0x38, 0x38, 0x00, 0x70, 0x70, 0x00, 0xe0, 0xe0, 0x01, 0xc0, 0xc0, 
0x03, 0x80
};

static const uint8_t bitmap_light2[] = {
0x00, 0x1c, 0x00, 0x00, 0x38, 0x00, 0x00, 0x70, 0x00, 0x00, 0xe0, 0x00, 0x01, 0xc0, 0x00, 0x03, 
0x80, 0x00, 0x07, 0x00, 0x01, 0x0e, 0x00, 0x03, 0x1c, 0x00, 0x07, 0x38, 0x00, 0x0e, 0x70, 0x00, 
0x1c, 0xe0, 0x00, 0x38, 0xc0, 0x00, 0x70, 0x80, 0x00, 0xe0, 0x00, 0x01, 0xc0, 0x00, 0x03, 0x80, 
0x00, 0x07, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x38, 0x00, 0x00, 0x70, 0x00, 0x00, 
0xe0, 0x00
};

const static int ANDON_PINS[SPEAKER+1] = { 10, 6, 3, 5, 7, 4 };

static uint8_t state = 0;

static SSD1306_t display = { 0 };
static bool lcd_init = false;

static char mac_str[16] = { 0 };
static char ip_str[17] = { 0 };
static bool wifi_connected = false;
static int wifi_connect_attempt = 0;

static EventGroupHandle_t wifi_event_group;

static void
lcd_task_main(void *_p)
{
	bool frame_state[SPEAKER] = { 0 };
	char mac_str_old[16] = { 0 };
	char mac_str_new[16] = { 0 };
	char ip_str_old[17] = { 0 };
	char ip_str_new[17] = { 0 };

	TickType_t last = 0;
	while (1) {
		TickType_t now = xTaskGetTickCount();
		if (!last || (now - last) * portTICK_PERIOD_MS >= 500) {
			uint8_t current = state;
			memcpy(ip_str_new, ip_str, sizeof(ip_str));
			memcpy(mac_str_new, mac_str, sizeof(mac_str));
			bool connected = wifi_connected;

			for (int i = 0; i < SPEAKER; i++) {
				if (current & (1 << i)) {
					frame_state[i] ^= true;
				}
			}
			last = now;

			if (!lcd_init) {
				_ssd1306_clear(&display, false);
				_ssd1306_bitmaps(&display, 0, 0, bitmap_frame, 128, 28, false);
				lcd_init = true;
			}

			const int xoff = (128 - 24 * 5 - 4) / 2;
			for (int i = 0; i < SPEAKER; i++) {
				if (current & (1 << i)) { 
					_ssd1306_bitmaps(&display, xoff + i * 25, 3,
						frame_state[i] ? bitmap_light2 : bitmap_light1,
						24, 22, false);
				} else {
					_ssd1306_clear_rect(&display, xoff + i * 25, 3, 24, 22, false);
				}
			}


			if (strcmp(mac_str_old, mac_str_new)) {
				_ssd1306_display_text(&display, 4, 40,
					mac_str_new, strlen(mac_str_new), false);
				strncpy(mac_str_old, mac_str_new, sizeof(mac_str));
			}

			if (strcmp(ip_str_old, ip_str_new)) {
				_ssd1306_display_text(&display, connected ? 4 : 0, 52,
					ip_str_new, strlen(ip_str_new), false);
				strncpy(ip_str_old, ip_str_new, sizeof(ip_str));
			}

			ssd1306_next_frame(&display);
		}

		vTaskDelay(TICKS_MS(50));
	}
}

static void
set_output(uint8_t bv)
{
	state = bv;
	for (int i = 0; i <= SPEAKER; i++) {
		gpio_set_level(ANDON_PINS[i], !(state & (1 << i)));
		vTaskDelay(10);
	}
}

static void
wifi_event_handler(void *arg, esp_event_base_t event_base,
	int32_t event_id, void *event_data)
{
	if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
		snprintf(ip_str, sizeof(ip_str), "  connecting..  ");
		vTaskDelay(1);
		wifi_connected = false;
		esp_wifi_connect();
	} else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
		wifi_connected = false;
		if (wifi_connect_attempt < WIFI_ATTEMPT_MAX) {
			esp_wifi_connect();
			wifi_connect_attempt++;
			LOG("retry to connect to the AP");
			snprintf(ip_str, sizeof(ip_str), "  connecting..%c ",
				'0' + wifi_connect_attempt);
		} else {
			snprintf(ip_str, sizeof(ip_str), "    no wifi    ");
			xEventGroupSetBits(wifi_event_group, WIFI_FAIL_BIT);
		}
		LOG("connect to the AP fail");
	} else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
		ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
		LOG("got ip:" IPSTR, IP2STR(&event->ip_info.ip));
		snprintf(ip_str, sizeof(ip_str), IPSTR, IP2STR(&event->ip_info.ip));
		wifi_connect_attempt = 0;
		wifi_connected = true;
		xEventGroupSetBits(wifi_event_group, WIFI_CONNECTED_BIT);
	}
}

static void
wifi_init(void)
{
	esp_err_t ret = nvs_flash_init();
	if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
		ESP_ERROR_CHECK(nvs_flash_erase());
		ret = nvs_flash_init();
	}
	ESP_ERROR_CHECK(ret);

	wifi_event_group = xEventGroupCreate();
	ESP_ERROR_CHECK(esp_netif_init());

	ESP_ERROR_CHECK(esp_event_loop_create_default());
	esp_netif_create_default_wifi_sta();

	wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
	ESP_ERROR_CHECK(esp_wifi_init(&cfg));

	uint8_t mac[6] = { 0 };
	esp_wifi_get_mac(WIFI_IF_STA, mac);
	snprintf(mac_str, sizeof(mac_str), ":%02X:%02X:%02X:%02X:%02X",
		mac[1], mac[2], mac[3], mac[4], mac[5]);

	esp_event_handler_instance_t instance_any_id;
	esp_event_handler_instance_t instance_got_ip;
	ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
		ESP_EVENT_ANY_ID, &wifi_event_handler, NULL, &instance_any_id));
	ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
		IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL, &instance_got_ip));

	wifi_config_t wifi_config = {
		.sta = {
			.ssid = WIFI_SSID,
			.password = WIFI_PASS,
			.threshold.authmode = WIFI_AUTH_WPA2_PSK,
			.sae_pwe_h2e = WPA3_SAE_PWE_BOTH,
			.sae_h2e_identifier = "",
		},
	};
	ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA) );
	ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config) );
	ESP_ERROR_CHECK(esp_wifi_start() );

	LOG("wifi_init_sta finished.");

	/* Waiting until either the connection is established (WIFI_CONNECTED_BIT)
	 * or connection failed for the maximum number of re-tries (WIFI_FAIL_BIT).
	 * The bits are set by event_handler() (see above) */
	EventBits_t bits = xEventGroupWaitBits(wifi_event_group,
	    WIFI_CONNECTED_BIT | WIFI_FAIL_BIT, pdFALSE, pdFALSE, portMAX_DELAY);

	/* xEventGroupWaitBits() returns the bits before the call returned,
	 * hence we can test which event actually happened. */
	if (bits & WIFI_CONNECTED_BIT) {
		LOG("connected to ap SSID:%s password:%s", WIFI_SSID, WIFI_PASS);
	} else if (bits & WIFI_FAIL_BIT) {
		LOG("Failed to connect to SSID:%s, password:%s", WIFI_SSID, WIFI_PASS);
	} else {
		LOG("UNEXPECTED EVENT");
	}
}

void
app_main(void)
{
	for (int i = 0; i <= SPEAKER; i++) {
		gpio_reset_pin(ANDON_PINS[i]);
		gpio_set_direction(ANDON_PINS[i], GPIO_MODE_OUTPUT_OD);
	 	gpio_set_level(ANDON_PINS[i], 1);
	}

	printf("DESK ANDON\n");
	i2c_master_init(&display, 0, 1, -1);
	ssd1306_init(&display, 128, 64);
	ssd1306_clear_screen(&display, false);
	ssd1306_display_text(&display, 3, 24, "DESK ANDON", 10, false);
	vTaskDelay(TICKS_MS(500));

	TaskHandle_t lcd_task = NULL;
	xTaskCreate(lcd_task_main, "lcd", 16 * 1024, NULL, tskIDLE_PRIORITY, &lcd_task);
	vTaskDelay(TICKS_MS(500));

	wifi_init();

	int tick = 0;
	uint8_t b = 0;
	while (1) {
		printf("Tick %i\n", ++tick);
		set_output(b);
		b = (b + 1) & ((1 << SPEAKER) - 1);
		vTaskDelay(TICKS_MS(1000));
	}
}