// remote_proxy.c — Удаленный TCP прокси (exit node)
#include "remote_proxy.h"
#include "tcp_proxy.h"
#include "udp_proxy.h"
#include "icmp_proxy.h"
#include "etcp.h"
#include "etcp_api.h"
#include "etcp_router.h"
#include "utun_instance.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 <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#endif
#ifndef MSG_NOSIGNAL
#define MSG_NOSIGNAL 0
#endif

static struct remote_proxy_ctx* g_rp_ctx = NULL;

static void rp_sock_read_cb (socket_t sock, void* arg);
static void rp_sock_write_cb(socket_t sock, void* arg);
static void rp_sock_error_cb(socket_t sock, void* arg);
static void rp_conn_free(struct remote_proxy_conn* rc);

static int rp_send_msg(struct UTUN_INSTANCE* inst, uint64_t dst, uint8_t subcmd,
                        uint64_t sid, uint16_t seq, const uint8_t* data, size_t len) {
    struct ll_entry* e = queue_entry_new(0);
    if (!e) return -1;
    e->dgram = u_malloc(TCP_PROXY_HDR_SIZE + len);
    if (!e->dgram) { queue_entry_free(e); return -1; }
    e->dgram[0] = ETCP_ID_TCP_PROXY;
    e->dgram[1] = subcmd;
    memcpy(e->dgram + 2, &sid, 8);
    memcpy(e->dgram + 10, &seq, 2);
    if (len > 0) memcpy(e->dgram + TCP_PROXY_HDR_SIZE, data, len);
    e->len = TCP_PROXY_HDR_SIZE + len;
    return etcp_route_send(inst, dst, e);
}

static int rp_send_connected(struct UTUN_INSTANCE* inst, uint64_t dst, uint64_t sid,
                              uint16_t local_port, uint8_t status) {
    uint8_t buf[3];
    memcpy(buf, &local_port, 2); buf[2] = status;
    return rp_send_msg(inst, dst, TCP_PROXY_SUBCMD_CONNECTED, sid, 0, buf, 3);
}

// ====================================================================
// Socket callbacks
// ====================================================================
static void rp_sock_read_cb(socket_t sock, void* arg) {
    (void)sock; struct remote_proxy_conn* rc = (struct remote_proxy_conn*)arg;
    if (!rc || rc->sock == SOCKET_INVALID) return;
    uint8_t buf[8192]; ssize_t n = recv(rc->sock, buf, sizeof(buf), 0);
    if (n > 0) {
        DEBUG_INFO(DEBUG_CATEGORY_TRAFFIC, "RP RECV ← stream=%016llx len=%zd seq=%u", (unsigned long long)rc->stream_id, n, rc->send_seq);
        struct UTUN_INSTANCE* inst = rc->ctx ? rc->ctx->inst : NULL;
        if (inst) {
            rp_send_msg(inst, rc->peer_node_id, TCP_PROXY_SUBCMD_DATA, rc->stream_id, rc->send_seq, buf, (size_t)n);
            rc->send_seq++;
        }
    } else {
        if (n == 0) DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy: EOF stream=%016llx", (unsigned long long)rc->stream_id);
        else DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: recv error %s", strerror(errno));
        struct UTUN_INSTANCE* inst = rc->ctx ? rc->ctx->inst : NULL;
    if (inst) rp_send_msg(inst, rc->peer_node_id, TCP_PROXY_SUBCMD_CLOSE, rc->stream_id, 0, NULL, 0);
        rc->connected = -1;
        rp_conn_free(rc);
    }
}

static void rp_sock_write_cb(socket_t sock, void* arg) {
    (void)sock; struct remote_proxy_conn* rc = (struct remote_proxy_conn*)arg;
    if (!rc || rc->sock == SOCKET_INVALID) return;
    if (!rc->connected) {
        if (!rc->connect_called) return;
        int err = 0; socklen_t len = sizeof(err);
        if (getsockopt(rc->sock, SOL_SOCKET, SO_ERROR, &err, &len) == 0 && err == 0) {
            rc->connected = 1;
            struct sockaddr_in local; socklen_t llen = sizeof(local);
            uint16_t local_port = 0;
            if (getsockname(rc->sock, (struct sockaddr*)&local, &llen) == 0) local_port = local.sin_port;
            DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy: connected stream=%016llx to %d.%d.%d.%d:%d local_port=%d",
                       (unsigned long long)rc->stream_id, rc->dest_ip[0], rc->dest_ip[1], rc->dest_ip[2], rc->dest_ip[3],
                       ntohs(rc->dest_port), ntohs(local_port));
            struct UTUN_INSTANCE* inst = rc->ctx ? rc->ctx->inst : NULL;
            if (inst) rp_send_connected(inst, rc->peer_node_id, rc->stream_id, local_port, TCP_PROXY_CONNECTED_OK);
        } else {
            DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: connect failed %s", err ? strerror(err) : "unknown");
            struct UTUN_INSTANCE* inst = rc->ctx ? rc->ctx->inst : NULL;
            if (inst) rp_send_connected(inst, rc->peer_node_id, rc->stream_id, 0, TCP_PROXY_CONNECTED_REFUSED);
            rc->connected = -1;
            rp_conn_free(rc);
        }
    }
}

