#include "WinsockSocket.h"

#include <winsock2.h>
#include <ws2tcpip.h>

#include "FileLogger.h"

WinsockSocket::WinsockSocket(const std::string& address, unsigned port)
    : Socket(address, port)
{

}

WinsockSocket::WinsockSocket(SOCKET handle)
    : Socket("", 0),
        mHandle(handle)
{
    mState.mConnectStatus = State::ConnectStatus::OK;
}

WinsockSocket::~WinsockSocket()
{
    MLOG_INFO("Socket being destroyed");
}

void WinsockSocket::initialize()
{
    addrinfo hints;
    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    auto result = ::getaddrinfo(mAddress.c_str(), std::to_string(mPort).c_str(), &hints, &mAddressInfo);
    if (result != 0)
    {
        mState.mConnectStatus = Socket::State::ConnectStatus::FAILED;
        mState.mErrorCode = result;
        mState.mErrorMessage = "WinsockSocket: getaddrinfo failed for connect";
        return;
    }

    mHandle = ::socket(mAddressInfo->ai_family, mAddressInfo->ai_socktype, mAddressInfo->ai_protocol);
    if (mHandle == INVALID_SOCKET)
    {
        onSockerError("WinsockSocket: Error at socket()");
        ::freeaddrinfo(mAddressInfo);
        return;
    }
}

void WinsockSocket::initializeForBind()
{
    addrinfo hints;
    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    hints.ai_flags = AI_PASSIVE;

    // https://learn.microsoft.com/en-us/windows/win32/winsock/windows-sockets-error-codes-2
    auto result = ::getaddrinfo(nullptr, std::to_string(mPort).c_str(), &hints, &mAddressInfo);
    if (result != 0)
    {
        mState.mBindStatus = Socket::State::BindStatus::FAILED;
        mState.mErrorCode = result;
        mState.mErrorMessage = "WinsockSocket: getaddrinfo failed for bind";
        return;
    }

    mHandle = ::socket(mAddressInfo->ai_family, mAddressInfo->ai_socktype, mAddressInfo->ai_protocol);
    if (mHandle == INVALID_SOCKET)
    {
        onSockerError("WinsockSocket: Error at socket()");
        ::freeaddrinfo(mAddressInfo);
        return;
    }
}

void WinsockSocket::doConnect()
{
    auto result = ::connect(mHandle, mAddressInfo->ai_addr, (int)mAddressInfo->ai_addrlen);
    if (result == SOCKET_ERROR)
    {
        ::closesocket(mHandle);
        mHandle = INVALID_SOCKET;
    }

    ::freeaddrinfo(mAddressInfo);

    if (mHandle == INVALID_SOCKET)
    {
        mState.mConnectStatus = Socket::State::ConnectStatus::FAILED;
        mState.mErrorCode = SOCKET_ERROR;
        mState.mErrorMessage = "WinsockSocket: Unable to connect to server.";
        return;
    }
    mState.mConnectStatus = Socket::State::ConnectStatus::OK;
}

void WinsockSocket::doBind()
{
    auto result = ::bind(mHandle, mAddressInfo->ai_addr, (int)mAddressInfo->ai_addrlen);
    if (result == SOCKET_ERROR)
    {
        ::closesocket(mHandle);
        mHandle = INVALID_SOCKET;
    }

    ::freeaddrinfo(mAddressInfo);

    if (mHandle == INVALID_SOCKET)
    {
        mState.mBindStatus = Socket::State::BindStatus::FAILED;
        mState.mErrorCode = ::WSAGetLastError();
        mState.mErrorMessage = "WinsockSocket: Unable to bind socket";
        return;
    }
    mState.mBindStatus = Socket::State::BindStatus::OK;
}

void WinsockSocket::doListen(onIncomingConnectionFunc connectionFunc)
{
    initializeForBind();
    if (mState.mBindStatus == Socket::State::BindStatus::FAILED)
    {
        return;
    }

    doBind();
    if (mState.mBindStatus == Socket::State::BindStatus::FAILED)
    {
        return;
    }

    if (::listen(mHandle, SOMAXCONN) == SOCKET_ERROR)
    {
        mState.mBindStatus = Socket::State::BindStatus::FAILED;
        mState.mErrorCode = ::WSAGetLastError();
        mState.mErrorMessage = "WinsockSocket: Listen failed";
        ::closesocket(mHandle);
        return;
    }

    while (true)
    {
        auto client_handle = ::accept(mHandle, NULL, NULL);
        if (client_handle == INVALID_SOCKET)
        {
            mState.mBindStatus = Socket::State::BindStatus::FAILED;
            mState.mErrorCode = ::WSAGetLastError();
            mState.mErrorMessage = "WinsockSocket: Accept failed";
            ::closesocket(mHandle);
            break;
        }
        else
        {
            connectionFunc(client_handle);
        }
    }
}

std::string WinsockSocket::send(const std::string& message)
{
    if (mState.mConnectStatus != Socket::State::ConnectStatus::OK)
    {
        initialize();
        if (mState.mConnectStatus == Socket::State::ConnectStatus::FAILED)
        {
            return {};
        }

        doConnect();
        if (mState.mConnectStatus == Socket::State::ConnectStatus::FAILED)
        {
            return {};
        }
    }

    auto result = ::send(mHandle, message.c_str(), static_cast<int>(message.size()), 0);
    if (result == SOCKET_ERROR)
    {
        onSockerError("WinsockSocket: Send failed.");
        return {};
    }

    result = ::shutdown(mHandle, SD_SEND);
    if (result == SOCKET_ERROR)
    {
        onSockerError("WinsockSocket: Post send shutdown failed.");
        return {};
    }

    std::string response;
    while (mState.mConnectStatus == Socket::State::ConnectStatus::OK)
    {
        response += recieve();
    }
    ::closesocket(mHandle);

    return response;
}

void WinsockSocket::respond(const std::string& message)
{
    auto result = ::send(mHandle, message.c_str(), static_cast<int>(message.size()), 0);
    if (result == SOCKET_ERROR)
    {
        onSockerError("WinsockSocket: Respond failed.");
    }
}

void WinsockSocket::onSockerError(const std::string& message)
{
    mState.mErrorCode = ::WSAGetLastError();
    mState.mErrorMessage = message;

    if (mState.mConnectStatus == Socket::State::ConnectStatus::OK)
    {
        ::closesocket(mHandle);
    }
    mState.mConnectStatus = Socket::State::ConnectStatus::FAILED;
}

std::string WinsockSocket::recieve()
{
    const int BUFFER_SIZE = 512;
    char buffer[BUFFER_SIZE];
    auto result = ::recv(mHandle, buffer, BUFFER_SIZE, 0);
    if (result > 0)
    {
        return std::string(buffer);
    }
    else if (result == 0)
    {
        mState.mConnectStatus = Socket::State::ConnectStatus::UNSET;
    }
    else
    {
        mState.mConnectStatus = Socket::State::ConnectStatus::FAILED;
        mState.mErrorCode = ::WSAGetLastError();
        mState.mErrorMessage = "WinsockSocket: recv failed.";
    }
    return {};
}
