#include "nat_transport.h"
#include "eim_nat.h"
#include "etcp.h"
#include "utun_instance.h"
#include "config_parser.h"
#include "tun_if.h"
#include "etcp_api.h"
#include "etcp_router.h"
#include "../routing_layer/topo_node.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "../lib/ll_queue.h"
#include <string.h>

#define NAT_SVC_HDR_SIZE 9  // svc_id(1) + src_node_id(8)

static ip_str_t ip_host_to_str(uint32_t ip_host) {
    struct in_addr a; a.s_addr = htonl(ip_host);
    return ip_to_str(&a, AF_INET);
}

// ==================== Callbacks ====================

// CLIENT: NAT TUN output → encapsulate in ETCP_RT_ID_NAT → send to provider via etcp_router
static void nat_transport_client_tun_out_cb(struct ll_queue* q, void* arg) {
    struct UTUN_INSTANCE* inst = (struct UTUN_INSTANCE*)arg;
    if (!inst) { queue_resume_callback(q); return; }
    struct nat_transport_ctx* tr = &inst->nat_tr;

    struct ll_entry* pkt = queue_data_get(q);
    if (!pkt) { queue_resume_callback(q); return; }
    if (!pkt->dgram || pkt->len < 2) { queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }

    size_t ip_len    = pkt->len - 1;
    size_t total_len = NAT_SVC_HDR_SIZE + ip_len;
    uint8_t* new_dgram = u_malloc(total_len);
    if (!new_dgram) { queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }

    new_dgram[0] = ETCP_RT_ID_NAT;
    memcpy(new_dgram + 1, &tr->self_node_id, 8);
    memcpy(new_dgram + 9, pkt->dgram + 1, ip_len);

    struct ll_entry* new_entry = queue_entry_new(0);
    if (!new_entry) { u_free(new_dgram); queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }
    new_entry->dgram = new_dgram;
    new_entry->len   = total_len;

    queue_dgram_free(pkt); queue_entry_free(pkt);
    queue_resume_callback(q);

    int ret = etcp_route_send(inst, TOPO_GROUP_UTUN, tr->nat_via_node_id, new_entry, 0);
    if (ret != 0) {
        DEBUG_WARN(DEBUG_CATEGORY_NAT, "NAT client: etcp_route_send to provider %016llx failed",
                   (unsigned long long)tr->nat_via_node_id);
        queue_dgram_free(new_entry); queue_entry_free(new_entry);
    }
}

// PROVIDER: NAT TUN output (internet response) → ingress NAT → send back via etcp_router
static void nat_transport_provider_tun_out_cb(struct ll_queue* q, void* arg) {
    struct UTUN_INSTANCE* inst = (struct UTUN_INSTANCE*)arg;
    if (!inst) { queue_resume_callback(q); return; }
    struct eim_nat_ctx* ctx = &inst->nat;

    struct ll_entry* pkt = queue_data_get(q);
    if (!pkt) { queue_resume_callback(q); return; }
    if (!pkt->dgram || pkt->len < 2) { queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }

    struct eim_nat_entry* entry = NULL;
    int ret = eim_nat_ingress(ctx, pkt->dgram + 1, pkt->len - 1, &entry);
    if (ret != 0 || !entry || entry->src_node_id == 0) {
        if (ret == 0) DEBUG_WARN(DEBUG_CATEGORY_NAT, "Ingress NAT: no matching entry, dropping");
        queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q);
        return;
    }

    size_t ip_len    = pkt->len - 1;
    size_t total_len = NAT_SVC_HDR_SIZE + ip_len;
    uint8_t* new_dgram = u_malloc(total_len);
    if (!new_dgram) { queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }

    new_dgram[0] = ETCP_RT_ID_NAT;
    memcpy(new_dgram + 1, &inst->nat_tr.self_node_id, 8);
    memcpy(new_dgram + 9, pkt->dgram + 1, ip_len);

    struct ll_entry* new_entry = queue_entry_new(0);
    if (!new_entry) { u_free(new_dgram); queue_dgram_free(pkt); queue_entry_free(pkt); queue_resume_callback(q); return; }
    new_entry->dgram = new_dgram;
    new_entry->len   = total_len;

    queue_dgram_free(pkt); queue_entry_free(pkt);
    queue_resume_callback(q);

    int send_ret = etcp_route_send(inst, TOPO_GROUP_UTUN, entry->src_node_id, new_entry, 0);
    if (send_ret != 0) {
        DEBUG_WARN(DEBUG_CATEGORY_NAT, "NAT provider: etcp_route_send back to node %016llx failed",
                   (unsigned long long)entry->src_node_id);
        queue_dgram_free(new_entry); queue_entry_free(new_entry);
    }
}

