#define AL_LOG_SECTION "socket"
#include <al/log.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/tcp.h>
#include <sys/fcntl.h>

#include "../util/error.h"

#include "socket.h"
#include "socket_internal.h"

void nn_fd_set_blocking(s32 fd, bool blocking)
{
    s32 ret = fcntl(fd, F_GETFL, 0); // ret = flags
    al_assert(ret != -1);
    ret = fcntl(fd, F_SETFL, blocking ? (ret & ~O_NONBLOCK) : (ret | O_NONBLOCK)); // ret = success/error
    al_assert(ret == 0);
}

// https://github.com/mpv-player/mpv/blob/e575ec4fc3654387c7358bd3640877ef32628d2c/osdep/poll_wrapper.c#L29
#define NNWT_TIME_S_TO_NS(s) ((s) * INT64_C(1000000000))
s32 nn_poll_fds(struct nn_pollfd *fds, nn_nfds nfds, s64 timeout_ns)
{
    struct timespec ts;
    ts.tv_sec = timeout_ns / NNWT_TIME_S_TO_NS(1);
    ts.tv_nsec = timeout_ns % NNWT_TIME_S_TO_NS(1);
    struct timespec *tsp = (timeout_ns >= 0) ? &ts : NULL;
    return ppoll(fds, nfds, tsp, NULL);
}

bool nn_socket_init(struct nn_socket *sock, s32 flags)
{
    switch (sock->type) {
    case NNWT_SOCKET_TCP:
        sock->fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        break;
    case NNWT_SOCKET_UDP:
        sock->fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        break;
    case NNWT_SOCKET_UNIX:
        sock->fd = socket(AF_UNIX, SOCK_STREAM, 0);
        break;
    default:
        al_assert_and_return(false);
    }

    if (sock->fd < 0) {
        log_error("socket() failed: %s (%d).", nn_strerror(errno), errno);
        return false;
    }

    sock->addrinfo = NULL;

    nn_socket_apply_flags(sock, flags);

    return true;
}

void nn_socket_set_blocking(struct nn_socket *sock, bool blocking)
{
    nn_fd_set_blocking(sock->fd, blocking);
}

void nn_socket_set_nodelay(struct nn_socket *sock, s32 nodelay)
{
    al_assert(sock->type == NNWT_SOCKET_TCP);
    setsockopt(sock->fd, IPPROTO_TCP, TCP_NODELAY, &nodelay, sizeof(nodelay));
}

void nn_socket_set_reuse_addr(struct nn_socket *sock, s32 reuse_addr)
{
    al_assert(sock->type == NNWT_SOCKET_TCP);
    setsockopt(sock->fd, SOL_SOCKET, SO_REUSEADDR, &reuse_addr, sizeof(reuse_addr));
}

u32 nn_socket_get_send_buf(struct nn_socket *sock)
{
    u32 send_queue_size;
    socklen_t optlen = sizeof(send_queue_size);
    getsockopt(sock->fd, SOL_SOCKET, SO_SNDBUF, &send_queue_size, &optlen);
    return send_queue_size / 2;
}

void nn_socket_set_send_buf(struct nn_socket *sock, u32 sndbuf)
{
    u32 send_queue_size = sndbuf;
    setsockopt(sock->fd, SOL_SOCKET, SO_SNDBUF, &send_queue_size, sizeof(send_queue_size));
    al_assert(nn_socket_get_send_buf(sock) == send_queue_size);
}

u32 nn_socket_get_recv_buf(struct nn_socket *sock)
{
    u32 receive_queue_size;
    socklen_t optlen = sizeof(receive_queue_size);
    getsockopt(sock->fd, SOL_SOCKET, SO_RCVBUF, &receive_queue_size, &optlen);
    return receive_queue_size / 2;
}

