// test_tcp_proxy.c — TCP proxy test: 1MB echo through 2 lwIP TCP instances + socketpair
// No root required — uses raw-fd mode instead of TUN
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/wait.h>
#include <errno.h>
#include <time.h>
#include "test_utils.h"
#include "../src/tcp_proxy.h"
#include "../src/config_parser.h"
#include "../src/lwip_tcp/lwip_tcp.h"
#include "../lib/u_async.h"
#include "../lib/debug_config.h"

#define TEST_PORT      9090
#define ECHO_PORT      9999
#define TEST_SIZE      (1024 * 1024)
#define POLL_TIMEOUT_MS 10000
#define DATA_TIMEOUT_MS 120000

static int g_test_ok = 0;
static pid_t echo_pid = 0;

static void echo_server(uint16_t port) {
    int srv = socket(AF_INET, SOCK_STREAM, 0);
    if(srv < 0) { perror("echo socket"); exit(1); }
    int opt = 1; setsockopt(srv, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    struct sockaddr_in addr = {.sin_family = AF_INET, .sin_port = htons(port)}; addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    if(bind(srv, (struct sockaddr*)&addr, sizeof(addr)) < 0) { perror("echo bind"); close(srv); exit(1); }
    if(listen(srv, 1) < 0) { perror("echo listen"); close(srv); exit(1); }
    int cli = accept(srv, NULL, NULL);
    if(cli < 0) { perror("echo accept"); close(srv); exit(1); }
    uint8_t buf[65536]; ssize_t n;
    while((n = recv(cli, buf, sizeof(buf), 0)) > 0) { ssize_t sent = 0; while(sent < n) { ssize_t s = send(cli, buf + sent, n - sent, 0); if(s < 0) goto done; sent += s; } }
done: close(cli); close(srv);
}

static void run_instance_b(int ip_fd) {
    struct UASYNC* ua = uasync_create();
    if(!ua) { fprintf(stderr, "B: uasync_create failed\n"); _exit(1); }
    struct tcp_proxy_mapping_config m = {.local_port = TEST_PORT, .remote_ip = "127.0.0.1", .remote_port = ECHO_PORT};
    struct tcp_proxy* b = tcp_proxy_create(NULL, ua, NULL, NULL, 0, 0, &m, 1, 0, 0, ip_fd, 0);
    if(!b) { fprintf(stderr, "B: tcp_proxy_create failed\n"); uasync_destroy(ua, 0); _exit(1); }
    while(1) uasync_poll(ua, 100);
}

int main(void) {
    // 1. Fork echo server
    echo_pid = fork();
    if(echo_pid == 0) { echo_server(ECHO_PORT); _exit(0); }
    if(echo_pid < 0) { perror("fork echo"); return 1; }
    usleep(100000);

    // 2. Socketpair
    int pair[2];
    if(socketpair(AF_UNIX, SOCK_STREAM, 0, pair) < 0) { perror("socketpair"); kill(echo_pid, SIGTERM); waitpid(echo_pid, NULL, 0); return 1; }

    // 3. Fork child = Instance B
    pid_t child = fork();
    if(child == 0) { close(pair[0]); run_instance_b(pair[1]); _exit(0); }
    if(child < 0) { perror("fork child"); close(pair[0]); close(pair[1]); kill(echo_pid, SIGTERM); waitpid(echo_pid, NULL, 0); return 1; }
    close(pair[1]);

    // 4. Parent = Instance A
    struct UASYNC* ua = uasync_create();
    if(!ua) { printf("[FAIL] uasync_create\n"); close(pair[0]); kill(child, SIGTERM); kill(echo_pid, SIGTERM); waitpid(child, NULL, 0); waitpid(echo_pid, NULL, 0); return 1; }

    struct tcp_proxy* a = tcp_proxy_create(NULL, ua, NULL, NULL, 0, 0, NULL, 0, 0, 0, pair[0], 0);
    if(!a) { printf("[FAIL] tcp_proxy_create\n"); uasync_destroy(ua, 0); close(pair[0]); kill(child, SIGTERM); kill(echo_pid, SIGTERM); waitpid(child, NULL, 0); waitpid(echo_pid, NULL, 0); return 1; }

    // 5. Active open
    struct proxy_conn* pc = tcp_proxy_active_open(a, "10.0.0.1", TEST_PORT);
    if(!pc) { printf("[FAIL] active_open\n"); tcp_proxy_destroy(a); uasync_destroy(ua, 0); close(pair[0]); kill(child, SIGTERM); kill(echo_pid, SIGTERM); waitpid(child, NULL, 0); waitpid(echo_pid, NULL, 0); return 1; }

    // 6. Wait for handshake
    int timeout_ms = POLL_TIMEOUT_MS;
    while(pc->pcb && pc->pcb->state != ESTABLISHED && timeout_ms > 0) {
        uasync_poll(ua, 100); timeout_ms -= 10;
        if(!pc->pcb || pc->pcb->state == CLOSED) break;
    }
    if(!pc->pcb || pc->pcb->state != ESTABLISHED) {
        printf("[FAIL] handshake\n");
        tcp_proxy_destroy(a); uasync_destroy(ua, 0); close(pair[0]); kill(child, SIGTERM); kill(echo_pid, SIGTERM); waitpid(child, NULL, 0); waitpid(echo_pid, NULL, 0); return 1;
    }

    // 7. Generate 1MB data
    uint8_t* send_buf = malloc(TEST_SIZE); uint8_t* recv_buf = malloc(TEST_SIZE);
    if(!send_buf || !recv_buf) { printf("[FAIL] malloc\n"); tcp_proxy_destroy(a); uasync_destroy(ua, 0); close(pair[0]); kill(child, SIGTERM); kill(echo_pid, SIGTERM); waitpid(child, NULL, 0); waitpid(echo_pid, NULL, 0); return 1; }
    srand(time(NULL)); int k; for(k = 0; k < TEST_SIZE; k++) send_buf[k] = (uint8_t)(rand() & 0xFF);

    // 8. Queue data for sending
    size_t off;
    for(off = 0; off < TEST_SIZE; ) { size_t chunk = TEST_SIZE - off; if(chunk > 1460) chunk = 1460;
        if(tcp_proxy_active_send(a, pc, send_buf + off, chunk) != 0) { printf("[FAIL] active_send at %zu\n", off); goto fail; }
        off += chunk; }

    // 9. Wait until all queued data is sent
    timeout_ms = DATA_TIMEOUT_MS;
    while(!tcp_proxy_active_send_done(a, pc) && timeout_ms > 0) {
        uasync_poll(ua, 100); timeout_ms -= 10;
    }
    if(!tcp_proxy_active_send_done(a, pc)) { printf("[FAIL] send timeout\n"); goto fail; }

    // 10. Half-close: all data sent, now signal FIN to echo server
    tcp_proxy_active_close(a, pc);

    // 11. Receive echoed data back
    size_t total_rcvd = 0; timeout_ms = DATA_TIMEOUT_MS;
    while(total_rcvd < TEST_SIZE && timeout_ms > 0) {
        uasync_poll(ua, 100); timeout_ms -= 10;
        ssize_t n = tcp_proxy_active_recv(a, pc, recv_buf + total_rcvd, TEST_SIZE - total_rcvd);
        if(n > 0) total_rcvd += n;
        if(pc->closing && total_rcvd < TEST_SIZE) break;
    }

    // 12. Verify
    if(total_rcvd == TEST_SIZE && memcmp(send_buf, recv_buf, TEST_SIZE) == 0) {
        printf("[PASS] test_tcp_proxy — 1MB echo verified\n"); g_test_ok = 1;
    } else {
        printf("[FAIL] test_tcp_proxy — %zu/%d bytes\n", total_rcvd, TEST_SIZE);
        if(total_rcvd == TEST_SIZE) for(k = 0; k < TEST_SIZE; k++) if(send_buf[k] != recv_buf[k]) { printf("  diff at %d: %02x/%02x\n", k, send_buf[k], recv_buf[k]); break; }
    }

fail:
    free(send_buf); free(recv_buf);
    tcp_proxy_destroy(a); uasync_destroy(ua, 0); close(pair[0]);
    kill(child, SIGTERM); waitpid(child, NULL, 0);
    kill(echo_pid, SIGTERM); waitpid(echo_pid, NULL, 0);
    return g_test_ok ? 0 : 1;
}