#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "lwip/err.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include "lwip/netdb.h"
#include "esp_netif.h"
#include <errno.h>

#define WIFI_SSID      "Wokwi-GUEST"
#define WIFI_PASS      ""
#define UDP_SERVER_IP  "8.8.8.8"
#define UDP_PORT       53
#define TAG "UDP_TEST"

static bool wifi_connected = false;

static void wifi_event_handler(void* arg, esp_event_base_t event_base, 
                                int32_t event_id, void* event_data) {
    if (event_id == WIFI_EVENT_STA_START) {
        esp_wifi_connect();
    } else if (event_id == WIFI_EVENT_STA_CONNECTED) {
        wifi_connected = true;
        ESP_LOGI(TAG, "WiFi connected");
    } else if (event_id == WIFI_EVENT_STA_DISCONNECTED) {
        wifi_connected = false;
        ESP_LOGE(TAG, "WiFi disconnected");
        esp_wifi_connect();
    }
}

static void ip_event_handler(void* arg, esp_event_base_t event_base, 
                              int32_t event_id, void* event_data) {
    if (event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
        ESP_LOGI(TAG, "Got IP: " IPSTR, IP2STR(&event->ip_info.ip));
    }
}

void udp_test_task(void *pvParameters) {
    // Wait for WiFi to be fully connected
    while (!wifi_connected) {
        vTaskDelay(pdMS_TO_TICKS(1000));
    }

    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(UDP_PORT);
    
    // Use inet_pton for IP conversion
    if (inet_pton(AF_INET, UDP_SERVER_IP, &dest_addr.sin_addr) != 1) {
        ESP_LOGE(TAG, "Invalid IP address");
        vTaskDelete(NULL);
        return;
    }

    int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock < 0) {
        ESP_LOGE(TAG, "Failed to create socket: %s", strerror(errno));
        vTaskDelete(NULL);
        return;
    }

    // Set socket timeout
    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    // Sample UDP packet (minimal DNS query)
    char buffer[256] = {
        0x00, 0x01,  // Transaction ID
        0x01, 0x00,  // Standard query
        0x00, 0x01,  // 1 question
        0x00, 0x00,  // 0 answer RRs
        0x00, 0x00,  // 0 authority RRs
        0x00, 0x00,  // 0 additional RRs
        // Query details would follow
    };

    while (1) {
        int err = sendto(sock, buffer, sizeof(buffer), 0, 
                         (struct sockaddr *)&dest_addr, sizeof(dest_addr));
        if (err < 0) {
            ESP_LOGE(TAG, "UDP send failed: %s", strerror(errno));
            break;
        }

        char rx_buffer[256];
        socklen_t addr_len = sizeof(dest_addr);
        int len = recvfrom(sock, rx_buffer, sizeof(rx_buffer), 0, 
                           (struct sockaddr *)&dest_addr, &addr_len);
        
        if (len < 0) {
            if (errno == EWOULDBLOCK || errno == EAGAIN) {
                ESP_LOGI(TAG, "Receive timeout");
            } else {
                ESP_LOGE(TAG, "UDP receive failed: %s", strerror(errno));
            }
        } else {
            ESP_LOGI(TAG, "Received UDP packet of %d bytes", len);
        }
        
        vTaskDelay(pdMS_TO_TICKS(5000));  // Wait 5 seconds between tests
    }

    close(sock);
    vTaskDelete(NULL);
}

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

    // Initialize networking
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    // WiFi initialization
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // Event handlers
    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL));
    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT, ESP_EVENT_ANY_ID, &ip_event_handler, NULL));

    // WiFi configuration
    wifi_config_t wifi_config = {
        .sta = {
            .ssid = WIFI_SSID,
            .password = WIFI_PASS,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    // Create UDP test task
    xTaskCreate(udp_test_task, "udp_test", 4096, NULL, 5, NULL);
}