// icmp_proxy.c — ICMP echo (ping) прокси: клиент ↔ exit через etcp_router
#include "icmp_proxy.h"
#include "etcp.h"
#include "etcp_api.h"
#include "etcp_router.h"
#include "utun_instance.h"
#include "tun_if.h"
#include "../lib/u_async.h"
#include "../lib/debug_config.h"
#include "../lib/ll_queue.h"
#include "../lib/mem.h"
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#ifndef _WIN32
#include <unistd.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#else
#define ICMP_MINLEN      8
#define ICMP_ECHO        8
#define ICMP_ECHOREPLY   0
struct icmp {
    uint8_t  icmp_type;
    uint8_t  icmp_code;
    uint16_t icmp_cksum;
    uint16_t icmp_id;
    uint16_t icmp_seq;
    uint8_t  icmp_data[1];
};
struct ip {
    uint8_t  ip_hl:4, ip_v:4;
    uint8_t  ip_tos;
    uint16_t ip_len;
    uint16_t ip_id;
    uint16_t ip_off;
    uint8_t  ip_ttl;
    uint8_t  ip_p;
    uint16_t ip_sum;
    struct in_addr ip_src, ip_dst;
};
#endif

#define ICMP_DEF_TTL        64
#define ICMP_TIMEOUT_TB     50000  // 5s for echo reply

static struct icmp_proxy_ctx* g_icmp_ctx = NULL;

static void req_expire_timer_cb(void* arg);
static void req_expire(struct icmp_proxy_ctx* ctx);

static struct icmp_request* req_find_by_id(struct icmp_request* head, uint16_t echo_id, uint16_t echo_seq) {
    struct icmp_request* r;
    for (r = head; r; r = r->next) if (r->echo_id == echo_id && r->echo_seq == echo_seq) return r;
    return NULL;
}

// Сборка и отправка raw ICMP echo запроса
static int exit_send_echo(struct UTUN_INSTANCE* inst, uint64_t client_node_id,
                          uint32_t dst_ip, uint32_t orig_src_ip, uint16_t echo_id, uint16_t echo_seq,
                          const uint8_t* payload, size_t payload_len) {
    if (!g_icmp_ctx || g_icmp_ctx->raw_sock == SOCKET_INVALID) return -1;
    size_t icmp_len = ICMP_MINLEN + payload_len;
    uint8_t* buf = u_malloc(icmp_len);
    if (!buf) return -1;

    struct icmp* icmp_hdr = (struct icmp*)buf;
    memset(icmp_hdr, 0, icmp_len);
    icmp_hdr->icmp_type = ICMP_ECHO;
    icmp_hdr->icmp_code = 0;
    icmp_hdr->icmp_id = echo_id;
    icmp_hdr->icmp_seq = echo_seq;
    if (payload_len > 0) memcpy(icmp_hdr->icmp_data, payload, payload_len);
    icmp_hdr->icmp_cksum = 0;
    uint32_t sum = 0; uint16_t* w = (uint16_t*)icmp_hdr;
    for (size_t i = 0; i < (icmp_len + 1) / 2; i++) sum += w[i];
    while (sum >> 16) sum = (sum & 0xFFFF) + (sum >> 16);
    icmp_hdr->icmp_cksum = ~(uint16_t)sum;

    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: sendto dst=0x%08x id=0x%04x seq=%u len=%zu",
               dst_ip, echo_id, echo_seq, icmp_len);
    struct sockaddr_in addr = {.sin_family = AF_INET, .sin_addr = {.s_addr = dst_ip}};
    ssize_t n = sendto(g_icmp_ctx->raw_sock, buf, icmp_len, 0, (struct sockaddr*)&addr, sizeof(addr));
    u_free(buf);
    if (n < 0) { DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: sendto failed: %s", strerror(errno)); return -1; }
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: sendto sent %zd bytes", n);

    struct icmp_request* r = u_calloc(1, sizeof(struct icmp_request));
    if (!r) return 0;
    r->client_node_id = client_node_id; r->dst_ip = dst_ip; r->orig_src_ip = orig_src_ip;
    r->echo_id = echo_id; r->echo_seq = echo_seq;
    r->payload_len = payload_len > sizeof(r->payload) ? sizeof(r->payload) : payload_len;
    if (payload_len > 0) memcpy(r->payload, payload, r->payload_len);
    r->sent_tb = get_time_tb();
    r->next = g_icmp_ctx->pending; g_icmp_ctx->pending = r;
    if (!g_icmp_ctx->expire_timer)
        g_icmp_ctx->expire_timer = uasync_set_timeout(g_icmp_ctx->ua, g_icmp_ctx->request_timeout_tb, g_icmp_ctx, req_expire_timer_cb, "icmp_expire");
    return 0;
}

