#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "revealer_1.h"
#include "util.h"


void allocateMemoryForRevealer1Addresses(Revealer1Addresses** addresses)
{
    Revealer1Addresses* p_addresses = (Revealer1Addresses*)malloc(sizeof(Revealer1Addresses));
    p_addresses->allocated_count = 10;
    p_addresses->count = 0;

    char** ip_addresses = (char**)malloc(p_addresses->allocated_count * sizeof(char*));
    for (int i = 0; i < p_addresses->allocated_count; i++) {
        ip_addresses[i] = (char*)malloc(ADDRESS_SIZE * sizeof(char));
    }
    p_addresses->addresses = ip_addresses;

    *addresses = p_addresses;
}


int bindSocket(SOCKET sock, const char* adapter_addr)
{
    struct sockaddr_in sock_addr;
	memset(&sock_addr, 0, sizeof(sock_addr));
    sock_addr.sin_family = AF_INET;
    sock_addr.sin_port = htons(0);
    inet_pton(AF_INET, adapter_addr, &(sock_addr.sin_addr));
    int result = bind(sock, (struct sockaddr*)&sock_addr, sizeof(sock_addr));
    if (result == SOCKET_ERROR) {
        log_warning(L"Failed to bind socket to address = '%hs'", adapter_addr);
    }
    return result;
}


void closeSocket(SOCKET sock)
{
    int result;
#if defined(_WIN32) || defined(WIN64)
    result = closesocket(sock);
#else
    result = close(sock);
#endif
    if (result == SOCKET_ERROR) {
        log_warning(L"Failed to close socket");
    }

#if defined(_WIN32) || defined(WIN64)
    WSACleanup();
#endif
}


SOCKET createUPDSocketForBroadcast()
{
#if defined(_WIN32) || defined(WIN64)
    // Declare and initialize variables
    WSADATA wsaData = { 0 };
    // Initialize Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        log_warning(L"WSAStartup failed");
        return INVALID_SOCKET;
    }
#endif

    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET) {
        log_warning(L"Failed to create socket");
        return INVALID_SOCKET;
    }

    unsigned long value = 1;
	if (setsockopt(sock, SOL_SOCKET, SO_BROADCAST, (char*)&value, sizeof(value)) == SOCKET_ERROR) {
        log_warning(L"Failed to set BROADCAST option to socket");
        closeSocket(sock);
        return INVALID_SOCKET;
    }

#if defined(_WIN32) || defined(WIN64)
    unsigned long rec_timeout = TIMEOUT_MS;
#else
    struct timeval rec_timeout;
    rec_timeout.tv_sec = 0;
    rec_timeout.tv_usec = TIMEOUT_MS * 1000;
#endif
	if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (char*)&rec_timeout, sizeof(rec_timeout)) == SOCKET_ERROR) {
        log_warning(L"Failed to set timeout to socket");
        closeSocket(sock);
        return INVALID_SOCKET;
    }

    return sock;
}


void freeMemoryForRevealer1Addresses(Revealer1Addresses* addresses)
{
	if (!addresses)
		return;

	for (int i = 0; i < addresses->allocated_count; i++) {
		free(addresses->addresses[i]);
		addresses->addresses[i] = NULL;
	}
	free(addresses);
	addresses = NULL;
}


void increaseMemoryForRevealer1Addresses(Revealer1Addresses* addresses)
{
	int old_allocated_coount = addresses->allocated_count;
    addresses->allocated_count = (int)(1.5 * addresses->allocated_count);
    addresses->addresses = (char**)realloc(addresses->addresses, addresses->allocated_count * sizeof(char*));
	for (int i = old_allocated_coount; i < addresses->allocated_count; i++) {
        addresses->addresses[i] = (char*)malloc(ADDRESS_SIZE * sizeof(char));
    }
}


void printRevealer1Addresses(Revealer1Addresses* addresses)
{
	if (!addresses) {
		log_debug(L"No addresses found by Revealer-1 protocol");
		return;
	}

	log_debug(L"Number of addresses found by Revealer-1 protocol: %d", addresses->count);
	for (int i = 0; i < addresses->count; i++) {
		log_debug(L"#%d. %hs\n", i + 1, addresses->addresses[i]);
	}
}


void receiveResponses(SOCKET sock, Revealer1Addresses* addresses)
{
    while (1) {
        char recv_buffer[BUFFER_SIZE];
        struct sockaddr_in sender_addr;
#if defined(_WIN32) || defined(WIN64)
		int sender_addr_size = sizeof(sender_addr);
#else
		socklen_t sender_addr_size = sizeof(sender_addr);
#endif
		ssize_t result = recvfrom(sock, recv_buffer, sizeof(recv_buffer), 0, (struct sockaddr*)&sender_addr, &sender_addr_size);
        if (result == SOCKET_ERROR) {
            break;
        }

        char sender_ip[ADDRESS_SIZE];
        inet_ntop(AF_INET, &sender_addr.sin_addr, sender_ip, sizeof(sender_ip));
        saveAddress(sender_ip, addresses);
    }
}


void saveAddress(const char* address, Revealer1Addresses* addresses)
{
    if (addresses->count >= addresses->allocated_count) {
        increaseMemoryForRevealer1Addresses(addresses);
    }

#if defined(_WIN32) || defined(WIN64)
    strcpy_s(addresses->addresses[addresses->count], ADDRESS_SIZE, address);
#else
    strncpy(addresses->addresses[addresses->count], address, ADDRESS_SIZE);
#endif
    addresses->count++;
}


void searchByRevealer1Protocol(const char* adapter_addr, Revealer1Addresses** addresses)
{
	SOCKET sock = createUPDSocketForBroadcast();
	if (sock == INVALID_SOCKET) {
		return;
	}

	if (bindSocket(sock, adapter_addr) == SOCKET_ERROR) {
		closeSocket(sock);
		return;
	}

	if (sendMessageToBroadcast(sock) == SOCKET_ERROR) {
		closeSocket(sock);
		return;
	}

	allocateMemoryForRevealer1Addresses(addresses);
	receiveResponses(sock, *addresses);
	closeSocket(sock);
}


ssize_t sendMessageToBroadcast(SOCKET sock)
{
    struct sockaddr_in sock_addr;
#if defined(_WIN32) || defined(WIN64)
	int len = sizeof(sock_addr);
#else
	socklen_t len = sizeof(sock_addr);
#endif
	if (getsockname(sock, (struct sockaddr*)&sock_addr, &len) == SOCKET_ERROR) {
        log_warning(L"Failed to get socket name");
		return SOCKET_ERROR;
    }

    char buffer[BUFFER_SIZE] = "";
#if defined(_WIN32) || defined(WIN64)
    sprintf_s(buffer, BUFFER_SIZE, "DISCOVER_CUBIELORD_REQUEST %u", ntohs(sock_addr.sin_port));
#else
    snprintf(buffer, BUFFER_SIZE, "DISCOVER_CUBIELORD_REQUEST %u", ntohs(sock_addr.sin_port));
#endif

    struct sockaddr_in to_addr;
	memset(&to_addr, 0, sizeof(to_addr));
    to_addr.sin_family = AF_INET;
    to_addr.sin_addr.s_addr = htonl(INADDR_BROADCAST);
    to_addr.sin_port = htons(DEFAULT_PORT);
	ssize_t result = sendto(sock, buffer, (int)strlen(buffer), 0, (struct sockaddr*)&to_addr, sizeof(to_addr));
    if (result == SOCKET_ERROR) {
        log_warning(L"Failed to send message");
    }

    return result;
}
