#include <WinSock2.h>
#include <WS2tcpip.h>
#include <stdio.h>
#define _CRT_SECURE_NO_WARNINGS
#include <stdlib.h>
#include <Windows.h>
#include <string.h>


// enum, typedef, union, endianess, memset

void Chat(SOCKET* sock);
DWORD WINAPI ChatReceive(void* param);

void ListenChat(void);

void ConnectChat(void);

struct in_addr GetIPFromString(char* str);

struct in_addr GetIPFromUser();

void InitWinsock(void)
{
    int iResult;
    WSADATA wsaData;

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed: %d\n", iResult);
        exit(1);
    }
}

int main(void)
{
    char c;
    int i = -1;

    printf("Enter 1 to connect or 2 to listen: ");
    c = getchar();
    getchar(); // to get the \n after the number

    InitWinsock();

    switch (c)
    {
    case '1':
        ConnectChat();
        break;
    case '2':
        i++;
        ListenChat(i);
        break;
    default:
        printf("Invalid key given! Exiting!");
        return 1;
        break;
    }
    return 0;
}

void ListenChat(int i)
{
    SOCKET sock;

    struct sockaddr_in sockaddr;
    int res;
    int len;
    SOCKET* psock=NULL;

    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == SOCKET_ERROR)
    {
        printf("Error on sock init!");
        exit(2);
    }

    sockaddr.sin_family = AF_INET;
    sockaddr.sin_port = htons(7777);
    sockaddr.sin_addr = GetIPFromString("0.0.0.0");

    memset(sockaddr.sin_zero, 0, 8);

    res = bind(sock, (struct sockaddr*)&sockaddr, sizeof(sockaddr));
    if (res == SOCKET_ERROR)
    {
        printf("Error on sock bind!");
        exit(2);
    }

    res = listen(sock, 1);
    if (res == SOCKET_ERROR)
    {
        printf("Error on sock listen!");
        exit(2);
    }

    psock[i] = accept(sock, NULL, NULL);
    psock[i + 1] = NULL;
    if (psock[i] == SOCKET_ERROR)
    {
        printf("Error on sock accept! %d", WSAGetLastError());
        exit(2);
    }

    res = send(psock[i], "Hello! I am connected!", 23, NULL);
    if (res == SOCKET_ERROR)
    {
        printf("Error on sock send!");
        exit(2);
    }

    Chat(psock);
}

struct in_addr GetIPFromString(char* str)
{
    struct in_addr in;

    InetPton(AF_INET, str, &in.S_un.S_addr);
    return in;
}

struct in_addr GetIPFromUser()
{
    char buffer[256];

    printf("Insert IP: ");
    scanf("%s", buffer);

    return GetIPFromString(buffer[0] == '\0' || buffer[0] == ' ' ? "127.0.0.1" : buffer);
}

void ConnectChat(void)
{
    SOCKET* sock;
    struct sockaddr_in sockaddr;
    int res;

    sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == SOCKET_ERROR)
    {
        printf("Error on sock init!");
        exit(2);
    }

    sockaddr.sin_family = AF_INET;
    sockaddr.sin_port = htons(7777);
    sockaddr.sin_addr = GetIPFromUser();

    memset(sockaddr.sin_zero, 0, 8);

    res = connect(sock, (struct sockaddr*)&sockaddr, sizeof(struct sockaddr));
    if (res == SOCKET_ERROR)
    {
        printf("Error on sock connect!");
        exit(2);
    }

    res = send(sock, "Hello! I am connected!", 23, NULL);
    if (res == SOCKET_ERROR)
    {
        printf("Error on sock send!");
        exit(2);
    }

    Chat(sock);

}

void Chat(SOCKET* sock)
{
    char buff[2048];
    int res;
    HANDLE thd;
    int i = 0, j;

    while (!(sock[i] == NULL))
    {
        thd = CreateThread(NULL, NULL, &ChatReceive, (void*)sock[i], NULL, NULL);
        i++;


        while (1)
        {
            fgets(buff, 2048, stdin);
            j = 0;

            while (!(sock[j] == NULL))
            {
                res = send(sock[j], buff, strlen(buff) + 1, NULL);
                j++;

                if (res == SOCKET_ERROR)
                {
                    printf("send stopped.");
                    break;
                }
            }
        }

        WaitForSingleObject(thd, INFINITE);
    }

    for (j = 0; j<i - 1; j++)
        closesocket(sock[j]);
}

DWORD WINAPI ChatReceive(void* param)
{
    char buff[2048];
    int res;
    SOCKET sock = (SOCKET)param;

    while (1)
    {
        res = recv(sock, buff, 2048, NULL);
        if (res == SOCKET_ERROR)
        {
            printf("recv stopped.");
            break;
        }
        printf("Other: %s\n", buff);
    }

}