/* SPDX-License-Identifier: GPL-3.0-only */
#include "htp_n3n.h"
#include "n3n_path.h"
#include "htp_packet_bridge.h"

#include <n3n/edge.h>
#include <n3n/initfuncs.h>
#include <n3n/mainloop.h>
#include <n3n/resolve.h>
#include <arpa/inet.h>
#include <errno.h>
#include <poll.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>
#include "n2n.h"

#define HTP_FRAME_CAPACITY (65535U + HTP_ETHERNET_HEADER_SIZE)

struct htp_n3n_edge {
    htp_n3n_config config;
    htp_n3n_callbacks callbacks;
    char *supernode;
    char *community;
    char *key;
    char *local_ipv4;
    char *peer_ipv4;
    uint32_t local_ipv4_host;
    uint32_t peer_ipv4_host;
    uint8_t local_mac[HTP_MAC_SIZE];
    int tap_fds[2];
    pthread_t thread;
    pthread_mutex_t state_mutex;
    pthread_mutex_t bridge_mutex;
    htp_packet_bridge bridge;
    struct n3n_runtime_data *runtime;
    bool keep_running;
    bool thread_started;
    bool joined;
    bool mutexes_initialized;
    htp_n3n_state state;
    htp_n3n_peer_path peer_path;
};

static pthread_mutex_t callback_mutex = PTHREAD_MUTEX_INITIALIZER;
static htp_n3n_edge *callback_edge;

static char *copy_string(const char *value) {
    return value ? strdup(value) : NULL;
}

static void set_state(htp_n3n_edge *edge, htp_n3n_state state, const char *detail) {
    htp_n3n_state_callback callback = NULL;
    void *context = NULL;
    pthread_mutex_lock(&edge->state_mutex);
    if (edge->state != state) {
        edge->state = state;
        callback = edge->callbacks.state_changed;
        context = edge->callbacks.context;
    }
    pthread_mutex_unlock(&edge->state_mutex);
    if (callback) callback(context, state, detail ? detail : "");
}

static const char *transport_name(htp_n3n_transport transport) {
    switch (transport) {
        case HTP_N3N_TRANSPORT_UDP_RELAY:
            return "udp-relay";
        case HTP_N3N_TRANSPORT_TCP_RELAY:
            return "tcp-relay";
        case HTP_N3N_TRANSPORT_UDP_P2P:
        default:
            return "udp-p2p";
    }
}

static void set_peer_path(htp_n3n_edge *edge, htp_n3n_peer_path path) {
    pthread_mutex_lock(&edge->state_mutex);
    edge->peer_path = path;
    pthread_mutex_unlock(&edge->state_mutex);
}

static void registration_updated(struct n3n_runtime_data *runtime, time_t now, const n3n_sock_t *supernode) {
    (void)runtime;
    (void)now;
    (void)supernode;
    pthread_mutex_lock(&callback_mutex);
    htp_n3n_edge *edge = callback_edge;
    if (edge) set_state(edge, HTP_N3N_CONNECTED, transport_name(edge->config.transport));
    pthread_mutex_unlock(&callback_mutex);
}

static void main_loop_period(struct n3n_runtime_data *runtime, time_t now) {
    (void)now;
    pthread_mutex_lock(&callback_mutex);
    htp_n3n_edge *edge = callback_edge;
    if (!edge || edge->runtime != runtime || edge->peer_ipv4_host == 0U) {
        pthread_mutex_unlock(&callback_mutex);
        return;
    }

    set_peer_path(edge, htp_n3n_classify_peer_path(
        edge->config.transport,
        htp_n3n_get_state(edge) == HTP_N3N_CONNECTED,
        runtime->last_p2p,
        runtime->last_sup));
    pthread_mutex_unlock(&callback_mutex);
}

static void socket_opened(struct n3n_runtime_data *runtime) {
    pthread_mutex_lock(&callback_mutex);
    htp_n3n_edge *edge = callback_edge;
    if (edge && edge->callbacks.protect_socket && runtime->sock >= 0) {
        edge->callbacks.protect_socket(edge->callbacks.context, runtime->sock);
    }
    pthread_mutex_unlock(&callback_mutex);
}

static int parse_ipv4(const char *text, uint32_t *host_order) {
    struct in_addr address;
    if (!text || inet_pton(AF_INET, text, &address) != 1) return -1;
    *host_order = ntohl(address.s_addr);
    return 0;
}

static void make_mac(uint32_t ipv4, uint8_t mac[HTP_MAC_SIZE]) {
    mac[0] = 0x02U;
    mac[1] = 0x48U;
    mac[2] = 0x54U;
    mac[3] = (uint8_t)(ipv4 >> 16U);
    mac[4] = (uint8_t)(ipv4 >> 8U);
    mac[5] = (uint8_t)ipv4;
}

