//  --- SECTION 1: Includes and Global Variables ---

#include <winsock2.h>
#include <ws2tcpip.h>
#include <iostream>
#include <string>
#include <vector>
#include <thread>
#include <mutex>
#include <chrono>
#include <sstream>
#include <iomanip>
#include <algorithm>

using namespace std;

#define PORT 54000

//  list of connected client sockets
vector<SOCKET> clients;

//  list of connected usernames
vector<string> usernames;

//  mutex to stop threads clashing
mutex clientsMutex;

//  ANSI RGB colour codes set to users
vector<string> colourList = {
    "\033[38;2;180;150;220m",   //  purple
    "\033[38;2;150;220;150m",   //  green
    "\033[38;2;255;200;150m",   //  orange
    "\033[38;2;130;190;255m",   //  blue
    "\033[38;2;255;230;130m",   //  yellow
    "\033[38;2;255;180;200m"    //  pink
};

//  red used for disconnect messages
const string colourRed = "\033[38;2;210;80;80m";

//  resets colour back to default
const string colourReset = "\033[0m";

//  purple reserved for the users own messages
int nextColourIndex = 1;



//  --- SECTION 2: Timestamp Helper Function ---

string GetTimestamp()
{
    auto n = chrono::system_clock::now();
    auto in_time_t = chrono::system_clock::to_time_t(n);

    tm buf;
    localtime_s(&buf, &in_time_t);

    stringstream ss;
    ss << "[" << put_time(&buf, "%H:%M:%S") << "]";

    return ss.str();
}



//  --- SECTION 3: Reliable Send Function ---

//  keeps sending until all bytes are sent
int Writen(SOCKET sockfd, const char* data, int bytesToWrite)
{
    int nleft = bytesToWrite;
    int nwritten;

    while (nleft > 0)
    {
        nwritten = send(sockfd, data, nleft, 0);

        if (nwritten <= 0)
        {
            return -1;
        }
        else
        {
            nleft -= nwritten;
            data += nwritten;
        }
    }

    return bytesToWrite;
}

//  sends the size of the message first
void WriteSize(SOCKET socket, int size)
{
    size = htonl(size);
    char* tosend = (char*)&size;
    Writen(socket, tosend, sizeof(size));
}



//  --- SECTION 4: Broadcast Function ---

//  sends a message to all connected users except the sender
void BroadcastMessage(const string& message, SOCKET senderSocket)
{
    lock_guard<mutex> lock(clientsMutex);

    for (SOCKET clientSocket : clients)
    {
        if (clientSocket != senderSocket)
        {
            WriteSize(clientSocket, (int)message.length());
            Writen(clientSocket, message.c_str(), (int)message.length());
        }
    }
}



//  --- SECTION 5: Send User List Function ---


//  builds the user list and sends it back to the user
void SendUserList(SOCKET requestingSocket)
{
    lock_guard<mutex> lock(clientsMutex);

    string userList = "Connected Users: ";

    for (int i = 0; i < (int)usernames.size(); i++)
    {
        userList += usernames[i];

        if (i < (int)usernames.size() - 1)
        {
            userList += ", ";
        }
    }

    WriteSize(requestingSocket, (int)userList.length());
    Writen(requestingSocket, userList.c_str(), (int)userList.length());
}



//  --- SECTION 6: User Handler ---