static void rp_sock_error_cb(socket_t sock, void* arg) {
    (void)sock; struct remote_proxy_conn* rc = (struct remote_proxy_conn*)arg;
    if (!rc || rc->sock == SOCKET_INVALID) return;
    DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: socket error stream=%016llx", (unsigned long long)(rc ? rc->stream_id : 0));
    rc->connected = -1;
    struct UTUN_INSTANCE* inst = rc->ctx ? rc->ctx->inst : NULL;
    if (inst) rp_send_msg(inst, rc->peer_node_id, TCP_PROXY_SUBCMD_CLOSE, rc->stream_id, 0, NULL, 0);
    rp_conn_free(rc);
}

static void rp_conn_free(struct remote_proxy_conn* rc) {
    if (!rc) return;
    if (rc->ctx) {
        struct remote_proxy_conn** prev = &rc->ctx->conns;
        while (*prev) { if (*prev == rc) { *prev = rc->next; break; } prev = &(*prev)->next; }
    }
    if (rc->sock != SOCKET_INVALID) {
        if (rc->read_id) { uasync_remove_socket_t(rc->ua, rc->sock); rc->read_id = NULL; }
        socket_close_wrapper(rc->sock); rc->sock = SOCKET_INVALID;
    }
    u_free(rc);
}

// ====================================================================
// Standalone handler (registered when tcp_proxy doesn't have remote mappings)
// ====================================================================
static void rp_standalone_recv_cb(struct ETCP_CONN* conn, struct ll_entry* entry) {
    struct UTUN_INSTANCE* i = conn ? conn->instance : (g_rp_ctx ? g_rp_ctx->inst : NULL);
    if (!i || !entry || entry->dgram == NULL || entry->len < TCP_PROXY_HDR_SIZE) {
        if (entry) { queue_dgram_free(entry); queue_entry_free(entry); }
        return;
    }
    uint8_t  subcmd = entry->dgram[1];
    uint64_t sid = 0; memcpy(&sid, entry->dgram + 2, 8);
    uint64_t src = conn ? conn->peer_node_id : i->node_id;
    if (subcmd == TCP_PROXY_SUBCMD_CONNECT)   { remote_proxy_handle_connect(i, entry, sid, src); return; }
    if (subcmd == TCP_PROXY_SUBCMD_DATA)      { remote_proxy_handle_data(i, entry, sid); return; }
    if (subcmd == TCP_PROXY_SUBCMD_CLOSE)     { remote_proxy_handle_close(i, sid); queue_dgram_free(entry); queue_entry_free(entry); return; }
    queue_dgram_free(entry); queue_entry_free(entry);
}

// ====================================================================
// Public API
// ====================================================================
struct remote_proxy_conn* remote_proxy_find_conn(struct remote_proxy_ctx* ctx, uint64_t stream_id) {
    struct remote_proxy_conn* c;
    for (c = ctx->conns; c; c = c->next) if (c->stream_id == stream_id) return c;
    return NULL;
}

