/*  WiFi softAP Example

   This example code is in the Public Domain (or CC0 licensed, at your option.)

   Unless required by applicable law or agreed to in writing, this
   software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
   CONDITIONS OF ANY KIND, either express or implied.
*/
#include <string.h>
#include <sys/param.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event_loop.h"
#include "esp_log.h"
#include "nvs_flash.h"
//#include "esp_vfs.h"
//#include "esp_vfs_dev.h"
#include "driver/uart.h"

#include "lwip/err.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include <lwip/netdb.h>

/* The examples use WiFi configuration that you can set via 'make menuconfig'.

   If you'd rather not, just change the below entries to strings with
   the config you want - ie #define EXAMPLE_WIFI_SSID "mywifissid"
*/
#define EXAMPLE_ESP_WIFI_SSID   "BMTIMPLANT"
#define EXAMPLE_ESP_WIFI_PASS   "lsbmt123"
#define EXAMPLE_MAX_STA_CONN    3
#define PORT                    5005

#define EX_UART_NUM UART_NUM_1
#define UART_BUF_SIZE (1024)

/* FreeRTOS event group to signal when we are connected*/
static EventGroupHandle_t s_wifi_event_group;

static const char *TAG0 = "wifi softAP";
static const char *TAG1 = "uart_events";

static esp_err_t event_handler(void *ctx, system_event_t *event)
{
    switch(event->event_id) {
    case SYSTEM_EVENT_AP_STACONNECTED:
        ESP_LOGI(TAG0, "station:"MACSTR" join, AID=%d",
                 MAC2STR(event->event_info.sta_connected.mac),
                 event->event_info.sta_connected.aid);
        break;
    case SYSTEM_EVENT_AP_STADISCONNECTED:
        ESP_LOGI(TAG0, "station:"MACSTR"leave, AID=%d",
                 MAC2STR(event->event_info.sta_disconnected.mac),
                 event->event_info.sta_disconnected.aid);
        break;
    default:
        break;
    }
    return ESP_OK;
}

void wifi_init_softap()
{
    s_wifi_event_group = xEventGroupCreate();

    tcpip_adapter_init();
    ESP_ERROR_CHECK(esp_event_loop_init(event_handler, NULL));

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    wifi_config_t wifi_config = {
        .ap = {
            .ssid = EXAMPLE_ESP_WIFI_SSID,
            .ssid_len = strlen(EXAMPLE_ESP_WIFI_SSID),
            .password = EXAMPLE_ESP_WIFI_PASS,
            .max_connection = EXAMPLE_MAX_STA_CONN,
            .authmode = strlen(EXAMPLE_ESP_WIFI_PASS) ? WIFI_AUTH_WPA_WPA2_PSK : WIFI_AUTH_OPEN
        },
    };

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(ESP_IF_WIFI_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG0, "wifi_init_softap finished.SSID:%s password:%s", EXAMPLE_ESP_WIFI_SSID, EXAMPLE_ESP_WIFI_PASS);
}

int setup_server_socket()
{
    struct sockaddr_in destAddr;
    destAddr.sin_addr.s_addr= htonl(INADDR_ANY);
    destAddr.sin_family     = AF_INET;
    destAddr.sin_port       = htons(PORT);

    int tcp_server_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_IP);
    if (tcp_server_sock < 0) {
        ESP_LOGE(TAG0, "Unable to create socket: errno %d", errno);
        return 0;
    }
    ESP_LOGI(TAG0, "Socket created");

    int err = bind(tcp_server_sock, (struct sockaddr *)&destAddr, sizeof(destAddr));
    if (err != 0) {
        ESP_LOGE(TAG0, "Socket unable to bind: errno %d", errno);
        return 0;
    }
    ESP_LOGI(TAG0, "Socket bound");

    err = listen(tcp_server_sock, 1);
    if (err != 0) {
        ESP_LOGE(TAG0, "Error occured during listen: errno %d", errno);
        return 0;
    }
    ESP_LOGI(TAG0, "Socket listening");

    return tcp_server_sock;
}

void setup_uart()
{
    uart_config_t uart_config = {
        .baud_rate = 4000000,
        .data_bits = UART_DATA_8_BITS,
        .parity    = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE
    };

    uart_param_config(EX_UART_NUM, &uart_config);
    // Set UART pins(TX: IO16 (UART2 default), RX: IO17 (UART2 default), RTS: IO18, CTS: IO19)
    ESP_ERROR_CHECK(uart_set_pin(EX_UART_NUM, 16, 17, UART_PIN_NO_CHANGE, UART_PIN_NO_CHANGE));
    ESP_ERROR_CHECK(uart_driver_install(EX_UART_NUM, UART_BUF_SIZE * 2, UART_BUF_SIZE * 2, 0, NULL, 0));
}

