// stcp_client.c — STCP client implementation
#include "stcp_client.h"
#include "secure_channel.h"
#include "crc32.h"
#include "../lib/u_async.h"
#include "../lib/socket_compat.h"
#include "../lib/ll_queue.h"
#include "../lib/mem.h"
#include "../lib/debug_config.h"
#include "../lib/platform_compat.h"
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#ifndef _WIN32
#include <unistd.h>
#include <fcntl.h>
#include <netinet/tcp.h>
#endif

struct stcp_client {
    struct stcp_conn conn;
    struct UASYNC *ua;
    stcp_ready_cb ready_cb;
    void *ready_arg;
    uint8_t peer_pubkey[SC_PUBKEY_SIZE];
    uint8_t my_ed25519_pubkey[SC_PUBKEY_SIZE];
};

static void client_connect_write_cb(socket_t sock, void *arg);
static void client_conn_read_cb(socket_t sock, void *arg);
static void client_hs_cb(struct stcp_conn *c, uint8_t *data, size_t len);
static void client_hs_padding_cb(struct stcp_conn *c, uint8_t *data, size_t len);
static void client_data_cb(struct stcp_conn *c, uint8_t *plain_data, size_t data_len);

static int client_derive_session(struct stcp_conn *c, const uint8_t *peer_pubkey) {
    struct secure_channel sc;
    sc_init_ctx(&sc, &c->my_keys);
    if (sc_set_peer_public_key(&sc, peer_pubkey, SC_PEER_PUBKEY_BIN) != SC_OK) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "client ECDH failed"); return -1; }
    memcpy(c->session_key, sc.session_key, SC_SESSION_KEY_SIZE);
    memcpy(c->peer_pubkey, peer_pubkey, SC_PUBKEY_SIZE); c->peer_pubkey_set = 1;
    log_dump(DEBUG_LEVEL_DEBUG, DEBUG_CATEGORY_CRYPTO, "stcp_client session_key", c->session_key, 16);
    if (sc_stream_init(&sc, &c->stream_send, STCP_STREAM_CLIENT_SEND) != SC_OK) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "client stream_send init failed"); return -1; }
    if (sc_stream_init(&sc, &c->stream_recv, STCP_STREAM_SERVER_SEND) != SC_OK) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "client stream_recv init failed"); return -1; }
    return 0;
}

static void client_send_handshake(struct stcp_conn *c, const uint8_t *server_pubkey, const uint8_t *my_ed25519) {
    uint8_t salt[SC_PUBKEY_ENC_SALT_SIZE];
    if (random_bytes(salt, SC_PUBKEY_ENC_SALT_SIZE) != 0) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "client random_bytes failed"); stcp_conn_do_close(c, 1); return; }
    uint16_t padding = 8;
    size_t total = SC_PUBKEY_ENC_SIZE + STCP_HS_ENC_CLIENT + padding;
    uint8_t *hs = u_malloc(total);
    if (!hs) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "client hs malloc failed"); stcp_conn_do_close(c, 1); return; }
    memcpy(hs, salt, SC_PUBKEY_ENC_SALT_SIZE);
    sc_obfuscate_pubkey(salt, server_pubkey, c->my_keys.public_key, hs + SC_PUBKEY_ENC_SALT_SIZE);

    uint8_t plain[42]; memcpy(plain, my_ed25519, 32);
    plain[32] = c->got_initial_pkt;
    memcpy(plain + 33, &c->session_id, 4);
    plain[37] = (uint8_t)padding; plain[38] = (uint8_t)(padding >> 8);
    plain[39] = c->device_type;
    *(uint16_t*)(plain + 40) = htobe16(c->keepalive_interval);
    uint32_t crc = crc32_calc(plain, 42);
    uint8_t *enc_dst = hs + SC_PUBKEY_ENC_SIZE;
    memcpy(enc_dst, plain, 42);
    enc_dst[42] = (uint8_t)(crc >> 0); enc_dst[43] = (uint8_t)(crc >> 8); enc_dst[44] = (uint8_t)(crc >> 16); enc_dst[45] = (uint8_t)(crc >> 24);
    if (sc_stream_xor(&c->stream_send, enc_dst, STCP_HS_ENC_CLIENT) != SC_OK) { u_free(hs); stcp_conn_do_close(c, 2); return; }
    for (int i = 0; i < padding; i++) hs[SC_PUBKEY_ENC_SIZE + STCP_HS_ENC_CLIENT + i] = (uint8_t)(salt[0] ^ i);

    c->state = STCP_STATE_HS_CLIENT_SENT;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "stcp_client: handshake sent (%zu bytes), entering HS_CLIENT_SENT", total);
    c->send_buf = hs; c->send_len = total; c->send_offset = 0;
    uasync_set_socket_write(c->ua, c->socket_id, 1);
}