// ETCP_RT_ID_NAT receive via etcp_router: CLIENT gets response, PROVIDER gets request
static void nat_transport_etcp_recv_cb(struct ETCP_CONN* conn, struct ll_entry* entry) {
    if (!conn || !entry || !entry->dgram || entry->len < NAT_SVC_HDR_SIZE) {
        if (entry) { queue_dgram_free(entry); queue_entry_free(entry); }
        return;
    }
    struct UTUN_INSTANCE* inst = conn->instance;
    if (!inst) { queue_dgram_free(entry); queue_entry_free(entry); return; }
    struct nat_transport_ctx* tr = &inst->nat_tr;
    struct eim_nat_ctx* ctx = &inst->nat;
    if (!ctx->initialized) { queue_dgram_free(entry); queue_entry_free(entry); return; }

    uint64_t src_node_id;
    memcpy(&src_node_id, entry->dgram + 1, 8);
    uint8_t* ip_data = entry->dgram + NAT_SVC_HDR_SIZE;
    size_t   ip_len  = entry->len - NAT_SVC_HDR_SIZE;

    if (tr->nat_via_node_id != 0) {
        if (tr->nat_tun) {
            tun_write(tr->nat_tun, entry->dgram + NAT_SVC_HDR_SIZE - 1, ip_len + 1);
            DEBUG_DEBUG(DEBUG_CATEGORY_NAT, "NAT client: received %zu bytes from provider", ip_len);
        }
    } else {
        int ret = eim_nat_egress(ctx, ip_data, ip_len, src_node_id, conn);
        if (ret < 0) {
            DEBUG_WARN(DEBUG_CATEGORY_NAT, "Egress NAT failed");
        } else if (ret == 0 && tr->nat_tun) {
            tun_write(tr->nat_tun, entry->dgram + NAT_SVC_HDR_SIZE - 1, ip_len + 1);
            DEBUG_DEBUG(DEBUG_CATEGORY_NAT, "NAT provider: sent %zu bytes to internet", ip_len);
        }
    }

    queue_dgram_free(entry); queue_entry_free(entry);
}

// ==================== Init / Destroy ====================

int nat_transport_init(struct UTUN_INSTANCE* inst) {
    if (!inst || !inst->config) return -1;
    struct nat_transport_ctx* tr = &inst->nat_tr;
    struct global_config* g = &inst->config->global;

    memset(tr, 0, sizeof(*tr));
    if (!g->nat_enabled) return 0;

    tr->self_node_id    = inst->node_id;
    tr->nat_via_node_id = g->nat_via_node_id;

    if (tr->nat_via_node_id == 0) {
        if (eim_nat_init_ctx(&inst->nat, g) != 0) {
            DEBUG_ERROR(DEBUG_CATEGORY_NAT, "Failed to init NAT engine");
            return -1;
        }
    } else {
        inst->nat.initialized = 1;
    }

    const char* tun_name = g->nat_tun_ifname[0] ? g->nat_tun_ifname : "tun_nat";
    char ip_str[64] = "";
    if (g->nat_tun_ip.family == AF_INET) {
        snprintf(ip_str, sizeof(ip_str), "%s", ip_to_str(&g->nat_tun_ip.addr.v4, AF_INET).str);
    } else if (g->tun_ip.family == AF_INET) {
        snprintf(ip_str, sizeof(ip_str), "%s", ip_to_str(&g->tun_ip.addr.v4, AF_INET).str);
    } else {
        DEBUG_ERROR(DEBUG_CATEGORY_NAT, "No NAT TUN IP");
        eim_nat_destroy_ctx(&inst->nat);
        return -1;
    }
    tr->nat_tun = tun_init_nat(inst->ua, tun_name, ip_str, g->mtu, g->tun_test_mode);
    if (!tr->nat_tun) {
        DEBUG_ERROR(DEBUG_CATEGORY_NAT, "Failed to create NAT TUN %s ip=%s", tun_name, ip_str);
        eim_nat_destroy_ctx(&inst->nat);
        return -1;
    }

    if (etcp_router_bind(inst, ETCP_RT_ID_NAT, nat_transport_etcp_recv_cb) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_NAT, "Failed to bind ETCP_RT_ID_NAT via etcp_router");
        tun_close(tr->nat_tun); tr->nat_tun = NULL;
        eim_nat_destroy_ctx(&inst->nat);
        return -1;
    }

    struct eim_nat_ctx* ctx = &inst->nat;
    if (tr->nat_via_node_id != 0) {
        if (tr->nat_tun->output_queue) queue_set_callback(tr->nat_tun->output_queue, nat_transport_client_tun_out_cb, inst);
        DEBUG_INFO(DEBUG_CATEGORY_NAT, "NAT client via node %016llx, TUN=%s gw=%s ports=%u-%u",
                   (unsigned long long)tr->nat_via_node_id, tun_name,
                   ip_host_to_str(ctx->gateway_ip).str, ctx->port_start, ctx->port_end);
    } else {
        if (tr->nat_tun->output_queue) queue_set_callback(tr->nat_tun->output_queue, nat_transport_provider_tun_out_cb, inst);
        DEBUG_INFO(DEBUG_CATEGORY_NAT, "NAT provider TUN=%s gw=%s ports=%u-%u",
                   tun_name, ip_host_to_str(ctx->gateway_ip).str, ctx->port_start, ctx->port_end);
    }

    tr->initialized = 1;
    return 0;
}

void nat_transport_destroy(struct UTUN_INSTANCE* inst) {
    if (!inst || !inst->nat_tr.initialized) return;
    struct nat_transport_ctx* tr = &inst->nat_tr;

    etcp_router_unbind(inst, ETCP_RT_ID_NAT);

    if (tr->nat_tun) { tun_close(tr->nat_tun); tr->nat_tun = NULL; }

    eim_nat_destroy_ctx(&inst->nat);

    memset(tr, 0, sizeof(*tr));
    DEBUG_INFO(DEBUG_CATEGORY_NAT, "NAT transport destroyed");
}