int remote_proxy_handle_connect(struct UTUN_INSTANCE* inst, struct ll_entry* entry,
                                 uint64_t stream_id, uint64_t src_node_id) {
    if (!inst || !inst->remote_proxy.enabled) { queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    struct remote_proxy_ctx* ctx = &inst->remote_proxy;
    if (entry->len < TCP_PROXY_CONNECT_HDR_SIZE) { queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    uint8_t* dest_ip    = entry->dgram + TCP_PROXY_HDR_SIZE;
    uint16_t dest_port  = 0; memcpy(&dest_port, dest_ip + 4, 2);

    struct remote_proxy_conn* rc = u_calloc(1, sizeof(struct remote_proxy_conn));
    if (!rc) { queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    rc->ctx = ctx; rc->stream_id = stream_id; rc->peer_node_id = src_node_id;
    memcpy(rc->dest_ip, dest_ip, 4); rc->dest_port = dest_port;
    rc->ua = inst->ua; rc->sock = SOCKET_INVALID;
    rc->send_seq = 0; rc->recv_seq_init = 0;

    rc->sock = socket(AF_INET, SOCK_STREAM, 0);
    if (rc->sock == SOCKET_INVALID) { DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: socket() failed"); rp_conn_free(rc); queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    socket_set_nonblocking(rc->sock);
    rc->read_id = uasync_add_socket_t(rc->ua, rc->sock, rp_sock_read_cb, rp_sock_write_cb, rp_sock_error_cb, rc);
    if (!rc->read_id) { DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: uasync_add_socket_t failed"); rp_conn_free(rc); queue_dgram_free(entry); queue_entry_free(entry); return -1; }

    struct sockaddr_in addr; memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET; memcpy(&addr.sin_addr.s_addr, dest_ip, 4); addr.sin_port = dest_port;
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy: connecting stream=%016llx to %d.%d.%d.%d:%d",
               (unsigned long long)stream_id, dest_ip[0], dest_ip[1], dest_ip[2], dest_ip[3], ntohs(dest_port));
    rc->connect_called = 1;
    int ret = connect(rc->sock, (struct sockaddr*)&addr, sizeof(addr));
    if (ret < 0 && errno != EINPROGRESS) {
        DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: connect() failed: %s", strerror(errno));
        rp_send_connected(inst, src_node_id, stream_id, 0, TCP_PROXY_CONNECTED_REFUSED);
        rp_conn_free(rc); queue_dgram_free(entry); queue_entry_free(entry); return -1;
    }

    rc->next = ctx->conns; ctx->conns = rc;
    queue_dgram_free(entry); queue_entry_free(entry);
    return 0;
}

int remote_proxy_handle_data(struct UTUN_INSTANCE* inst, struct ll_entry* entry, uint64_t stream_id) {
    if (!inst) { queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    struct remote_proxy_ctx* ctx = &inst->remote_proxy;
    struct remote_proxy_conn* rc = remote_proxy_find_conn(ctx, stream_id);
    if (!rc || rc->sock == SOCKET_INVALID || rc->connected != 1) { queue_dgram_free(entry); queue_entry_free(entry); return -1; }
    uint16_t seq; memcpy(&seq, entry->dgram + 10, 2);
    if (!rc->recv_seq_init) {
        rc->recv_last_seq = seq; rc->recv_seq_init = 1;
    } else {
        int16_t delta = (int16_t)(seq - rc->recv_last_seq);
        if (delta <= 0) {
            DEBUG_WARN(DEBUG_CATEGORY_SOCKET, "remote_proxy: duplicate DATA seq=%u stream=%016llx, discard", seq, (unsigned long long)stream_id);
            queue_dgram_free(entry); queue_entry_free(entry); return 0;
        }
        if (delta > 1) {
            DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: seq gap seq=%u last=%u stream=%016llx, tearing down", seq, rc->recv_last_seq, (unsigned long long)stream_id);
            rp_send_msg(inst, rc->peer_node_id, TCP_PROXY_SUBCMD_CLOSE, rc->stream_id, 0, NULL, 0);
            rp_conn_free(rc);
            queue_dgram_free(entry); queue_entry_free(entry); return -1;
        }
        rc->recv_last_seq = seq;
    }
    size_t data_len = entry->len - TCP_PROXY_HDR_SIZE;
    DEBUG_INFO(DEBUG_CATEGORY_TRAFFIC, "RP SEND → stream=%016llx seq=%u len=%zu", (unsigned long long)stream_id, seq, data_len);
    uint8_t* data = entry->dgram + TCP_PROXY_HDR_SIZE;
    ssize_t n = send(rc->sock, data, data_len, MSG_NOSIGNAL);
    if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
        DEBUG_ERROR(DEBUG_CATEGORY_SOCKET, "remote_proxy: send error %s", strerror(errno));
    }
    queue_dgram_free(entry); queue_entry_free(entry);
    return 0;
}

void remote_proxy_handle_close(struct UTUN_INSTANCE* inst, uint64_t stream_id) {
    if (!inst) return;
    struct remote_proxy_ctx* ctx = &inst->remote_proxy;
    struct remote_proxy_conn** prev = &ctx->conns;
    while (*prev) {
        struct remote_proxy_conn* rc = *prev;
        if (rc->stream_id == stream_id) {
            DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy: close stream=%016llx", (unsigned long long)stream_id);
            *prev = rc->next; rp_conn_free(rc); return;
        }
        prev = &rc->next;
    }
}

int remote_proxy_init(struct UTUN_INSTANCE* inst) {
    if (!inst) return -1;
    struct remote_proxy_ctx* ctx = &inst->remote_proxy;
    memset(ctx, 0, sizeof(*ctx));
    ctx->enabled = inst->config && inst->config->global.remote_proxy_enabled;
    ctx->inst = inst;
    if (!ctx->enabled) return 0;
    g_rp_ctx = ctx;
    // Register standalone handler for when tcp_proxy is not using remote mappings
    // (tcp_proxy_create will overwrite this handler if it has remote mappings)
    etcp_router_bind(inst, ETCP_ID_TCP_PROXY, rp_standalone_recv_cb);
    if (!inst->config || !inst->config->global.tcp_proxy_enabled) {
        udp_proxy_init(inst, inst->ua);
        icmp_proxy_init(inst, inst->ua);
    }
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy initialized on node %016llx", (unsigned long long)inst->node_id);
    return 0;
}

void remote_proxy_destroy(struct UTUN_INSTANCE* inst) {
    if (!inst) return;
    struct remote_proxy_ctx* ctx = &inst->remote_proxy;
    struct remote_proxy_conn* rc = ctx->conns;
    while (rc) { struct remote_proxy_conn* next = rc->next; rp_conn_free(rc); rc = next; }
    ctx->conns = NULL; ctx->enabled = 0;
    if (g_rp_ctx == ctx) g_rp_ctx = NULL;
    udp_proxy_destroy(inst);
    icmp_proxy_destroy(inst);
    DEBUG_INFO(DEBUG_CATEGORY_SOCKET, "remote_proxy destroyed");
}