#include "Internal.h"
#include <stdlib.h>
#include <string.h>
#include <assert.h>

int CgeNetInit(void) {
    WSADATA wsaData;
    return (WSAStartup(MAKEWORD(2, 2), &wsaData) == 0) ? 1 : 0;
}

void CgeNetCleanup(void) {
    WSACleanup();
}

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

    case WSAEHOSTUNREACH:
        return CGE_NET_EHOSTUNREACH;

    case WSAENETUNREACH:
        return CGE_NET_ENETUNREACH;

    case WSAEINVAL:
        return CGE_NET_EINVAL;

    case WSAENOBUFS:
    case WSA_NOT_ENOUGH_MEMORY:
        return CGE_NET_EOOM;

    case WSAEADDRINUSE:
        return CGE_NET_EADDRINUSE;

    case WSAEADDRNOTAVAIL:
        return CGE_NET_EADDRNOTAVAIL;

    case WSAECONNREFUSED:
        return CGE_NET_ECONNREFUSED;

    case WSAECONNRESET:
        return CGE_NET_ECONNRESET;

    case WSAECONNABORTED:
        return CGE_NET_ECONNABORTED;

    case WSAETIMEDOUT:
        return CGE_NET_ETIMEDOUT;

    case WSAEWOULDBLOCK:
        return CGE_NET_EWOULDBLOCK;

    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 == INVALID_SOCKET)
        return wsaToErrorCode(WSAGetLastError());

    return CGE_NET_EOK;
}

CgeNetSocket *CgeNetOpen(int domain, int type, int *result) {
    CgeNetSocket *socket = NULL;
    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);

    closesocket(socket->handle);
    free(socket);
}

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

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

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

    if (getsockopt(socket->handle, SOL_SOCKET, SO_ERROR, (char *)&result, &len)) {
        socket->lastError = wsaToErrorCode(WSAGetLastError());
        return 0;
    }

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

int CgeNetSetBlocking(CgeNetSocket *socket, int value) {
    u_long mode;

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

    mode = value ? 0 : 1;
    if (ioctlsocket(socket->handle, FIONBIO, &mode) == SOCKET_ERROR) {
        socket->lastError = wsaToErrorCode(WSAGetLastError());
        return 0;
    }

    return 1;
}

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 != INVALID_SOCKET);
    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, (int)salen);
    if (result == SOCKET_ERROR)
        socket->lastError = wsaToErrorCode(WSAGetLastError());

    return result != SOCKET_ERROR;
}

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

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

    result = listen(socket->handle, backlog);
    if (result == SOCKET_ERROR)
        socket->lastError = wsaToErrorCode(WSAGetLastError());

    return result != SOCKET_ERROR;
}

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

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

    fd = accept(socket->handle, (struct sockaddr*)&ss, &slen);
    if (fd == INVALID_SOCKET) {
        socket->lastError = wsaToErrorCode(WSAGetLastError());
        return NULL;
    }

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

    client->handle = fd;
    client->domain = socket->domain;
    client->type = socket->type;
    client->lastError = CGE_NET_EOK;
    client->pollObject = NULL;
    client->pollIndex = 0;
    client->context = NULL;

    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 if (ss.ss_family == AF_INET6) {
            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;

    assert(socket != NULL);
    assert(socket->handle != INVALID_SOCKET);
    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, (int)salen);
    if (result == SOCKET_ERROR)
        socket->lastError = wsaToErrorCode(WSAGetLastError());

    return result != SOCKET_ERROR;
}

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

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

    result = send(socket->handle, (const char *)buffer, (int)size, 0);
    if (actual)
        *actual = (size_t)result;
    socket->lastError = (result == SOCKET_ERROR) ? wsaToErrorCode(WSAGetLastError()) : CGE_NET_EOK;

    return result != SOCKET_ERROR;
}

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

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

    result = recv(socket->handle, (char *)buffer, (int)size, 0);
    if (actual)
        *actual = (size_t)result;

    if (result == 0 && socket->type == CGE_NET_TCP) {
        socket->lastError = CGE_NET_EDISCONNECT;
        return 0;
    }
    socket->lastError = (result == SOCKET_ERROR) ? wsaToErrorCode(WSAGetLastError()) : CGE_NET_EOK;

    return result != SOCKET_ERROR;
}

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

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

    result = recv(socket->handle, (char *)buffer, (int)size, MSG_PEEK);
    if (actual)
        *actual = (size_t)result;

    if (result == 0 && socket->type == CGE_NET_TCP) {
        socket->lastError = CGE_NET_EDISCONNECT;
        return 0;
    }
    socket->lastError = (result == SOCKET_ERROR) ? wsaToErrorCode(WSAGetLastError()) : CGE_NET_EOK;

    return result != SOCKET_ERROR;
}

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

    assert(socket != NULL);
    assert(socket->handle != INVALID_SOCKET);
    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);
    }

    result = sendto(socket->handle, (const char *)buffer, (int)size, 0, sa, (int)salen);
    if (actual)
        *actual = (size_t)result;
    socket->lastError = (result == SOCKET_ERROR) ? wsaToErrorCode(WSAGetLastError()) : CGE_NET_EOK;

    return result != SOCKET_ERROR;
}

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

    assert(socket != NULL);
    assert(socket->handle != INVALID_SOCKET);
    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;
    }

    result = recvfrom(socket->handle, (char *)buffer, (int)size, 0, (struct sockaddr *)&ss, &sslen);
    if (result == SOCKET_ERROR) {
        socket->lastError = wsaToErrorCode(WSAGetLastError());
        return 0;
    }

    if (actual)
        *actual = (size_t)result;

    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);
    }
    socket->lastError = CGE_NET_EOK;

    return 1;
}

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

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