// Чтение echo ответа из raw сокета, сопоставление по echo_id
static void raw_read_cb(socket_t sock, void* arg) {
    (void)sock; (void)arg;
    if (!g_icmp_ctx || g_icmp_ctx->raw_sock == SOCKET_INVALID) return;
    uint8_t buf[65536];
    struct sockaddr_in from; socklen_t flen = sizeof(from);
    ssize_t n = recvfrom(g_icmp_ctx->raw_sock, buf, sizeof(buf), 0, (struct sockaddr*)&from, &flen);
    if (n <= 0) { if (n < 0) DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: recvfrom error: %s", strerror(errno)); return; }
    if (n < (ssize_t)(sizeof(struct ip) + ICMP_MINLEN)) return;

    struct ip* ip_hdr = (struct ip*)buf;
    if (ip_hdr->ip_p != IPPROTO_ICMP) return;
    size_t ip_hdr_len = ip_hdr->ip_hl * 4;
    if (n < (ssize_t)(ip_hdr_len + ICMP_MINLEN)) return;
    struct icmp* icmp_hdr = (struct icmp*)(buf + ip_hdr_len);
    if (icmp_hdr->icmp_type != ICMP_ECHOREPLY) return;

    struct icmp_request* r = req_find_by_id(g_icmp_ctx->pending, icmp_hdr->icmp_id, icmp_hdr->icmp_seq);
    if (!r) { DEBUG_WARN(DEBUG_CATEGORY_SOCKET, "icmp_proxy: unclaimed echo reply id=0x%04x seq=%u from=0x%08x",
                          icmp_hdr->icmp_id, icmp_hdr->icmp_seq, from.sin_addr.s_addr); return; }
    {   struct icmp_request** rp = &g_icmp_ctx->pending;
        while (*rp) { if (*rp == r) { *rp = r->next; break; } rp = &(*rp)->next; } }
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: echo reply id=0x%04x seq=%u from=0x%08x",
               icmp_hdr->icmp_id, icmp_hdr->icmp_seq, from.sin_addr.s_addr);

    size_t payload_len = n - ip_hdr_len - ICMP_MINLEN;
    if (payload_len > 1500) payload_len = 1500;
    uint8_t* payload = (payload_len > 0) ? (uint8_t*)(icmp_hdr->icmp_data) : NULL;

    struct ll_entry* e = queue_entry_new(0);
    if (!e) return;
    e->dgram = u_malloc(ICMP_PROXY_HDR_SIZE + payload_len);
    if (!e->dgram) { queue_entry_free(e); return; }
    e->dgram[0] = ETCP_RT_ID_ICMP_PROXY;
    e->dgram[1] = ICMP_PROXY_SUBCMD_REPLY;
    memcpy(e->dgram + 2,  &r->client_node_id, 8);
    memcpy(e->dgram + 10, &r->dst_ip, 4);
    memcpy(e->dgram + 14, &r->orig_src_ip, 4);
    memcpy(e->dgram + 18, &icmp_hdr->icmp_id, 2);
    memcpy(e->dgram + 20, &icmp_hdr->icmp_seq, 2);
    if (payload_len > 0) memcpy(e->dgram + ICMP_PROXY_HDR_SIZE, payload, payload_len);
    e->len = ICMP_PROXY_HDR_SIZE + payload_len;
    int ret = etcp_route_send(g_icmp_ctx->inst, r->client_node_id, e, 0);
    if (ret != 0) DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: etcp_route_send reply failed: %d", ret);
    else DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: reply forwarded to client %016llx", (unsigned long long)r->client_node_id);
    u_free(r);
}

