#include "Internal.h"
#include <arpa/inet.h>
#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>

int CgeNetInit(void) {
    return 1;
}

void CgeNetCleanup(void) {
}

static int errnoToErrorCode(int value) {
    switch (value) {
    case 0:
        return CGE_NET_EOK;

    case EHOSTUNREACH:
        return CGE_NET_EHOSTUNREACH;

    case ENETUNREACH:
        return CGE_NET_ENETUNREACH;

    case EINVAL:
        return CGE_NET_EINVAL;

    case ENOMEM:
        return CGE_NET_EOOM;

    case ETIMEDOUT:
        return CGE_NET_ETIMEDOUT;

    case EWOULDBLOCK:
        return CGE_NET_EWOULDBLOCK;

    case EADDRINUSE:
        return CGE_NET_EADDRINUSE;

    case EADDRNOTAVAIL:
        return CGE_NET_EADDRNOTAVAIL;

    case ECONNREFUSED:
        return CGE_NET_ECONNREFUSED;

    case ECONNRESET:
        return CGE_NET_ECONNRESET;

    case ECONNABORTED:
        return CGE_NET_ECONNABORTED;

    default:
        return CGE_NET_EUNKNOWN;
    }
}

static int socketInit(CgeNetSocket *s, int domain, int type) {
    int protocol;

    assert(s != NULL);

    s->domain = domain;
    s->type = type;
    s->pollObject = NULL;
    s->pollIndex = 0;

    switch (domain) {
        case CGE_NET_IPV4: domain = AF_INET; break;
        case CGE_NET_IPV6: domain = AF_INET6; break;
        default: return CGE_NET_EINVAL;
    }

    switch (type) {
        case CGE_NET_TCP: type = SOCK_STREAM; protocol = IPPROTO_TCP; break;
        case CGE_NET_UDP: type = SOCK_DGRAM; protocol = IPPROTO_UDP; break;
        default: return CGE_NET_EINVAL;
    }

    s->handle = socket(domain, type, protocol);
    if (s->handle == -1)
        return errnoToErrorCode(errno);
    fcntl(s->handle, F_SETFD, FD_CLOEXEC);

    return CGE_NET_EOK;
}

CgeNetSocket *CgeNetOpen(int domain, int type, int *result) {
    CgeNetSocket *socket;
    int code = CGE_NET_EOOM;

    if ((socket = malloc(sizeof(*socket)))) {
        if ((code = socketInit(socket, domain, type))) {
            free(socket);
            socket = NULL;
        }
    }
    if (result)
        *result = code;

    return socket;
}

void CgeNetClose(CgeNetSocket *socket) {
    if (!socket)
        return;

    if (socket->pollObject)
        CgeNetPollRemove(socket->pollObject, socket);

    assert(socket->handle != -1);
    close(socket->handle);
    free(socket);
}

int CgeNetErrorCode(CgeNetSocket *socket) {
    assert(socket != NULL);
    return socket->lastError;
}

int CgeNetFetchError(CgeNetSocket *socket) {
    int result;
    socklen_t len = sizeof(result);

    assert(socket != NULL);
    assert(socket->handle != -1);

    if (getsockopt(socket->handle, SOL_SOCKET, SO_ERROR, &result, &len) < 0) {
        socket->lastError = errnoToErrorCode(errno);
        return 0;
    }

    socket->lastError = errnoToErrorCode(result);
    return result != 0;
}

int CgeNetSetBlocking(CgeNetSocket *socket, int value) {
    int flags;

    assert(socket != NULL);
    assert(socket->handle != -1);

    if ((flags = fcntl(socket->handle, F_GETFL, 0)) == -1)
        goto error;
    flags = value ? (flags | O_NONBLOCK) : (flags & ~O_NONBLOCK);
    if (fcntl(socket->handle, F_SETFL, flags) == -1)
        goto error;
    return 1;

error:
    socket->lastError = errnoToErrorCode(errno);
    return 0;
}

int CgeNetBind(CgeNetSocket *socket, const CgeNetAddr *addr) {
    struct sockaddr_in sa4;
    struct sockaddr_in6 sa6;
    const struct sockaddr *sa;
    socklen_t salen;
    int result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(addr != NULL);
    assert(addr->type == CGE_NET_IPV4 || addr->type == CGE_NET_IPV6);

    if (addr->type == CGE_NET_IPV4) {
        memset(&sa4, 0, sizeof(sa4));
        sa4.sin_family = AF_INET;
        sa4.sin_port = htons(addr->ipv4.port);
        memcpy(&sa4.sin_addr, addr->ipv4.addr, 4);
        sa = (struct sockaddr*)&sa4;
        salen = sizeof(sa4);
    } else {
        memset(&sa6, 0, sizeof(sa6));
        sa6.sin6_family = AF_INET6;
        sa6.sin6_port = htons(addr->ipv6.port);
        memcpy(&sa6.sin6_addr, addr->ipv6.addr, 16);
        sa = (struct sockaddr*)&sa6;
        salen = sizeof(sa6);
    }

    result = bind(socket->handle, sa, salen);
    if (result == -1)
        socket->lastError = errnoToErrorCode(errno);

    return result != -1;
}