static void client_hs_cb(struct stcp_conn *c, uint8_t *data, size_t len) {
    (void)len;
    log_dump(DEBUG_LEVEL_DEBUG, DEBUG_CATEGORY_ETCP, "stcp_client process_srv_resp recv_buf", data, len);

    const uint8_t *salt = data;
    const uint8_t *enc_pubkey = salt + SC_PUBKEY_ENC_SALT_SIZE;
    uint8_t server_pubkey[SC_PUBKEY_SIZE];
    sc_obfuscate_pubkey(salt, c->my_keys.public_key, enc_pubkey, server_pubkey);
    if (c->peer_pubkey_set && memcmp(server_pubkey, c->peer_pubkey, SC_PUBKEY_SIZE) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "stcp_client: server pubkey mismatch — possible MITM");
        stcp_conn_do_close(c, 4); return;
    }

    uint8_t enc_hs[STCP_HS_ENC_SERVER];
    memcpy(enc_hs, data + SC_PUBKEY_ENC_SIZE, STCP_HS_ENC_SERVER);
    log_dump(DEBUG_LEVEL_DEBUG, DEBUG_CATEGORY_CRYPTO, "stcp_client enc_hs BEFORE xor", enc_hs, STCP_HS_ENC_SERVER);
    size_t hs_data_len;
    if (stcp_frame_decrypt(enc_hs, STCP_HS_ENC_SERVER, &c->stream_recv, &hs_data_len)) { stcp_conn_do_close(c, 3); return; }
    log_dump(DEBUG_LEVEL_DEBUG, DEBUG_CATEGORY_CRYPTO, "stcp_client enc_hs AFTER xor", enc_hs, STCP_HS_ENC_SERVER);
    memcpy(c->peer_ed25519_pubkey, enc_hs, SC_PUBKEY_SIZE); c->peer_ed25519_set = 1;
    c->peer_got_initial_pkt = enc_hs[32];
    memcpy(&c->peer_session_id, enc_hs + 33, 4);
    uint16_t padding_size = (uint16_t)enc_hs[37] | ((uint16_t)enc_hs[38] << 8);
    c->peer_device_type = enc_hs[39];
    c->peer_keepalive_interval = ((uint16_t)enc_hs[40] << 8) | enc_hs[41];
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "stcp_client: server response OK gop=%d sid=%08x padding=%u dev=%d ka=%u",
               c->peer_got_initial_pkt, c->peer_session_id, padding_size, c->peer_device_type, c->peer_keepalive_interval);
    stcp_recv_set(c, padding_size, 0, client_hs_padding_cb);
}

static void client_hs_padding_cb(struct stcp_conn *c, uint8_t *data, size_t len) {
    (void)data; (void)len;
    if (c->hs_timer) { uasync_cancel_timeout(c->ua, c->hs_timer); c->hs_timer = NULL; }
    c->state = STCP_STATE_DATA;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "stcp_client: handshake OK, entering DATA state");
    if (c->on_ready) { void (*cb)(struct stcp_conn*, void*) = c->on_ready; c->on_ready = NULL; cb(c, c->ready_arg); }
    stcp_recv_set(c, 0, 1, client_data_cb);
}

static void client_data_cb(struct stcp_conn *c, uint8_t *plain_data, size_t data_len) {
    stcp_rx_push(c, plain_data, data_len);
}

static void client_conn_read_cb(socket_t sock, void *arg) {
    struct stcp_conn *c = (struct stcp_conn *)arg;
    (void)sock;
    int r = stcp_conn_read(c);
    if (r == -1) return;
    if (r == 0) return;
    stcp_recv_try(c);
}

static void client_connect_write_cb(socket_t sock, void *arg) {
    struct stcp_client *cli = (struct stcp_client *)arg;
    struct stcp_conn *c = &cli->conn;
    int err = 0;
    socklen_t len = sizeof(err);
    if (getsockopt(sock, SOL_SOCKET, SO_ERROR, (char *)&err, &len) < 0 || err != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "stcp_client connect failed err=%d", err);
        stcp_conn_do_close(c, err); return;
    }
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "stcp_client: TCP connected, starting handshake");
    uasync_remove_socket_t(cli->ua, sock);
    c->socket_id = uasync_add_socket_t(cli->ua, sock, client_conn_read_cb, stcp_write_cb, NULL, c);
    if (!c->socket_id) { stcp_conn_do_close(c, ENOMEM); return; }
    int opt = 1; setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (const char *)&opt, sizeof(opt));
    if (client_derive_session(c, cli->peer_pubkey)) { stcp_conn_do_close(c, 1); return; }
    client_send_handshake(c, cli->peer_pubkey, cli->my_ed25519_pubkey);
    // after handshake sent, wait for server response
    stcp_recv_set(c, SC_PUBKEY_ENC_SIZE + STCP_HS_ENC_SERVER, 0, client_hs_cb);
}