bool handle_uart_rx(int tcp_client_sock)
{
    uint8_t buf[UART_BUF_SIZE];

    int len = uart_read_bytes(EX_UART_NUM, buf, sizeof(buf), 0);
    //ESP_LOGI(TAG1, "uart_rx: %d %d", len, tcp_client_sock);

    if(tcp_client_sock)
    {
        int err = send(tcp_client_sock, buf, len, 0);
        //ESP_LOGI(TAG1, "tcp sent");
        if (err < 0)
        {
            ESP_LOGE(TAG0, "Error occured during sending: errno %d", errno);
            return false;
        }
    }

    return true;
}

int handle_tcp_connect(int tcp_server_sock)
{
    struct sockaddr_in6 sourceAddr; // Large enough for both IPv4 or IPv6
    uint addrLen = sizeof(sourceAddr);
    int sock = accept(tcp_server_sock, (struct sockaddr *)&sourceAddr, &addrLen);
    if (sock < 0)
    {
        ESP_LOGE(TAG0, "Unable to accept connection: errno %d", errno);
        return 0;
    }
    ESP_LOGI(TAG0, "Socket accepted");

    return sock;
}

void close_tcp_client_sock(int sock)
{
    if (sock != -1)
    {
        //ESP_LOGE(TAG0, "Shutting down socket");
        //shutdown(sock, 2);
        //vTaskDelay(20 / portTICK_PERIOD_MS);
        ESP_LOGE(TAG0, "closing socket");
        close(sock);
        ESP_LOGE(TAG0, "now...");
    }
}

bool handle_tcp_recv(int sock)
{
    char rx_buffer[128];

    int len = recv(sock, rx_buffer, sizeof(rx_buffer), 0);
    // Error occured during receiving
    if (len < 0)
    {
        ESP_LOGE(TAG0, "recv failed: errno %d", errno);
        return false;
    }
    // Connection closed
    else if (len == 0)
    {
        ESP_LOGI(TAG0, "Connection closed");
        return false;
    }
    // Data received
    else
    {
        uart_write_bytes(EX_UART_NUM, rx_buffer, len);
    }
    return true;
}

int check_for_actions(int tcp_server_sock, int tcp_client_sock)
{
    fd_set rfds;

    FD_ZERO(&rfds);
    //FD_SET(uart_sock, &rfds);
    FD_SET(tcp_client_sock ? tcp_client_sock : tcp_server_sock, &rfds);//either accept a new client or serve the existing one

    struct timeval tv = {
        .tv_sec = 0,
        .tv_usec = 0,
    };

    int nfs = MAX(tcp_server_sock, tcp_client_sock);
    int ret = select(nfs + 1, &rfds, NULL, NULL, &tv);

    size_t size;
    ESP_ERROR_CHECK(uart_get_buffered_data_len(EX_UART_NUM, &size));

    if (ret < 0)
    {
        ESP_LOGE(TAG0, "Select failed: errno %d", errno);
        return tcp_client_sock;
    }

    if (FD_ISSET(tcp_server_sock, &rfds))
        tcp_client_sock = handle_tcp_connect(tcp_server_sock);
    else if ((size && !handle_uart_rx(tcp_client_sock)) || (FD_ISSET(tcp_client_sock, &rfds)    && !handle_tcp_recv(tcp_client_sock)))
    {
        close_tcp_client_sock(tcp_client_sock);
        return 0;
    }

    if (ret == 0 && size == 0)
    {
        vTaskDelay(10 / portTICK_PERIOD_MS);
    }

    return tcp_client_sock;
}

static void tcp_server_task(void *pvParameters)
{
    int tcp_client_sock = 0;
    int tcp_server_sock = setup_server_socket();
    setup_uart();

    while (tcp_server_sock)
    {
        tcp_client_sock = check_for_actions(tcp_server_sock, tcp_client_sock);
    }
    vTaskDelete(NULL);
}

void app_main()
{
    //Initialize NVS
    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);

    ESP_LOGI(TAG0, "ESP_WIFI_MODE_AP");
    wifi_init_softap();
    xTaskCreate(tcp_server_task, "tcp_server", 4096, NULL, 5, NULL);
}