#pragma once

#ifdef _WIN32
#include <ws2tcpip.h>
#include <winsock2.h>
#include <cstdio>
#endif

#define DEBUG

typedef unsigned char byte;

#ifdef _WIN32
typedef SOCKET Socket;
bool init_socket_lib() {
    // TODO: Should do some sort of smart version upgrade via data
    // See https://learn.microsoft.com/en-us/windows/win32/api/winsock/ns-winsock-wsadata
    WSADATA data;
    if (WSAStartup((2 << 8) | 2, &data) != 0) {
	printf("init_socket_lib failed: %d\n", WSAGetLastError());
	WSACleanup();
	return false;
    }
    return true;
}
#endif

extern bool socket_lib_initialized = init_socket_lib();

namespace tcp {

Socket make_socket();
bool connect(Socket &s, const char* hostname, int port);
bool send_transmit_stop(Socket &s);
bool send_read_stop(Socket &s);
bool disconnect(Socket &s);
bool send_data(Socket &s, byte* buf, int len);
bool get_data(Socket &s, byte* out_buf, int buf_size, int& out_num_written);

#ifdef _WIN32

Socket make_socket() { return INVALID_SOCKET; }

bool connect(Socket& s, const char *hostname, int port) {
    addrinfo *result, *ptr, hints;
    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    char port_str[128];
    sprintf(port_str, "%d", port);
    int code = getaddrinfo(hostname, port_str, &hints, &result);
    if (code != 0) {
	printf("getaddrinfo failed: %d\n", code);
	WSACleanup();
	return false;
    }

    s = INVALID_SOCKET;
    ptr = result;
    s = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
    if (s == INVALID_SOCKET) {
	printf("Error at socket(): %ld\n", WSAGetLastError());
	freeaddrinfo(result);
	WSACleanup();
	return false;
    }

    // Iterate through the linked list of addresses returned by getaddrinfo
    do {
	code = connect(s, ptr->ai_addr, (int)ptr->ai_addrlen);
	if (code == SOCKET_ERROR) {
	    closesocket(s);	
	    s = INVALID_SOCKET;
	} else {
	    break;
	}
	ptr = ptr->ai_next;
    } while (ptr);

    freeaddrinfo(result);

    if (s == INVALID_SOCKET) {
	printf("Unable to connect to %s:%d\n", hostname, port);
	WSACleanup();
	return false;
    }

    return true;
}

bool send_transmit_stop(Socket &s) {
    int code = shutdown(s, SD_SEND);
    if (code == SOCKET_ERROR) {
	printf("send_transmit_stop failed: %d\n", WSAGetLastError());
	closesocket(s);
	return false;
    }
    return true;
}

bool send_read_stop(Socket &s) {
    int code = shutdown(s, SD_RECEIVE);
    if (code == SOCKET_ERROR) {
	printf("send_read_stop failed: %d\n", WSAGetLastError());
	closesocket(s);
	return false;
    }
    return true;
}

bool disconnect(Socket &s) {
    int code = shutdown(s, SD_BOTH);
    if (code == SOCKET_ERROR) {
	printf("disconnect failed: %d\n", WSAGetLastError());
	closesocket(s);
	WSACleanup();
    }
    return code != SOCKET_ERROR;
}

// TODO: Retries? Should this be a param?
bool send_data(Socket &s, byte *buf, int len) {
    int code = send(s, (char*)buf, len, 0);
    if (code == SOCKET_ERROR) {
	printf("send_data failed: %d\n", WSAGetLastError());
	closesocket(s);
	WSACleanup();
    }
    return code != SOCKET_ERROR;
}

bool get_data(Socket &s, byte *out_buf, int buf_size, int &out_num_written) {
    int code = recv(s, (char*)out_buf, buf_size, 0);
    if (code >= 0) {
	out_num_written = code;
    } else {
	printf("get_data failed: %d\n", WSAGetLastError());
    }
    return code >= 0;
}

#endif


} // namespace tcp