/*
 * MTG RFID Companion - Captive Portal DNS Server
 *
 * Lightweight UDP DNS server on port 53. Resolves every domain query
 * to 192.168.4.1 to trigger automatic captive portal detection on smartphones.
 */

#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_log.h"
#include "lwip/sockets.h"
#include "lwip/netdb.h"
#include "dns_server.h"

static const char *TAG = "dns";

#define DNS_PORT 53
#define DNS_MAX_LEN 512

static TaskHandle_t s_dns_task = NULL;
static int s_sock = -1;
static volatile bool s_running = false;

typedef struct __attribute__((packed)) {
    uint16_t id;
    uint16_t flags;
    uint16_t qdcount;
    uint16_t ancount;
    uint16_t nscount;
    uint16_t arcount;
} dns_header_t;

static void dns_server_task(void *pvParameters)
{
    uint8_t rx_buffer[DNS_MAX_LEN];
    uint8_t tx_buffer[DNS_MAX_LEN];
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof(client_addr);

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

    s_sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (s_sock < 0) {
        ESP_LOGE(TAG, "Unable to create DNS socket: errno %d", errno);
        s_running = false;
        vTaskDelete(NULL);
        return;
    }

    struct timeval timeout = {
        .tv_sec = 1,
        .tv_usec = 0,
    };
    setsockopt(s_sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    if (bind(s_sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        ESP_LOGE(TAG, "DNS socket bind failed: errno %d", errno);
        close(s_sock);
        s_sock = -1;
        s_running = false;
        vTaskDelete(NULL);
        return;
    }

    ESP_LOGI(TAG, "Captive portal DNS server listening on UDP port 53");

    while (s_running) {
        int len = recvfrom(s_sock, rx_buffer, sizeof(rx_buffer), 0,
                           (struct sockaddr *)&client_addr, &addr_len);

        if (len < (int)sizeof(dns_header_t)) {
            continue;
        }

        dns_header_t *req_hdr = (dns_header_t *)rx_buffer;
        int qdcount = ntohs(req_hdr->qdcount);
        if (qdcount < 1) {
            continue;
        }

        /* Parse Question section to find end of question */
        int idx = sizeof(dns_header_t);
        while (idx < len && rx_buffer[idx] != 0) {
            idx += 1 + rx_buffer[idx];
        }
        idx += 1; /* null terminator */
        idx += 4; /* QTYPE + QCLASS */

        if (idx > len || idx + 16 > DNS_MAX_LEN) {
            continue;
        }

        /* Build Response Packet */
        memcpy(tx_buffer, rx_buffer, idx);
        dns_header_t *resp_hdr = (dns_header_t *)tx_buffer;
        resp_hdr->flags = htons(0x8180); /* Standard response, recursion available, no error */
        resp_hdr->ancount = htons(1);
        resp_hdr->nscount = 0;
        resp_hdr->arcount = 0;

        /* Answer Section */
        uint8_t *ans = tx_buffer + idx;
        /* Name pointer to question at offset 12 (0xC00C) */
        *ans++ = 0xC0;
        *ans++ = 0x0C;
        /* Type: A (IPv4) */
        *ans++ = 0x00;
        *ans++ = 0x01;
        /* Class: IN */
        *ans++ = 0x00;
        *ans++ = 0x01;
        /* TTL: 60 seconds */
        *ans++ = 0x00;
        *ans++ = 0x00;
        *ans++ = 0x00;
        *ans++ = 0x3C;
        /* Data length: 4 bytes */
        *ans++ = 0x00;
        *ans++ = 0x04;
        /* IP Address: 192.168.4.1 */
        *ans++ = 192;
        *ans++ = 168;
        *ans++ = 4;
        *ans++ = 1;

        int tx_len = ans - tx_buffer;
        sendto(s_sock, tx_buffer, tx_len, 0,
               (struct sockaddr *)&client_addr, addr_len);
    }

    if (s_sock >= 0) {
        close(s_sock);
        s_sock = -1;
    }
    ESP_LOGI(TAG, "DNS server stopped");
    s_dns_task = NULL;
    vTaskDelete(NULL);
}

esp_err_t dns_server_start(void)
{
    if (s_running) {
        return ESP_OK;
    }
    s_running = true;
    BaseType_t ret = xTaskCreatePinnedToCore(
        dns_server_task, "dns_server", 3072, NULL, 4, &s_dns_task, 0);
    if (ret != pdPASS) {
        s_running = false;
        return ESP_ERR_NO_MEM;
    }
    return ESP_OK;
}

void dns_server_stop(void)
{
    s_running = false;
}