int CgeNetListen(CgeNetSocket *socket, unsigned int backlog) {
    int result;

    assert(socket != NULL);
    assert(socket->handle != -1);

    result = listen(socket->handle, backlog);
    if (result < 0)
        socket->lastError = errnoToErrorCode(errno);
    return result == 0;
}

CgeNetSocket *CgeNetAccept(CgeNetSocket *socket, CgeNetAddr *addr) {
    int fd;
    CgeNetSocket *client;
    struct sockaddr_storage ss;
    socklen_t slen = sizeof(ss);

    assert(socket != NULL);
    assert(socket->handle != -1);

    do {
        fd = accept(socket->handle, (struct sockaddr*)&ss, &slen);
    } while (fd == -1 && errno == EINTR);

    if (fd == -1) {
        socket->lastError = errnoToErrorCode(errno);
        return NULL;
    }

    if (!(client = malloc(sizeof(*client)))) {
        close(fd);
        socket->lastError = CGE_NET_EOOM;
        return NULL;
    }

    client->handle = fd;
    client->domain = socket->domain;
    client->type = socket->type;
    client->lastError = CGE_NET_EOK;
    fcntl(fd, F_SETFD, FD_CLOEXEC);

    if (addr) {
        if (ss.ss_family == AF_INET) {
            struct sockaddr_in *sa4 = (struct sockaddr_in*)&ss;
            addr->ipv4.type = CGE_NET_IPV4;
            memcpy(addr->ipv4.addr, &sa4->sin_addr, 4);
            addr->ipv4.port = ntohs(sa4->sin_port);
        } else {
            struct sockaddr_in6 *sa6 = (struct sockaddr_in6*)&ss;
            addr->ipv6.type = CGE_NET_IPV6;
            memcpy(addr->ipv6.addr, &sa6->sin6_addr, 16);
            addr->ipv6.port = ntohs(sa6->sin6_port);
        }
    }

    return client;
}

int CgeNetConnect(CgeNetSocket *socket, const CgeNetAddr *addr) {
    struct sockaddr_in sa4;
    struct sockaddr_in6 sa6;
    const struct sockaddr *sa;
    socklen_t salen;
    int result;
    int pollResult;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(addr != NULL);
    assert(addr->type == CGE_NET_IPV4 || addr->type == CGE_NET_IPV6);

    if (addr->type == CGE_NET_IPV4) {
        memset(&sa4, 0, sizeof(sa4));
        sa4.sin_family = AF_INET;
        sa4.sin_port = htons(addr->ipv4.port);
        memcpy(&sa4.sin_addr, addr->ipv4.addr, 4);
        sa = (struct sockaddr*)&sa4;
        salen = sizeof(sa4);
    } else {
        memset(&sa6, 0, sizeof(sa6));
        sa6.sin6_family = AF_INET6;
        sa6.sin6_port = htons(addr->ipv6.port);
        memcpy(&sa6.sin6_addr, addr->ipv6.addr, 16);
        sa = (struct sockaddr*)&sa6;
        salen = sizeof(sa6);
    }

    result = connect(socket->handle, sa, salen);
    if (result < 0 && errno == EINTR) {
        socklen_t len = sizeof(result);

        result = 0;
        do {
            struct pollfd pfd;
            pfd.fd = socket->handle;
            pfd.events = POLLOUT;
            pollResult = poll(&pfd, 1, -1);
        } while (pollResult < 0 && errno == EINTR);

        if (pollResult < 0) {
            socket->lastError = errnoToErrorCode(errno);
            return 0;
        }

        if (getsockopt(socket->handle, SOL_SOCKET, SO_ERROR, &result, &len) < 0) {
            socket->lastError = errnoToErrorCode(errno);
            return 0;
        }
        errno = result;
        result = result != 0 ? -1 : 0;
    }

    if (result == -1)
        socket->lastError = errnoToErrorCode(errno);
    return result != -1;
}

int CgeNetWrite(CgeNetSocket *socket, const void *buffer, size_t size,
                size_t *actual) {
    ssize_t result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(buffer != NULL);

    do {
        result = send(socket->handle, buffer, size, 0);
    } while (result < 0 && errno == EINTR);

    if (actual)
        *actual = result;
    socket->lastError = result >= 0 ? CGE_NET_EOK : errnoToErrorCode(errno);
    return result >= 0;
}