void nn_socket_set_recv_buf(struct nn_socket *sock, u32 rcvbuf)
{
    u32 receive_queue_size = rcvbuf;
    setsockopt(sock->fd, SOL_SOCKET, SO_RCVBUF, &receive_queue_size, sizeof(receive_queue_size));
    al_assert(nn_socket_get_recv_buf(sock) == receive_queue_size);
}

static bool parse_address(struct nn_socket *sock, str *addr, u16 port)
{
    struct addrinfo hints = { 0 };
    hints.ai_family = AF_INET;
    switch (sock->type) {
    case NNWT_SOCKET_TCP:
        hints.ai_socktype = SOCK_STREAM;
        hints.ai_protocol = IPPROTO_TCP;
        break;
    case NNWT_SOCKET_UDP:
        hints.ai_socktype = SOCK_DGRAM;
        hints.ai_protocol = IPPROTO_UDP;
        break;
    }

    char *c_str = addr ? al_str_to_c_str(addr) : NULL;
    if (!c_str) hints.ai_flags = AI_PASSIVE;
    char port_str[6];
    al_snprintf(port_str, sizeof(port_str), "%hu", port);
    s32 status = getaddrinfo(c_str, port_str, &hints, &sock->addrinfo);
    al_free(c_str);
    if (status != 0) {
        log_error("Failed to parse address (%s).", gai_strerror(status));
        return false;
    }

    return true;
}

bool nn_socket_set(struct nn_socket *sock, str *addr, u16 port)
{
    al_assert(sock->type == NNWT_SOCKET_UDP);
    return parse_address(sock, addr, port);
}

bool nn_socket_bind(struct nn_socket *sock, str *addr, u16 port)
{
    struct sockaddr *saddr = NULL;
    socklen_t addrlen = 0;
    switch (sock->type) {
    case NNWT_SOCKET_TCP:
    case NNWT_SOCKET_UDP:
        if (!parse_address(sock, addr, port)) {
            return false;
        }
        saddr = sock->addrinfo->ai_addr;
        addrlen = sock->addrinfo->ai_addrlen;
        break;
    case NNWT_SOCKET_UNIX:
        al_assert(addr);
        sock->addr_un.sun_family = AF_UNIX;
        char *c_str = al_str_to_c_str(addr);
        al_memcpy(sock->addr_un.sun_path, c_str, al_strlen(c_str) + 1);
        unlink(c_str);
        al_free(c_str);
        saddr = (struct sockaddr *)&sock->addr_un;
        addrlen = sizeof(sock->addr_un);
        break;
    default:
        al_assert_and_return(false);
    }

    if (bind(sock->fd, saddr, addrlen) < 0) {
        log_error("bind(%.*s:%hu) failed: %s (%d).", nn_addr_x(addr), port, nn_strerror(errno), errno);
        return false;
    }

    log_debug("Socket bound to %.*s:%hu.", nn_addr_x(addr), port);

    return true;
}

bool nn_socket_listen(struct nn_socket *sock)
{
    al_assert(sock->type == NNWT_SOCKET_TCP || sock->type == NNWT_SOCKET_UNIX);

    if (listen(sock->fd, SOMAXCONN) < 0) {
        log_error("listen() failed: %s (%d).", nn_strerror(errno), errno);
        return false;
    }

    socklen_t optlen = sizeof(sock->acceptconn);
    getsockopt(sock->fd, SOL_SOCKET, SO_ACCEPTCONN, &sock->acceptconn, &optlen);
    al_assert(sock->acceptconn);

    log_debug("Listening.");

    return true;
}