// ====================================================================
// Exit узел: принять REQUEST, отправить echo через raw сокет
// ====================================================================
static void exit_handle_request(struct ETCP_CONN* conn, struct ll_entry* entry) {
    struct UTUN_INSTANCE* inst = conn ? conn->instance : (g_icmp_ctx ? g_icmp_ctx->inst : NULL);
    if (!inst || !g_icmp_ctx || entry->len < ICMP_PROXY_HDR_SIZE + 1) goto drop;

    uint64_t client_node_id; memcpy(&client_node_id, entry->dgram + 2, 8);
    uint32_t dst_ip;       memcpy(&dst_ip,       entry->dgram + 10, 4);
    uint32_t orig_src_ip;  memcpy(&orig_src_ip,  entry->dgram + 14, 4);
    uint16_t echo_id;      memcpy(&echo_id,      entry->dgram + 18, 2);
    uint16_t echo_seq;     memcpy(&echo_seq,     entry->dgram + 20, 2);
    uint8_t* payload = entry->dgram + ICMP_PROXY_HDR_SIZE;
    size_t payload_len = entry->len - ICMP_PROXY_HDR_SIZE;

    if (g_icmp_ctx->raw_sock != SOCKET_INVALID) {
        exit_send_echo(inst, client_node_id, dst_ip, orig_src_ip, echo_id, echo_seq, payload, payload_len);
    } else if (g_icmp_ctx->test_loopback) {
        DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: test loopback reply to 0x%08x", dst_ip);
        struct ll_entry* e = queue_entry_new(0);
        if (e) {
            e->dgram = u_malloc(ICMP_PROXY_HDR_SIZE + payload_len);
            if (e->dgram) {
                e->dgram[0] = ETCP_RT_ID_ICMP_PROXY;
                e->dgram[1] = ICMP_PROXY_SUBCMD_REPLY;
                memcpy(e->dgram + 2,  &client_node_id, 8);
                memcpy(e->dgram + 10, &dst_ip, 4);
                memcpy(e->dgram + 14, &orig_src_ip, 4);
                memcpy(e->dgram + 18, &echo_id, 2);
                memcpy(e->dgram + 20, &echo_seq, 2);
                if (payload_len > 0) memcpy(e->dgram + ICMP_PROXY_HDR_SIZE, payload, payload_len);
                e->len = ICMP_PROXY_HDR_SIZE + payload_len;
                etcp_route_send(inst, client_node_id, e, 0);
            } else queue_entry_free(e);
        }
    } else {
        DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: no raw socket, dropping echo request to 0x%08x", dst_ip);
    }

drop:
    queue_dgram_free(entry); queue_entry_free(entry);
}

// ====================================================================
// Сторона клиента: принять REPLY, доставить echo ответ в TUN
// ====================================================================
static void client_handle_reply(struct ETCP_CONN* conn, struct ll_entry* entry) {
    if (entry->len < ICMP_PROXY_HDR_SIZE + 1) { queue_dgram_free(entry); queue_entry_free(entry); return; }
    uint32_t src_ip;
    uint32_t orig_src_ip;
    uint16_t echo_id, echo_seq;
    memcpy(&src_ip,      entry->dgram + 10, 4);
    memcpy(&orig_src_ip, entry->dgram + 14, 4);
    memcpy(&echo_id,     entry->dgram + 18, 2);
    memcpy(&echo_seq,    entry->dgram + 20, 2);
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: client got reply id=0x%04x seq=%u from=0x%08x dst=0x%08x",
               echo_id, echo_seq, src_ip, orig_src_ip);
    struct UTUN_INSTANCE* inst = conn ? conn->instance : (g_icmp_ctx ? g_icmp_ctx->inst : NULL);
    icmp_proxy_deliver_reply(inst, orig_src_ip, src_ip, echo_id, echo_seq,
        entry->dgram + ICMP_PROXY_HDR_SIZE, entry->len - ICMP_PROXY_HDR_SIZE);
    queue_dgram_free(entry); queue_entry_free(entry);
}

// ====================================================================
// Единый etcp_router коллбэк
// ====================================================================
void icmp_proxy_recv_cb(struct ETCP_CONN* conn, struct ll_entry* entry) {
    if (!entry || !entry->dgram || entry->len < 2) {
        if (entry) { queue_dgram_free(entry); queue_entry_free(entry); }
        return;
    }
    uint8_t subcmd = entry->dgram[1];
    if (subcmd == ICMP_PROXY_SUBCMD_REQUEST) { exit_handle_request(conn, entry); return; }
    if (subcmd == ICMP_PROXY_SUBCMD_REPLY)   { client_handle_reply(conn, entry); return; }
    queue_dgram_free(entry); queue_entry_free(entry);
}

