#include <Utilities/Socket.hpp>
#include <Protocol/Protocol.hpp>
#include <Utilities/Utilities.hpp>
#include <string>

void parseNetAddr(std::string addr, std::string &hostname, std::string &portstr)
{
    portstr = "14201";

    if (addr.size() >= 2 && addr.front() == '[' && addr.back() == ']')
    {
        hostname = addr.substr(1, addr.size() - 2);
        return;
    }

    size_t colonPos = addr.rfind(':');
    if (colonPos != std::string::npos)
    {
        portstr = addr.substr(colonPos + 1);
        addr = addr.substr(0, colonPos);
    }

    if (addr.size() >= 2 && addr.front() == '[' && addr.back() == ']')
    {
        hostname = addr.substr(1, addr.size() - 2);
        return;
    }

    hostname = addr;
}

SOCKET connectSocket(std::string addrstr, std::string &error)
{
    std::string server, portstr;
    parseNetAddr(addrstr, server, portstr);

    struct addrinfo hints, *addrList, *addr;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET6; 
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags = AI_V4MAPPED | AI_ADDRCONFIG | AI_ALL;
    hints.ai_protocol = IPPROTO_TCP;

    int s = getaddrinfo(server.c_str(), portstr.c_str(), &hints, &addrList);
    if (s != 0)
    {
        error = "Failed to resolve the server address.";
        return INVALID_SOCKET;
    }

    SOCKET sock = INVALID_SOCKET;
    for (addr = addrList; addr != NULL; addr = addr->ai_next)
    {
        sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
        if (sock == INVALID_SOCKET) continue;

        int v6only = 0;
        setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, (char*) &v6only, sizeof(v6only));

        int nodelay = 1;
        setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (char*) &nodelay, sizeof(nodelay));

        // Set the socket to non-blocking mode
        if (!setSocketBlocking(sock, false))
        {
            closesocket(sock);
            sock = INVALID_SOCKET;
            continue;
        }

        if (connect(sock, addr->ai_addr, addr->ai_addrlen) == SOCKET_ERROR)
        {
            int errorCode = lastSocketError();
            if (errorCode != SOCKET_EWOULDBLOCK && errorCode != SOCKET_EINPROGRESS)
            {
                closesocket(sock);
                sock = INVALID_SOCKET;
                continue;
            }
        }

        // Use select() to check for writeability and thus connection success.
        fd_set set;
        struct timeval timeout;

        FD_ZERO(&set);
        FD_SET(sock, &set);
        timeout.tv_sec = Protocol::timeout;
        timeout.tv_usec = 0;

        int selRes = select(sock + 1, NULL, &set, NULL, &timeout);
        if (selRes > 0)
        {
            int so_error;
            socklen_t len = sizeof(so_error);

            getsockopt(sock, SOL_SOCKET, SO_ERROR, (char*) &so_error, &len);
            if (so_error == 0)
            {
                break; // Connection established successfully
            }
        }

        closesocket(sock);
        sock = INVALID_SOCKET;
    }

    freeaddrinfo(addrList);

    if (sock == INVALID_SOCKET)
    {
        error = "Connection failed.";
        return INVALID_SOCKET;
    }

    return sock;
}

bool setSocketBlocking(SOCKET fd, bool blocking)
{
#ifdef _WIN32
    unsigned long mode = blocking ? 0 : 1;
    return ioctlsocket(fd, FIONBIO, &mode) == 0;
#else
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) return false;
    flags = blocking ? (flags & ~O_NONBLOCK) : (flags | O_NONBLOCK);
    return fcntl(fd, F_SETFL, flags) == 0;
#endif
};

bool readSocket(SOCKET sock, char *buffer, uint32_t size)
{
    while (size != 0)
    {
        fd_set readfds;
        FD_ZERO(&readfds);
        FD_SET(sock, &readfds);

        struct timeval tv;
        tv.tv_sec = Protocol::timeout;
        tv.tv_usec = 0;

        if (select((unsigned int) sock+1, &readfds, NULL, NULL, &tv) != 1)
        {
            return false;
        };

        int bytesRead = recv(sock, buffer, size, 0);
        if (bytesRead == SOCKET_ERROR || bytesRead == 0)
        {
            return false;
        };

        buffer += bytesRead;
        size -= bytesRead;
    };

    return true;
};

bool writeSocket(SOCKET sock, const char *buffer, uint32_t size)
{
    while (size != 0)
    {
        fd_set writefds;
        FD_ZERO(&writefds);
        FD_SET(sock, &writefds);

        struct timeval tv;
        tv.tv_sec = Protocol::timeout;
        tv.tv_usec = 0;

        if (select((unsigned int) sock+1, NULL, &writefds, NULL, &tv) != 1)
        {
            return false;
        };

        int bytesSent = send(sock, buffer, size, 0);
        if (bytesSent == SOCKET_ERROR || bytesSent == 0)
        {
            return false;
        };

        buffer += bytesSent;
        size -= bytesSent;
    };

    return true;
};

std::string getSocketRemoteAddr(SOCKET sock)
{
    struct sockaddr_in6 addr;
    socklen_t len = sizeof(addr);

    if (getpeername(sock, reinterpret_cast<struct sockaddr*>(&addr), &len) != 0)
    {
        crash("getpeername failed on a socket.");
    }

    if (addr.sin6_family != AF_INET6)
    {
        crash("getSocketRemoteAddr received an address family other than AF_INET6.");
    }

    char ip[INET6_ADDRSTRLEN];
    if (!inet_ntop(AF_INET6, &addr.sin6_addr, ip, sizeof(ip)))
    {
        crash("getSocketRemoteAddr: failed to convert address to printable form.");
    }

    uint16_t port = ntohs(addr.sin6_port);

    return formatString("[$+]:$+", ip, port);
}