#include <string.h>
#include <stdlib.h>

#include "pico/stdlib.h"
#include "pico/cyw43_arch.h"

#include "lwip/netif.h"
#include "lwip/tcp.h"
#include "lwip/ip4_addr.h"

#include "wifi-settings.h"

#define MAX_PAYLOAD 256
#define TCP_PORT 5505

extern struct netif *netif_default;

enum {
	READ_CMD,
	READ_LEN0,
	READ_LEN1,
	READ_PAYLOAD
};

struct conn_state {
	uint8_t state;
	uint8_t cmd;
	uint16_t len;
	uint16_t received;
	uint8_t buf[MAX_PAYLOAD];
};

static void handle_frame(struct tcp_pcb *pcb,
                         uint8_t cmd,
                         uint8_t *payload,
                         uint16_t len)
{
	(void)cmd;
	(void)payload;
	(void)len;

	const char reply_payload[] = "Hello";
	const uint16_t reply_len = 5;

	uint8_t frame[3 + reply_len];

	frame[0] = 0x01;                 // ACK command
	frame[1] = (uint8_t)(reply_len & 0xFF);
	frame[2] = (uint8_t)(reply_len >> 8);
	memcpy(&frame[3], reply_payload, reply_len);

	tcp_write(pcb, frame, sizeof(frame), TCP_WRITE_FLAG_COPY);
	tcp_output(pcb);
}

static err_t on_recv(void *arg, struct tcp_pcb *pcb, struct pbuf *p, err_t err) {
	(void)err;

	struct conn_state *st = (struct conn_state *)arg;

	if (!p) {
		free(st);
		tcp_close(pcb);
		return ERR_OK;
	}

	struct pbuf *q = p;

	while (q) {
		uint8_t *data = (uint8_t *)q->payload;
		uint16_t i = 0;

		while (i < q->len) {
			uint8_t b = data[i++];

			if (st->state == READ_CMD) {
				st->cmd = b;
				st->len = 0;
				st->received = 0;
				st->state = READ_LEN0;
			}
			else if (st->state == READ_LEN0) {
				st->len = (uint16_t)b;
				st->state = READ_LEN1;
			}
			else if (st->state == READ_LEN1) {
				st->len |= ((uint16_t)b << 8);

				if (st->len > MAX_PAYLOAD) {
					pbuf_free(p);
					free(st);
					tcp_abort(pcb);
					return ERR_ABRT;
				}

				st->received = 0;
				st->state = (st->len == 0) ? READ_CMD : READ_PAYLOAD;

				if (st->len == 0) {
					handle_frame(pcb, st->cmd, st->buf, 0);
				}
			}
			else { /* READ_PAYLOAD */
				st->buf[st->received++] = b;

				if (st->received == st->len) {
					handle_frame(pcb, st->cmd, st->buf, st->len);
					st->state = READ_CMD;
				}
			}
		}

		q = q->next;
	}

	tcp_recved(pcb, p->tot_len);
	pbuf_free(p);
	return ERR_OK;
}

static err_t on_accept(void *arg, struct tcp_pcb *newpcb, err_t err) {
	(void)arg;
	(void)err;

	struct conn_state *st = (struct conn_state *)calloc(1, sizeof(struct conn_state));
	if (!st) {
		tcp_abort(newpcb);
		return ERR_ABRT;
	}

	st->state = READ_CMD;

	tcp_arg(newpcb, st);
	tcp_recv(newpcb, on_recv);

	return ERR_OK;
}

int main() {
	stdio_init_all();

	if (cyw43_arch_init()) {
		return -1;
	}

	cyw43_arch_enable_sta_mode();

	if (cyw43_arch_wifi_connect_timeout_ms(
			WIFI_SSID,
			WIFI_PASS,
			CYW43_AUTH_WPA2_AES_PSK,
			30000)) {
		return -1;
	}

	struct tcp_pcb *listen_pcb = tcp_new_ip_type(IPADDR_TYPE_V4);
	if (!listen_pcb) {
		return -1;
	}

	if (tcp_bind(listen_pcb, IP_ADDR_ANY, TCP_PORT) != ERR_OK) {
		tcp_close(listen_pcb);
		return -1;
	}

	listen_pcb = tcp_listen(listen_pcb);
	tcp_accept(listen_pcb, on_accept);

	bool dhcp_done = false;

	while (true) {
		cyw43_arch_poll();

		if (!dhcp_done) {
			if (netif_default &&
				netif_is_up(netif_default) &&
				!ip4_addr_isany_val(*netif_ip4_addr(netif_default))) {
				dhcp_done = true;
				cyw43_arch_gpio_put(CYW43_WL_GPIO_LED_PIN, 1);
			}
		}

		sleep_ms(1);
	}
}