#include <string.h>
#include <stdbool.h>

#include "lwip/sockets.h"
#include "lwip/inet.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "dns_server.h"

static const char *TAG = "DNS";
#define DNS_PORT 53

static TaskHandle_t s_dns_task = NULL;
static volatile bool s_dns_running = false;

static void dns_task(void *arg)
{
    (void)arg;

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        ESP_LOGE(TAG, "❌ socket DNS falhou");
        s_dns_running = false;
        s_dns_task = NULL;
        vTaskDelete(NULL);
        return;
    }

    int yes = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

    struct sockaddr_in addr = {0};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(DNS_PORT);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        ESP_LOGE(TAG, "❌ bind DNS falhou (porta %d)", DNS_PORT);
        close(sock);
        s_dns_running = false;
        s_dns_task = NULL;
        vTaskDelete(NULL);
        return;
    }

    ESP_LOGI(TAG, "✅ DNS captive a correr na porta %d (task)", DNS_PORT);

    while (s_dns_running) {
        uint8_t buf[512];
        struct sockaddr_in source_addr = {0};
        socklen_t slen = sizeof(source_addr);

        int len = recvfrom(sock, buf, sizeof(buf), 0,
                           (struct sockaddr*)&source_addr, &slen);

        if (len <= 0) {
            vTaskDelay(pdMS_TO_TICKS(10));
            continue;
        }

        buf[2] |= 0x80; // resposta
        buf[3] |= 0x80; // RA
        buf[6]  = 0x00;
        buf[7]  = 0x01; // ANCOUNT=1

        if (len + 16 < (int)sizeof(buf)) {
            buf[len++] = 0xC0; buf[len++] = 0x0C;
            buf[len++] = 0x00; buf[len++] = 0x01;
            buf[len++] = 0x00; buf[len++] = 0x01;
            buf[len++] = 0x00; buf[len++] = 0x00;
            buf[len++] = 0x00; buf[len++] = 0x3C;
            buf[len++] = 0x00; buf[len++] = 0x04;
            buf[len++] = 192;  buf[len++] = 168;
            buf[len++] = 4;    buf[len++] = 1;

            sendto(sock, buf, len, 0, (struct sockaddr*)&source_addr, slen);
        }
    }

    close(sock);
    s_dns_task = NULL;
    ESP_LOGI(TAG, "DNS task terminou");
    vTaskDelete(NULL);
}

void start_dns_server(void)
{
    if (s_dns_task) {
        ESP_LOGW(TAG, "DNS já está a correr");
        return;
    }
    s_dns_running = true;
    xTaskCreate(dns_task, "dns_task", 4096, NULL, 4, &s_dns_task);
    ESP_LOGI(TAG, "✅ DNS captive portal iniciado na porta %d", DNS_PORT);
}

void stop_dns_server(void)
{
    s_dns_running = false;
}