#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"

#define WIFI_SSID      "Wokwi-GUEST"
#define WIFI_PASS      ""
#define DNS_PORT       53
#define DNS_SERVER     "1.1.1.1"

static const char *TAG = "dns_test";

// DNS query packet matching the Erlang definition
static const uint8_t dns_query[] = {
    0x07, 0xe8,  // ID (2024 - matching your test)
    0x01, 0x00,  // Flags (RD bit set)
    0x00, 0x01,  // QDCOUNT
    0x00, 0x00,  // ANCOUNT
    0x00, 0x00,  // NSCOUNT
    0x00, 0x00,  // ARCOUNT
    0x06,        // Length of "github"
    'g', 'i', 't', 'h', 'u', 'b',
    0x03,        // Length of "com"
    'c', 'o', 'm',
    0x00,        // Null terminator
    0x00, 0x01,  // QTYPE (A record)
    0x00, 0x01   // QCLASS (IN)
};

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) {
        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_base == IP_EVENT && 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));
    }
}

static void wifi_init(void)
{
    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));

    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, IP_EVENT_STA_GOT_IP, &ip_event_handler, NULL));

    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(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
}

static void dns_query_task(void *pvParameters)
{

    //vTaskDelay(pdMS_TO_TICKS(5000));  // Wait 5 seconds before next query
    int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0) {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        vTaskDelete(NULL);
        return;
    }

    struct sockaddr_in dest_addr;
    dest_addr.sin_addr.s_addr = inet_addr(DNS_SERVER);
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(DNS_PORT);

    while (1) {
        ESP_LOGI(TAG, "About to send");
        int err = sendto(sock, dns_query, sizeof(dns_query), 0, 
                        (struct sockaddr *)&dest_addr, sizeof(dest_addr));
        if (err < 0) {
            ESP_LOGE(TAG, "Error sending DNS query: errno %d", errno);
        } else {
            ESP_LOGI(TAG, "DNS query sent");
            
            // Wait for response
            char rx_buffer[512];
            struct sockaddr_in source_addr;
            socklen_t socklen = sizeof(source_addr);
            int len = recvfrom(sock, rx_buffer, sizeof(rx_buffer) - 1, 0,
                             (struct sockaddr *)&source_addr, &socklen);
            
            if (len > 0) {
                ESP_LOGI(TAG, "Received %d bytes from %s", len,
                         inet_ntoa(source_addr.sin_addr));
                // Print first few bytes of response
                ESP_LOG_BUFFER_HEX(TAG, rx_buffer, len > 16 ? 16 : len);
            }
        }
        vTaskDelay(pdMS_TO_TICKS(5000));  // Wait 5 seconds before next query
    }
}

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 WiFi
    wifi_init();

    // Create DNS query task
    //xTaskCreate(dns_query_task, "dns_query_task", 4096, NULL, 5, NULL);
}