#include "omnisocket_client.h"

int omnisocket_session_init(omnisocket_session_t *session) {
    int rc;

    if (session == NULL) {
        errno = EINVAL;
        return -1;
    }
    memset(session, 0, sizeof(*session));
    rc = pthread_mutex_init(&session->mutex, NULL);
    if (rc != 0) {
        errno = rc;
        return -1;
    }
    rc = pthread_cond_init(&session->idle_cond, NULL);
    if (rc != 0) {
        pthread_mutex_destroy(&session->mutex);
        errno = rc;
        return -1;
    }
    return 0;
}

void omnisocket_session_destroy(omnisocket_session_t *session) {
    if (session == NULL) {
        return;
    }
    (void) omnisocket_session_close(session);
    pthread_cond_destroy(&session->idle_cond);
    pthread_mutex_destroy(&session->mutex);
}

static int omnisocket_session_begin_client_op(omnisocket_session_t *session, kcp_client_t **out_client) {
    if (session == NULL || out_client == NULL) {
        errno = EINVAL;
        return -1;
    }

    pthread_mutex_lock(&session->mutex);
    if (session->closing) {
        pthread_mutex_unlock(&session->mutex);
        errno = ECANCELED;
        return -1;
    }
    if (session->client == NULL) {
        pthread_mutex_unlock(&session->mutex);
        errno = ENOTCONN;
        return -1;
    }
    *out_client = session->client;
    session->active_ops += 1;
    pthread_mutex_unlock(&session->mutex);
    return 0;
}

int omnisocket_session_connect(
    omnisocket_session_t *session,
    const char *server_addr,
    const char *relay_via,
    const char *peer_id,
    const char *bind_ip,
    const char *bind_device,
    const kcp_conn_options_t *options,
    int stats_interval_ms
) {
    kcp_client_t *client;

    if (session == NULL || server_addr == NULL || peer_id == NULL) {
        errno = EINVAL;
        return -1;
    }

    pthread_mutex_lock(&session->mutex);
    while (session->closing) {
        pthread_cond_wait(&session->idle_cond, &session->mutex);
    }
    if (session->client != NULL) {
        pthread_mutex_unlock(&session->mutex);
        errno = EISCONN;
        return -1;
    }
    client = kcp_client_dial_with_options(
        server_addr,
        relay_via,
        peer_id,
        bind_ip,
        bind_device,
        options,
        NULL,
        NULL,
        NULL,
        stats_interval_ms
    );
    if (client == NULL) {
        pthread_mutex_unlock(&session->mutex);
        return -1;
    }
    session->client = client;
    session->stats.connected = 1;
    pthread_mutex_unlock(&session->mutex);
    return 0;
}

int omnisocket_session_close(omnisocket_session_t *session) {
    kcp_client_t *client;

    if (session == NULL) {
        errno = EINVAL;
        return -1;
    }

    pthread_mutex_lock(&session->mutex);
    while (session->closing) {
        pthread_cond_wait(&session->idle_cond, &session->mutex);
    }
    client = session->client;
    if (client != NULL) {
        session->closing = 1;
        session->client = NULL;
    }
    session->stats.connected = 0;
    pthread_mutex_unlock(&session->mutex);

    if (client != NULL) {
        kcp_client_close(client);
        pthread_mutex_lock(&session->mutex);
        while (session->active_ops > 0) {
            pthread_cond_wait(&session->idle_cond, &session->mutex);
        }
        pthread_mutex_unlock(&session->mutex);
        kcp_client_free(client);
        pthread_mutex_lock(&session->mutex);
        session->closing = 0;
        pthread_cond_broadcast(&session->idle_cond);
        pthread_mutex_unlock(&session->mutex);
    }
    return 0;
}

int omnisocket_session_send(omnisocket_session_t *session, const char *to, const void *data, size_t data_len) {
    kcp_client_t *client;
    int rc;

    if (session == NULL || to == NULL || (data == NULL && data_len > 0)) {
        errno = EINVAL;
        return -1;
    }

    if (omnisocket_session_begin_client_op(session, &client) != 0) {
        return -1;
    }
    rc = kcp_client_send_binary(client, to, data, data_len);
    pthread_mutex_lock(&session->mutex);
    if (rc == 0) {
        session->stats.send_calls += 1;
        session->stats.send_bytes += (uint64_t) data_len;
    } else {
        session->stats.send_errors += 1;
    }
    if (session->active_ops > 0) {
        session->active_ops -= 1;
    }
    if (session->closing && session->active_ops == 0) {
        pthread_cond_broadcast(&session->idle_cond);
    }
    pthread_mutex_unlock(&session->mutex);
    return rc;
}