bool nn_socket_accept(struct nn_socket *sock, struct nn_socket *cl, s32 flags)
{
    al_assert(sock->type == NNWT_SOCKET_TCP || sock->type == NNWT_SOCKET_UNIX);

    cl->type = sock->type;
    struct sockaddr_storage addr = { 0 };
    socklen_t addrlen = sizeof(addr);
    if ((cl->fd = accept(sock->fd, (struct sockaddr *)&addr, &addrlen)) == -1) {
        log_error("accept() failed: %s (%d).", nn_strerror(errno), errno);
        return false;
    }

    cl->addrinfo = NULL;

    nn_socket_apply_flags(cl, flags);

    switch (cl->type) {
    case NNWT_SOCKET_TCP: {
        u16 port = 0;
        char addr_str[INET6_ADDRSTRLEN] = { 0 };
        if (addr.ss_family == AF_INET) {
            struct sockaddr_in *saddr = (struct sockaddr_in *)&addr;
            port = ntohs(saddr->sin_port);
            inet_ntop(AF_INET, &saddr->sin_addr, addr_str, sizeof(addr_str));
        } else if (addr.ss_family == AF_INET6) {
            struct sockaddr_in6 *saddr = (struct sockaddr_in6 *)&addr;
            port = ntohs(saddr->sin6_port);
            inet_ntop(AF_INET6, &saddr->sin6_addr, addr_str, sizeof(addr_str));
        }
        log_info("Connection from %s:%hu.", addr_str, port);
        break;
    }
    case NNWT_SOCKET_UNIX:
        log_info("Connection on UNIX socket.");
        break;
    }

    return true;
}

bool nn_socket_connect(struct nn_socket *sock, str *addr, u16 port)
{
    s32 ret;
    switch (sock->type) {
    case NNWT_SOCKET_TCP:
    case NNWT_SOCKET_UDP:
        if (!parse_address(sock, addr, port)) {
            return false;
        }
        ret = connect(sock->fd, sock->addrinfo->ai_addr, sock->addrinfo->ai_addrlen);
        break;
    case NNWT_SOCKET_UNIX:
        sock->addr_un.sun_family = AF_UNIX;
        al_memcpy(sock->addr_un.sun_path, addr->data, addr->length);
        sock->addr_un.sun_path[addr->length] = '\0';
        ret = connect(sock->fd, (struct sockaddr *)&sock->addr_un, sizeof(sock->addr_un));
        break;
    default:
        al_assert_and_return(false);
    }
    if (ret != 0 && errno != EINPROGRESS) {
        log_error("connect(%.*s:%hu) failed: %s (%d).", nn_addr_x(addr), port, nn_strerror(errno), errno);
        return false;
    }
    return true;
}

ssize_t nn_socket_read(struct nn_socket *sock, void *buf, size_t size)
{
    return recv(sock->fd, buf, size, 0);
}

ssize_t nn_socket_write(struct nn_socket *sock, void *buf, size_t size)
{
    return send(sock->fd, buf, size, MSG_NOSIGNAL);
}

ssize_t nn_socket_sendto(struct nn_socket *sock, void *buf, size_t size)
{
    return sendto(sock->fd, buf, size, 0, sock->addrinfo->ai_addr, sock->addrinfo->ai_addrlen);
}

ssize_t nn_socket_recvfrom(struct nn_socket *sock, void *buf, size_t size)
{
    return recvfrom(sock->fd, buf, size, 0, sock->addrinfo->ai_addr, &sock->addrinfo->ai_addrlen);
}

// This would normally be a log_warn() but we rely on expected errors for control flow.
bool nn_socket_check_error(ssize_t ret)
{
    if (ret < 0) {
        log_debug("Socket error: %s (%d).", nn_strerror(errno), errno);
        return true;
    }
    return false;
}

s32 nn_socket_get_fd(struct nn_socket *sock)
{
    return sock->fd;
}

void nn_socket_shutdown(struct nn_socket *sock)
{
    shutdown(sock->fd, SHUT_RDWR);
}

void nn_socket_close(struct nn_socket *sock)
{
    close(sock->fd);
    if (sock->addrinfo) {
        freeaddrinfo(sock->addrinfo);
        sock->addrinfo = NULL;
    }
}

void nn_socket_cleanup(struct nn_socket *sock)
{
    if (sock->type == NNWT_SOCKET_UNIX && sock->acceptconn) {
        unlink(sock->addr_un.sun_path);
    }
}