// ====================================================================
// Сторона клиента: отправить ICMP echo запрос в exit узел
// ====================================================================
int icmp_proxy_send_to_exit(struct UTUN_INSTANCE* inst, uint64_t exit_node_id,
                            uint32_t dst_ip, uint32_t orig_src_ip, uint16_t echo_id, uint16_t echo_seq,
                            const uint8_t* payload, size_t payload_len) {
    if (!inst) return -1;
    struct ll_entry* e = queue_entry_new(0);
    if (!e) return -1;
    e->dgram = u_malloc(ICMP_PROXY_HDR_SIZE + payload_len);
    if (!e->dgram) { queue_entry_free(e); return -1; }
    e->dgram[0] = ETCP_RT_ID_ICMP_PROXY;
    e->dgram[1] = ICMP_PROXY_SUBCMD_REQUEST;
    memcpy(e->dgram + 2,  &inst->node_id, 8);
    memcpy(e->dgram + 10, &dst_ip, 4);
    memcpy(e->dgram + 14, &orig_src_ip, 4);
    memcpy(e->dgram + 18, &echo_id, 2);
    memcpy(e->dgram + 20, &echo_seq, 2);
    if (payload_len > 0) memcpy(e->dgram + ICMP_PROXY_HDR_SIZE, payload, payload_len);
    e->len = ICMP_PROXY_HDR_SIZE + payload_len;
    return etcp_route_send(inst, exit_node_id, e, 0);
}

// ====================================================================
// Сторона клиента: доставить ICMP echo ответ в TUN
// ====================================================================
int icmp_proxy_deliver_reply(struct UTUN_INSTANCE* inst,
                             uint32_t dst_ip, uint32_t src_ip, uint16_t echo_id, uint16_t echo_seq,
                             const uint8_t* payload, size_t payload_len) {
    if (!inst || !inst->tcp_proxy_client || !inst->tcp_proxy_client->tun) {
        DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: deliver_reply failed — no %s",
                     !inst ? "inst" : !inst->tcp_proxy_client ? "tcp_proxy_client" : "TUN");
        return -1;
    }
    struct tun_if* tun = inst->tcp_proxy_client->tun;

    size_t icmp_len = ICMP_MINLEN + payload_len;
    size_t pkt_len = 20 + icmp_len;
    uint8_t* pkt = u_malloc(pkt_len);
    if (!pkt) return -1;

    memset(pkt, 0, 20);
    pkt[0] = 0x45; pkt[8] = 64; pkt[9] = IPPROTO_ICMP;
    uint16_t total_len = htons(20 + icmp_len);
    memcpy(pkt + 2, &total_len, 2);
    memcpy(pkt + 12, &src_ip, 4);
    memcpy(pkt + 16, &dst_ip, 4);

    {   uint32_t ip_sum = 0; for (int i = 0; i < 10; i++) { uint16_t w; memcpy(&w, pkt + i*2, 2); ip_sum += w; }
        ip_sum = (ip_sum & 0xFFFF) + (ip_sum >> 16); ip_sum += (ip_sum >> 16);
        uint16_t cs = (uint16_t)(~ip_sum); memcpy(pkt + 10, &cs, 2); }

    struct icmp* icmp_hdr = (struct icmp*)(pkt + 20);
    icmp_hdr->icmp_type = ICMP_ECHOREPLY; icmp_hdr->icmp_code = 0;
    icmp_hdr->icmp_id = echo_id; icmp_hdr->icmp_seq = echo_seq;
    if (payload_len > 0) memcpy(icmp_hdr->icmp_data, payload, payload_len);
    icmp_hdr->icmp_cksum = 0;
    uint32_t sum = 0; uint16_t* w = (uint16_t*)icmp_hdr;
    for (size_t i = 0; i < (icmp_len + 1) / 2; i++) sum += w[i];
    while (sum >> 16) sum = (sum & 0xFFFF) + (sum >> 16);
    icmp_hdr->icmp_cksum = ~(uint16_t)sum;

    struct ll_entry* e = queue_entry_new(0);
    if (!e) { u_free(pkt); return -1; }
    e->dgram = u_malloc(1 + pkt_len);
    e->dgram[0] = 4; memcpy(e->dgram + 1, pkt, pkt_len); e->len = 1 + pkt_len;
    u_free(pkt);
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: reply delivered to TUN id=0x%04x seq=%u", echo_id, echo_seq);
    queue_data_put(tun->input_queue, e);
    return 0;
}