int omnisocket_session_recv(omnisocket_session_t *session, message_t *out_msg, int timeout_ms) {
    kcp_client_t *client;
    int rc;

    if (session == NULL || out_msg == NULL) {
        errno = EINVAL;
        return -1;
    }

    if (omnisocket_session_begin_client_op(session, &client) != 0) {
        return -1;
    }
    rc = kcp_client_receive_timed(client, out_msg, timeout_ms);
    pthread_mutex_lock(&session->mutex);
    if (rc == 0) {
        session->stats.recv_calls += 1;
        session->stats.recv_bytes += (uint64_t) out_msg->body_len;
    } else if (rc == 1) {
        session->stats.recv_timeouts += 1;
    } else {
        session->stats.recv_errors += 1;
    }
    if (session->active_ops > 0) {
        session->active_ops -= 1;
    }
    if (session->closing && session->active_ops == 0) {
        pthread_cond_broadcast(&session->idle_cond);
    }
    pthread_mutex_unlock(&session->mutex);
    return rc;
}

int omnisocket_session_recv_into(
    omnisocket_session_t *session,
    void *buffer,
    size_t buffer_len,
    kcp_client_recv_meta_t *out_meta,
    int timeout_ms
) {
    kcp_client_t *client;
    int rc;

    if (session == NULL || out_meta == NULL || (buffer == NULL && buffer_len > 0)) {
        errno = EINVAL;
        return -1;
    }

    if (omnisocket_session_begin_client_op(session, &client) != 0) {
        return -1;
    }
    rc = kcp_client_receive_binary_into(client, buffer, buffer_len, out_meta, timeout_ms);
    pthread_mutex_lock(&session->mutex);
    if (rc == 0) {
        session->stats.recv_calls += 1;
        session->stats.recv_bytes += (uint64_t) out_meta->body_len;
    } else if (rc == 1) {
        session->stats.recv_timeouts += 1;
    } else {
        session->stats.recv_errors += 1;
    }
    if (session->active_ops > 0) {
        session->active_ops -= 1;
    }
    if (session->closing && session->active_ops == 0) {
        pthread_cond_broadcast(&session->idle_cond);
    }
    pthread_mutex_unlock(&session->mutex);
    return rc;
}

void omnisocket_session_stats_snapshot(omnisocket_session_t *session, omnisocket_session_stats_t *out_stats) {
    if (session == NULL || out_stats == NULL) {
        return;
    }
    pthread_mutex_lock(&session->mutex);
    *out_stats = session->stats;
    pthread_mutex_unlock(&session->mutex);
}

int omnisocket_session_kcp_metrics_snapshot(
    omnisocket_session_t *session,
    omnisocket_session_kcp_metrics_t *out_metrics
) {
    kcp_client_t *client = NULL;
    kcp_conn_metrics_t metrics;
    int rc = 0;

    if (session == NULL || out_metrics == NULL) {
        errno = EINVAL;
        return -1;
    }

    memset(out_metrics, 0, sizeof(*out_metrics));

    pthread_mutex_lock(&session->mutex);
    if (session->client != NULL && !session->closing) {
        client = session->client;
        session->active_ops += 1;
    }
    pthread_mutex_unlock(&session->mutex);

    if (client == NULL) {
        return 0;
    }

    memset(&metrics, 0, sizeof(metrics));
    rc = kcp_client_metrics_snapshot(client, &metrics);

    pthread_mutex_lock(&session->mutex);
    if (session->active_ops > 0) {
        session->active_ops -= 1;
    }
    if (session->closing && session->active_ops == 0) {
        pthread_cond_broadcast(&session->idle_cond);
    }
    pthread_mutex_unlock(&session->mutex);

    if (rc != 0) {
        return rc;
    }

    out_metrics->connected = metrics.connected;
    out_metrics->has_conv = metrics.has_conv;
    out_metrics->conv = metrics.conv;
    snprintf(out_metrics->local_addr, sizeof(out_metrics->local_addr), "%s", metrics.local_addr);
    snprintf(out_metrics->remote_addr, sizeof(out_metrics->remote_addr), "%s", metrics.remote_addr);
    out_metrics->rto_ms = metrics.rto_ms;
    out_metrics->srtt_ms = metrics.srtt_ms;
    out_metrics->srttvar_ms = metrics.srttvar_ms;
    out_metrics->bytes_sent = metrics.bytes_sent;
    out_metrics->bytes_received = metrics.bytes_received;
    out_metrics->in_pkts = metrics.in_pkts;
    out_metrics->out_pkts = metrics.out_pkts;
    out_metrics->in_segs = metrics.in_segs;
    out_metrics->out_segs = metrics.out_segs;
    out_metrics->retrans_segs = metrics.retrans_segs;
    out_metrics->fast_retrans_segs = metrics.fast_retrans_segs;
    out_metrics->early_retrans_segs = metrics.early_retrans_segs;
    out_metrics->lost_segs = metrics.lost_segs;
    out_metrics->repeat_segs = metrics.repeat_segs;
    out_metrics->in_errs = metrics.in_errs;
    out_metrics->kcp_in_errs = metrics.kcp_in_errs;
    out_metrics->ring_buffer_snd_queue = metrics.ring_buffer_snd_queue;
    out_metrics->ring_buffer_rcv_queue = metrics.ring_buffer_rcv_queue;
    out_metrics->ring_buffer_snd_buffer = metrics.ring_buffer_snd_buffer;
    return 0;
}