//  runs on its own thread for each connected user
void HandleClient(SOCKET clientSocket)
{
    //  receive the username from the client
    int usernameSize = 0;
    int recvResult = recv(clientSocket, (char*)&usernameSize, sizeof(int), 0);

    //  if the user disconnects before sending a username, clean up and stop
    if (recvResult <= 0)
    {
        cout << GetTimestamp()
            << " ######### New User Disconnected #########" << endl;

        lock_guard<mutex> lock(clientsMutex);
        clients.erase(remove(clients.begin(), clients.end(), clientSocket),
            clients.end());
        closesocket(clientSocket);
        return;
    }

    usernameSize = ntohl(usernameSize);

    char usernameBuffer[256];
    recv(clientSocket, usernameBuffer, usernameSize, 0);
    usernameBuffer[usernameSize] = '\0';
    string username(usernameBuffer);

    string clientColour;

    //  keep asking for a valid username until one is accepted
    while (true)
    {
        //  check if the username is blank
        if (username.empty())
        {
            string rejectMessage = "REJECTED:Username Cannot Be Blank!";
            WriteSize(clientSocket, (int)rejectMessage.length());
            Writen(clientSocket, rejectMessage.c_str(),
                (int)rejectMessage.length());
        }
        else
        {
            //  check if the username is already taken
            bool isTaken = false;
            {
                lock_guard<mutex> lock(clientsMutex);
                for (const string& existing : usernames)
                {
                    if (existing == username)
                    {
                        isTaken = true;
                        break;
                    }
                }
            }

            if (isTaken)
            {
                string rejectMessage = "REJECTED:Username Already Taken";
                WriteSize(clientSocket, (int)rejectMessage.length());
                Writen(clientSocket, rejectMessage.c_str(),
                    (int)rejectMessage.length());
            }
            else
            {
                break;
            }
        }

        //  if the user disconnects, clean up and stop
        int retrySize = 0;
        int retryResult = recv(clientSocket, (char*)&retrySize, sizeof(int), 0);

        if (retryResult <= 0)
        {
            cout << GetTimestamp()
                << " ######### New User Disconnected #########" << endl;

            lock_guard<mutex> lock(clientsMutex);
            clients.erase(remove(clients.begin(), clients.end(), clientSocket),
                clients.end());
            closesocket(clientSocket);
            return;
        }

        retrySize = ntohl(retrySize);

        char retryBuffer[256];
        recv(clientSocket, retryBuffer, retrySize, 0);
        retryBuffer[retrySize] = '\0';
        username = retryBuffer;
    }

    //  username is valid > accept the client
    string acceptMessage = "ACCEPTED";
    WriteSize(clientSocket, (int)acceptMessage.length());
    Writen(clientSocket, acceptMessage.c_str(), (int)acceptMessage.length());

    //  assign a colour and add to the username list
    {
        lock_guard<mutex> lock(clientsMutex);
        clientColour = colourList[nextColourIndex % colourList.size()];
        nextColourIndex++;
        usernames.push_back(username);
    }

    cout << GetTimestamp() << " *** " << username
        << " has joined the chat ***" << endl;

    //  tell all other users someone has joined
    string joinMessage = clientColour + username
        + " has joined the chat!" + colourReset;
    BroadcastMessage(joinMessage, clientSocket);

    //  keep receiving messages until the user disconnects
    while (true)
    {
        int messageSize = 0;
        int result = recv(clientSocket, (char*)&messageSize, sizeof(int), 0);

        //  if result is 0 or less the user has disconnected
        if (result <= 0)
        {
            break;
        }

        messageSize = ntohl(messageSize);

        char messageBuffer[4096];
        recv(clientSocket, messageBuffer, messageSize, 0);
        messageBuffer[messageSize] = '\0';
        string message(messageBuffer);

        //  handle the /users command
        if (message == "/users")
        {
            cout << GetTimestamp() << " *** " << username
                << " requested the user list ***" << endl;
            SendUserList(clientSocket);
            continue;
        }

        //  broadcast the message to all other users
        string fullMessage = clientColour + username + ": "
            + message + colourReset;

        cout << GetTimestamp() << " " << username << ": " << message << endl;

        BroadcastMessage(fullMessage, clientSocket);
    }

    //  user has disconnected > remove them from both lists
    {
        lock_guard<mutex> lock(clientsMutex);
        clients.erase(remove(clients.begin(), clients.end(), clientSocket),
            clients.end());
        usernames.erase(remove(usernames.begin(), usernames.end(), username),
            usernames.end());
    }

    //  tell all other users someone has left
    string leaveMessage = colourRed + "* " + username
        + " has left the chat! *" + colourReset;
    cout << GetTimestamp() << " *** " << username
        << " has left the chat ***" << endl;
    BroadcastMessage(leaveMessage, clientSocket);

    closesocket(clientSocket);
}



//  --- SECTION 7: Main Function ---

int main()
{
    cout << GetTimestamp()
        << " =========================================" << endl;
    cout << GetTimestamp()
        << "                  PingLink                " << endl;
    cout << GetTimestamp()
        << "                Chat Server               " << endl;
    cout << GetTimestamp()
        << " =========================================" << endl;

    //  step 1: start up winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        cout << GetTimestamp()
            << " ~~~~~ WinSock Failed to Initialise ~~~~~~" << endl;
        return -1;
    }
    cout << GetTimestamp()
        << " ~~~~~~~~~~ WinSock Initialised ~~~~~~~~~~" << endl;

    //  step 2: create the socket
    SOCKET serverSocket = socket(AF_INET, SOCK_STREAM, 0);
    if (serverSocket == INVALID_SOCKET)
    {
        cout << GetTimestamp()
            << " ~~~~~~~~ Failed to Create Socket ~~~~~~~~" << endl;
        WSACleanup();
        return -1;
    }
    cout << GetTimestamp()
        << " ~~~~~~~~~~~~~ Socket Created ~~~~~~~~~~~~" << endl;

    //  step 3: set up the address and port
    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(PORT);
    serverAddr.sin_addr.s_addr = INADDR_ANY;

    //  step 4: bind the socket to the port
    if (bind(serverSocket, (sockaddr*)&serverAddr, sizeof(serverAddr))
        == SOCKET_ERROR)
    {
        cout << GetTimestamp()
            << " ~~~~~~~~~~~~~~~ Bind Failed ~~~~~~~~~~~~~" << endl;
        closesocket(serverSocket);
        WSACleanup();
        return -1;
    }
    cout << GetTimestamp()
        << " ~~~~~~ Socket Bound to Port: " << PORT << " ~~~~~~" << endl;

    //  step 5: start listening for connections
    if (listen(serverSocket, SOMAXCONN) == SOCKET_ERROR)
    {
        cout << GetTimestamp()
            << " ~~~~~~~~~~~~~~ Listen Failed ~~~~~~~~~~~~" << endl;
        closesocket(serverSocket);
        WSACleanup();
        return -1;
    }
    cout << GetTimestamp()
        << " ~~~~ Server Listening For Connection ~~~~" << endl;
    cout << GetTimestamp()
        << " =========================================" << endl;
    cout << endl;

    //  step 6: wait for users to connect
    while (true)
    {
        SOCKET clientSocket = accept(serverSocket, nullptr, nullptr);

        if (clientSocket == INVALID_SOCKET)
        {
            cout << GetTimestamp()
                << " ~~~~~~~~~~ Failed to Accept User ~~~~~~~~" << endl;
            continue;
        }

        {
            lock_guard<mutex> lock(clientsMutex);
            clients.push_back(clientSocket);
        }

        cout << GetTimestamp()
            << " ########## New User Connected ###########" << endl;

        //  give each user their own thread
        thread clientThread(HandleClient, clientSocket);
        clientThread.detach();
    }

    //  step 7: clean up
    closesocket(serverSocket);
    WSACleanup();
    return 0;
}