int CgeNetRead(CgeNetSocket *socket, void *buffer, size_t size,
               size_t *actual) {
    ssize_t result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(buffer != NULL);

    do {
        result = recv(socket->handle, buffer, size, 0);
    } while (result < 0 && errno == EINTR);

    if (size > 0 && result == 0 && socket->type == CGE_NET_TCP) {
        socket->lastError = CGE_NET_EDISCONNECT;
        return 0;
    }

    if (actual)
        *actual = result;
    socket->lastError = result >= 0 ? CGE_NET_EOK : errnoToErrorCode(errno);
    return result >= 0;
}

int CgeNetPeek(CgeNetSocket *socket, void *buffer, size_t size, size_t *actual) {
    ssize_t result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(buffer != NULL);

    do {
        result = recv(socket->handle, buffer, size, MSG_PEEK);
    } while (result < 0 && errno == EINTR);

    if (size > 0 && result == 0 && socket->type == CGE_NET_TCP) {
        socket->lastError = CGE_NET_EDISCONNECT;
        return 0;
    }

    if (actual)
        *actual = result;
    socket->lastError = result >= 0 ? CGE_NET_EOK : errnoToErrorCode(errno);
    return result >= 0;
}

int CgeNetWriteTo(CgeNetSocket *socket, const CgeNetAddr *addr,
                  const void *buffer, size_t size, size_t *actual) {
    struct sockaddr_in sa4;
    struct sockaddr_in6 sa6;
    const struct sockaddr *sa;
    socklen_t salen;
    ssize_t result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(addr != NULL);
    assert(addr->type == CGE_NET_IPV4 || addr->type == CGE_NET_IPV6);
    assert(buffer != NULL);

    if (socket->type != CGE_NET_UDP) {
        socket->lastError = CGE_NET_EINVAL;
        return -1;
    }

    if (addr->type == CGE_NET_IPV4) {
        memset(&sa4, 0, sizeof(sa4));
        sa4.sin_family = AF_INET;
        sa4.sin_port = htons(addr->ipv4.port);
        memcpy(&sa4.sin_addr, addr->ipv4.addr, 4);
        sa = (struct sockaddr*)&sa4;
        salen = sizeof(sa4);
    } else {
        memset(&sa6, 0, sizeof(sa6));
        sa6.sin6_family = AF_INET6;
        sa6.sin6_port = htons(addr->ipv6.port);
        memcpy(&sa6.sin6_addr, addr->ipv6.addr, 16);
        sa = (struct sockaddr*)&sa6;
        salen = sizeof(sa6);
    }

    do {
        result = sendto(socket->handle, buffer, size, 0, sa, salen);
    } while (result < 0 && errno == EINTR);

    if (actual)
        *actual = result;
    socket->lastError = result >= 0 ? CGE_NET_EOK : errnoToErrorCode(errno);
    return result >= 0;
}

int CgeNetReadFrom(CgeNetSocket *socket, CgeNetAddr* addr, void *buffer,
                   size_t size, size_t *actual) {
    struct sockaddr_storage ss;
    socklen_t slen = sizeof(ss);
    ssize_t result;

    assert(socket != NULL);
    assert(socket->handle != -1);
    assert(addr != NULL);
    assert(addr->type == CGE_NET_IPV4 || addr->type == CGE_NET_IPV6);
    assert(buffer != NULL);

    if (socket->type != CGE_NET_UDP) {
        socket->lastError = CGE_NET_EINVAL;
        return -1;
    }

    do {
        result = recvfrom(socket->handle, buffer, size, 0,
                          (struct sockaddr*)&ss, &slen);
    } while (result < 0 && errno == EINTR);

    if (result < 0) {
        socket->lastError = errnoToErrorCode(errno);
        return 0;
    }

    if (ss.ss_family == AF_INET) {
        struct sockaddr_in *sa4 = (struct sockaddr_in*)&ss;
        addr->ipv4.type = CGE_NET_IPV4;
        memcpy(addr->ipv4.addr, &sa4->sin_addr, 4);
        addr->ipv4.port = ntohs(sa4->sin_port);
    } else {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6*)&ss;
        addr->ipv6.type = CGE_NET_IPV6;
        memcpy(addr->ipv6.addr, &sa6->sin6_addr, 16);
        addr->ipv6.port = ntohs(sa6->sin6_port);
    }

    if (actual)
        *actual = result;
    return 1;
}

void *CgeNetContext(CgeNetSocket *socket) {
    assert(socket != NULL);
    return socket->context;
}

void CgeNetSetContext(CgeNetSocket *socket, void *data) {
    assert(socket != NULL);
    socket->context = data;
}