void icmp_proxy_set_test_loopback(struct UTUN_INSTANCE* inst, int enabled) {
    (void)inst;
    if (g_icmp_ctx) g_icmp_ctx->test_loopback = enabled;
}

static void req_expire_timer_cb(void* arg) {
    struct icmp_proxy_ctx* ctx = (struct icmp_proxy_ctx*)arg;
    req_expire(ctx);
    ctx->expire_timer = NULL;
    if (ctx->pending)
        ctx->expire_timer = uasync_set_timeout(ctx->ua, ctx->request_timeout_tb, ctx, req_expire_timer_cb, "icmp_expire");
}

static void req_expire(struct icmp_proxy_ctx* ctx) {
    uint64_t now = get_time_tb(); struct icmp_request** prev = &ctx->pending;
    while (*prev) {
        struct icmp_request* r = *prev;
        if (now - r->sent_tb > ctx->request_timeout_tb) { *prev = r->next; u_free(r); }
        else prev = &r->next;
    }
}

// ====================================================================
// Публичные init/destroy
// ====================================================================
int icmp_proxy_init(struct UTUN_INSTANCE* inst, struct UASYNC* ua) {
    if (!inst) return -1;
    struct icmp_proxy_ctx* ctx = u_calloc(1, sizeof(struct icmp_proxy_ctx));
    if (!ctx) return -1;
    ctx->inst = inst; ctx->ua = ua; ctx->raw_sock = SOCKET_INVALID;
    ctx->request_timeout_tb = ICMP_TIMEOUT_TB;
    ctx->is_exit = inst->tcp_proxy_server.enabled;
    ctx->test_loopback = 0;
    ctx->expire_timer = NULL;
    g_icmp_ctx = ctx;

    if (ctx->is_exit) {
        ctx->raw_sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
        if (ctx->raw_sock == SOCKET_INVALID)
            DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "icmp_proxy: raw socket(SOCK_RAW) failed: %s", strerror(errno));
        else {
            DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy: raw socket created fd=%d", ctx->raw_sock);
            ctx->raw_read_id = uasync_add_socket_t(ua, ctx->raw_sock, raw_read_cb, NULL, NULL, NULL);
            if (!ctx->raw_read_id) { socket_close_wrapper(ctx->raw_sock); ctx->raw_sock = SOCKET_INVALID; }
        }
    }

    etcp_router_bind(inst, ETCP_RT_ID_ICMP_PROXY, icmp_proxy_recv_cb);
    ctx->initialized = 1;
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy initialized (exit=%d raw_sock=%d)", ctx->is_exit, (int)ctx->raw_sock);
    return 0;
}

void icmp_proxy_destroy(struct UTUN_INSTANCE* inst) {
    if (!inst || !g_icmp_ctx) return;
    etcp_router_unbind(inst, ETCP_RT_ID_ICMP_PROXY);
    if (g_icmp_ctx->expire_timer) { uasync_cancel_timeout(g_icmp_ctx->ua, g_icmp_ctx->expire_timer); g_icmp_ctx->expire_timer = NULL; }
    if (g_icmp_ctx->raw_sock != SOCKET_INVALID) {
        if (g_icmp_ctx->raw_read_id) { uasync_remove_socket_t(g_icmp_ctx->ua, g_icmp_ctx->raw_sock); g_icmp_ctx->raw_read_id = NULL; }
        socket_close_wrapper(g_icmp_ctx->raw_sock);
    }
    struct icmp_request* r = g_icmp_ctx->pending;
    while (r) { struct icmp_request* n = r->next; u_free(r); r = n; }
    u_free(g_icmp_ctx); g_icmp_ctx = NULL;
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "icmp_proxy destroyed");
}