struct stcp_client *stcp_client_connect(struct UASYNC *ua, const char *addr, uint16_t port,
                                         struct SC_MYKEYS *keys, const uint8_t *peer_pubkey,
                                         const uint8_t *my_ed25519_pubkey,
                                         uint8_t got_initial_pkt, uint32_t session_id,
                                         uint8_t device_type, uint16_t keepalive_interval,
                                         stcp_ready_cb ready_cb, void *arg,
                                         stcp_close_cb close_cb, void *close_arg,
                                         const struct sockaddr_storage *local_addr) {
    if (!ua || !addr || !keys || !peer_pubkey || !ready_cb) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "invalid args"); return NULL; }
    struct stcp_client *cli = u_calloc(1, sizeof(struct stcp_client));
    if (!cli) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "calloc failed"); return NULL; }
    cli->ua = ua; cli->ready_cb = ready_cb; cli->ready_arg = arg;
    memcpy(cli->peer_pubkey, peer_pubkey, SC_PUBKEY_SIZE);
    if (my_ed25519_pubkey) memcpy(cli->my_ed25519_pubkey, my_ed25519_pubkey, SC_PUBKEY_SIZE);
    struct stcp_conn *c = &cli->conn;
    c->ua = ua; c->state = STCP_STATE_INIT; c->is_server = 0; c->my_keys = *keys;
    c->on_ready = ready_cb; c->ready_arg = arg;
    c->on_close = close_cb; c->close_arg = close_arg;
    c->on_write_error = stcp_conn_do_close;
    c->tx_cb = stcp_tx_queue_cb;
    c->got_initial_pkt = got_initial_pkt;
    c->session_id = session_id;
    c->device_type = device_type;
    c->keepalive_interval = keepalive_interval;

    struct addrinfo hints = {0};
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    char port_str[16]; snprintf(port_str, sizeof(port_str), "%u", port);
    struct addrinfo *res;
    if (getaddrinfo(addr, port_str, &hints, &res) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "stcp_client getaddrinfo %s:%u failed", addr, port);
        u_free(cli); return NULL;
    }
    c->sock = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (c->sock == SOCKET_INVALID) { freeaddrinfo(res); u_free(cli); return NULL; }
    socket_set_nonblocking(c->sock);
    if (local_addr && local_addr->ss_family) {
        if (bind(c->sock, (const struct sockaddr*)local_addr, local_addr->ss_family == AF_INET6 ? sizeof(struct sockaddr_in6) : sizeof(struct sockaddr_in)) < 0)
            DEBUG_WARN(DEBUG_CATEGORY_ETCP, "stcp_client bind to %s port=0 failed err=%d(%s), continuing to %s:%u",
                       sockaddr_storage_to_str(local_addr).str, socket_get_error(), socket_strerror(socket_get_error()), addr, port);
    }
    int conn_ret = connect(c->sock, res->ai_addr, res->ai_addrlen);
    freeaddrinfo(res);
    if (conn_ret < 0) {
        int sock_err = socket_get_error();
        if (sock_err != EINPROGRESS && sock_err != ERR_WOULDBLOCK) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "stcp_client connect to %s:%u failed err=%d(%s)", addr, port, sock_err, socket_strerror(sock_err));
            socket_close_wrapper(c->sock); u_free(cli); return NULL;
        }
        c->socket_id = uasync_add_socket_t(ua, c->sock, NULL, client_connect_write_cb, NULL, cli);
        if (!c->socket_id) { socket_close_wrapper(c->sock); u_free(cli); return NULL; }
        c->hs_timer = uasync_set_timeout(ua, STCP_CONNECT_TIMEOUT, c, hs_timeout_cb, "stcp_hs");
    } else {
        int opt = 1; setsockopt(c->sock, IPPROTO_TCP, TCP_NODELAY, (const char *)&opt, sizeof(opt));
        c->socket_id = uasync_add_socket_t(ua, c->sock, client_conn_read_cb, stcp_write_cb, NULL, c);
        if (!c->socket_id) { socket_close_wrapper(c->sock); u_free(cli); return NULL; }
        c->hs_timer = uasync_set_timeout(ua, STCP_CONNECT_TIMEOUT, c, hs_timeout_cb, "stcp_hs");
        if (client_derive_session(c, cli->peer_pubkey)) { c->free_on_close = cli; stcp_conn_do_close(c, 1); return NULL; }
        client_send_handshake(c, cli->peer_pubkey, cli->my_ed25519_pubkey);
        stcp_recv_set(c, SC_PUBKEY_ENC_SIZE + STCP_HS_ENC_SERVER, 0, client_hs_cb);
    }
    return cli;
}

void stcp_client_destroy(struct stcp_client *cli) {
    if (!cli) return;
    cli->conn.free_on_close = cli;
    stcp_conn_do_close(&cli->conn, 0);
}

struct stcp_conn *stcp_client_get_conn(struct stcp_client *cli) {
    return cli ? &cli->conn : NULL;
}