static void mac_string(const uint8_t mac[HTP_MAC_SIZE], char output[18]) {
    snprintf(output, 18U, "%02X:%02X:%02X:%02X:%02X:%02X",
        mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
}

static int configure_runtime(htp_n3n_edge *edge, n2n_edge_conf_t *conf) {
    edge_init_conf_defaults(conf, "ht-mobile-vpn");
    conf->allow_routing = true;
    conf->allow_multicast = false;
    conf->allow_p2p = edge->config.transport == HTP_N3N_TRANSPORT_UDP_P2P;
    conf->connect_tcp = edge->config.transport == HTP_N3N_TRANSPORT_TCP_RELAY;
    conf->register_interval = 5U;
    conf->register_ttl = 0U;
    conf->transop_id = N2N_TRANSFORM_ID_AES;
    conf->tuntap_ip_mode = TUNTAP_IP_MODE_STATIC;
    conf->tuntap_v4.net_addr = htonl(edge->local_ipv4_host);
    conf->tuntap_v4.net_bitlen = edge->config.prefix_length;
    conf->mtu = edge->config.mtu;
#ifndef _WIN32
    conf->userid = getuid();
    conf->groupid = getgid();
#endif
    snprintf((char *)conf->community_name, sizeof(conf->community_name), "%s", edge->community);
    mac_string(edge->local_mac, conf->device_mac);
    conf->encrypt_key = strdup(edge->key);
    if (!conf->encrypt_key) return -1;
    resolve_hostnames_str_add(RESOLVE_LIST_SUPERNODE, edge->supernode);
    return edge_verify_conf(conf);
}

static void *edge_thread(void *context) {
    htp_n3n_edge *edge = context;
    n2n_edge_conf_t conf;
    memset(&conf, 0, sizeof(conf));
    int result = -1;
    n3n_initfuncs();
    if (configure_runtime(edge, &conf) != 0) {
        edge_term_conf(&conf);
        n3n_deinitfuncs();
        set_state(edge, HTP_N3N_FAILED, "invalid n3n configuration");
        return NULL;
    }

    struct n3n_runtime_data *runtime = edge_init(&conf, &result);
    if (!runtime) {
        n3n_deinitfuncs();
        set_state(edge, HTP_N3N_FAILED, "n3n initialization failed");
        return NULL;
    }
    runtime->device.fd = edge->tap_fds[0];
    runtime->device.ip_addr = htonl(edge->local_ipv4_host);
    runtime->device.mtu = edge->config.mtu;
    memcpy(runtime->device.mac_addr, edge->local_mac, HTP_MAC_SIZE);
    snprintf(runtime->device.dev_name, sizeof(runtime->device.dev_name), "%s", "ht-mobile-vpn");
    runtime->keep_running = &edge->keep_running;
    runtime->embed_callbacks.sn_registration_updated = registration_updated;
    runtime->embed_callbacks.main_loop_period = main_loop_period;
    runtime->embed_callbacks.sock_opened = socket_opened;
    mainloop_register_fd(runtime->device.fd, fd_info_proto_tuntap);
    edge->runtime = runtime;

    pthread_mutex_lock(&callback_mutex);
    callback_edge = edge;
    if (edge->callbacks.protect_socket && runtime->sock >= 0) {
        edge->callbacks.protect_socket(edge->callbacks.context, runtime->sock);
    }
    pthread_mutex_unlock(&callback_mutex);

    result = run_edge_loop(runtime);

    pthread_mutex_lock(&callback_mutex);
    if (callback_edge == edge) callback_edge = NULL;
    pthread_mutex_unlock(&callback_mutex);
    edge->runtime = NULL;
    mainloop_unregister_fd(runtime->device.fd);
    edge_term(runtime);
    n3n_deinitfuncs();

    if (htp_n3n_get_state(edge) == HTP_N3N_STOPPING || !edge->keep_running) {
        set_state(edge, HTP_N3N_STOPPED, "stopped");
    } else if (result != 0) {
        set_state(edge, HTP_N3N_FAILED, "n3n edge loop failed");
    } else {
        set_state(edge, HTP_N3N_STOPPED, "stopped");
    }
    return NULL;
}

htp_n3n_edge *htp_n3n_create(const htp_n3n_config *config, const htp_n3n_callbacks *callbacks) {
    if (!config || !config->supernode || !config->community || !config->key ||
        !config->local_ipv4 || config->prefix_length == 0U || config->prefix_length > 32U) return NULL;
    htp_n3n_edge *edge = calloc(1U, sizeof(*edge));
    if (!edge) return NULL;
    edge->tap_fds[0] = -1;
    edge->tap_fds[1] = -1;
    pthread_mutex_init(&edge->state_mutex, NULL);
    pthread_mutex_init(&edge->bridge_mutex, NULL);
    edge->mutexes_initialized = true;
    edge->supernode = copy_string(config->supernode);
    edge->community = copy_string(config->community);
    edge->key = copy_string(config->key);
    edge->local_ipv4 = copy_string(config->local_ipv4);
    edge->peer_ipv4 = copy_string(config->peer_ipv4);
    if (!edge->supernode || !edge->community || !edge->key || !edge->local_ipv4 ||
        parse_ipv4(edge->local_ipv4, &edge->local_ipv4_host) != 0 ||
        (edge->peer_ipv4 && parse_ipv4(edge->peer_ipv4, &edge->peer_ipv4_host) != 0)) {
        htp_n3n_destroy(edge);
        return NULL;
    }
    edge->config = *config;
    edge->config.supernode = edge->supernode;
    edge->config.community = edge->community;
    edge->config.key = edge->key;
    edge->config.local_ipv4 = edge->local_ipv4;
    edge->config.peer_ipv4 = edge->peer_ipv4;
    if (edge->config.transport < HTP_N3N_TRANSPORT_UDP_P2P ||
        edge->config.transport > HTP_N3N_TRANSPORT_TCP_RELAY) {
        edge->config.transport = HTP_N3N_TRANSPORT_UDP_P2P;
    }
    if (edge->config.mtu == 0U) edge->config.mtu = 1290U;
    if (callbacks) edge->callbacks = *callbacks;
    make_mac(edge->local_ipv4_host, edge->local_mac);
    htp_packet_bridge_init(&edge->bridge, edge->local_ipv4_host, edge->local_mac);
    edge->state = HTP_N3N_IDLE;
    edge->peer_path = HTP_N3N_PEER_PATH_UNKNOWN;
    return edge;
}

int htp_n3n_start(htp_n3n_edge *edge) {
    if (!edge || edge->thread_started || socketpair(AF_UNIX, SOCK_DGRAM, 0, edge->tap_fds) != 0) return -1;
    edge->keep_running = true;
    edge->thread_started = true;
    set_peer_path(edge, HTP_N3N_PEER_PATH_UNKNOWN);
    set_state(edge, HTP_N3N_CONNECTING, transport_name(edge->config.transport));
    if (pthread_create(&edge->thread, NULL, edge_thread, edge) != 0) {
        edge->thread_started = false;
        close(edge->tap_fds[0]);
        close(edge->tap_fds[1]);
        edge->tap_fds[0] = edge->tap_fds[1] = -1;
        set_state(edge, HTP_N3N_FAILED, "unable to start n3n thread");
        return -1;
    }
    return 0;
}

int htp_n3n_send_ip_packet(htp_n3n_edge *edge, const uint8_t *packet, size_t packet_length) {
    if (!edge || edge->tap_fds[1] < 0 || packet_length > 65535U) return -1;
    uint8_t frame[HTP_FRAME_CAPACITY];
    pthread_mutex_lock(&edge->bridge_mutex);
    const int frame_length = htp_packet_bridge_ip_to_ethernet(&edge->bridge, packet, packet_length, frame, sizeof(frame));
    pthread_mutex_unlock(&edge->bridge_mutex);
    if (frame_length < 0) return -1;
    return send(edge->tap_fds[1], frame, (size_t)frame_length, 0) == frame_length ? 0 : -1;
}

int htp_n3n_receive_ip_packet(
    htp_n3n_edge *edge,
    uint8_t *packet,
    size_t packet_capacity,
    size_t *packet_length,
    int timeout_ms) {
    if (packet_length) *packet_length = 0U;
    if (!edge || edge->tap_fds[1] < 0 || !packet || !packet_length) return -1;
    struct pollfd descriptor = {.fd = edge->tap_fds[1], .events = POLLIN, .revents = 0};
    const int ready = poll(&descriptor, 1U, timeout_ms);
    if (ready <= 0) return ready;
    uint8_t frame[HTP_FRAME_CAPACITY];
    const ssize_t frame_length = recv(edge->tap_fds[1], frame, sizeof(frame), 0);
    if (frame_length <= 0) return -1;
    uint8_t reply[64];
    size_t reply_length = 0U;
    pthread_mutex_lock(&edge->bridge_mutex);
    const htp_bridge_result result = htp_packet_bridge_ethernet_to_ip(
        &edge->bridge, frame, (size_t)frame_length, packet, packet_capacity, packet_length,
        reply, sizeof(reply), &reply_length);
    pthread_mutex_unlock(&edge->bridge_mutex);
    if (result == HTP_BRIDGE_CONTROL_REPLY && reply_length > 0U) {
        send(edge->tap_fds[1], reply, reply_length, 0);
        return 0;
    }
    return result == HTP_BRIDGE_IP_PACKET ? 1 : (result == HTP_BRIDGE_ERROR ? -1 : 0);
}

int htp_n3n_resolve_peer(htp_n3n_edge *edge, const char *peer_ipv4, int timeout_ms) {
    uint32_t peer_ipv4_host;
    if (!edge || edge->tap_fds[1] < 0 || timeout_ms <= 0 ||
        parse_ipv4(peer_ipv4, &peer_ipv4_host) != 0) {
        return -1;
    }

    const int interval_ms = 250;
    int elapsed_ms = 0;
    uint8_t request[HTP_ETHERNET_HEADER_SIZE + 28U];
    uint8_t packet[2048];

    while (elapsed_ms < timeout_ms) {
        pthread_mutex_lock(&edge->bridge_mutex);
        const int resolved = htp_packet_bridge_has_peer(&edge->bridge, peer_ipv4_host);
        const int request_length = resolved ? 0 : htp_packet_bridge_build_arp_request(
            &edge->bridge, peer_ipv4_host, request, sizeof(request));
        pthread_mutex_unlock(&edge->bridge_mutex);
        if (resolved) return 1;
        if (request_length <= 0 ||
            send(edge->tap_fds[1], request, (size_t)request_length, 0) != request_length) {
            return -1;
        }

        const int wait_ms = (timeout_ms - elapsed_ms) < interval_ms
            ? timeout_ms - elapsed_ms
            : interval_ms;
        size_t packet_length = 0U;
        const int receive_result = htp_n3n_receive_ip_packet(
            edge, packet, sizeof(packet), &packet_length, wait_ms);
        if (receive_result < 0 && htp_n3n_get_state(edge) != HTP_N3N_CONNECTED) return -1;
        elapsed_ms += wait_ms;

        pthread_mutex_lock(&edge->bridge_mutex);
        const int learned = htp_packet_bridge_has_peer(&edge->bridge, peer_ipv4_host);
        pthread_mutex_unlock(&edge->bridge_mutex);
        if (learned) return 1;
    }

    return 0;
}

htp_n3n_state htp_n3n_get_state(const htp_n3n_edge *edge) {
    if (!edge) return HTP_N3N_FAILED;
    pthread_mutex_lock((pthread_mutex_t *)&edge->state_mutex);
    const htp_n3n_state state = edge->state;
    pthread_mutex_unlock((pthread_mutex_t *)&edge->state_mutex);
    return state;
}

htp_n3n_peer_path htp_n3n_get_peer_path(const htp_n3n_edge *edge) {
    if (!edge) return HTP_N3N_PEER_PATH_UNKNOWN;
    pthread_mutex_lock((pthread_mutex_t *)&edge->state_mutex);
    const htp_n3n_peer_path path = edge->peer_path;
    pthread_mutex_unlock((pthread_mutex_t *)&edge->state_mutex);
    return path;
}

void htp_n3n_stop(htp_n3n_edge *edge) {
    if (!edge || !edge->thread_started || edge->joined) return;
    set_state(edge, HTP_N3N_STOPPING, "stopping");
    edge->keep_running = false;
    if (edge->tap_fds[1] >= 0) {
        const uint8_t wake_frame[HTP_ETHERNET_HEADER_SIZE] = {0};
        send(edge->tap_fds[1], wake_frame, sizeof(wake_frame), 0);
    }
    pthread_join(edge->thread, NULL);
    edge->joined = true;
}

void htp_n3n_destroy(htp_n3n_edge *edge) {
    if (!edge) return;
    htp_n3n_stop(edge);
    if (edge->tap_fds[0] >= 0) close(edge->tap_fds[0]);
    if (edge->tap_fds[1] >= 0) close(edge->tap_fds[1]);
    if (edge->mutexes_initialized) {
        pthread_mutex_destroy(&edge->state_mutex);
        pthread_mutex_destroy(&edge->bridge_mutex);
    }
    free(edge->supernode);
    free(edge->community);
    free(edge->key);
    free(edge->local_ipv4);
    free(edge->peer_ipv4);
    free(edge);
}

const char *htp_n3n_version(void) {
    return "3.4.6+htportal-mobile";
}
