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

#include <winsock2.h>
#include <ws2tcpip.h>
#include <iostream>
#include <string>
#include <thread>
#include <chrono>

using namespace std;

#define SERVER_IP "127.0.0.1"
#define PORT 54000

//  flag to track connection
bool running = true;

//  ANSI RGB colour codes for the users messages
const string colourPurple = "\033[38;2;180;150;220m";
const string colourReset = "\033[0m";



//  --- SECTION 2: Reliable Send Function ---

//  keeps sending until all bytes have been 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 3: Receive Messages Function ---

//  runs on a different thread and prints incoming messages
void ReceiveMessages(SOCKET serverSocket)
{
    while (running)
    {
        //  receive the size of the incoming message
        int messageSize = 0;
        int result = recv(serverSocket, (char*)&messageSize, sizeof(int), 0);

        //  if result is 0 or less the server has been disconnected
        if (result <= 0)
        {
            if (running)
            {
                cout << "\r\033[2K" << endl;
                cout << "\r\033[2K******************************" << endl;
                cout << "     Server Disconnected!     " << endl;
                cout << "******************************" << endl;
                running = false;
            }
            break;
        }

        messageSize = ntohl(messageSize);

        //  receive the message
        char buffer[4096];
        recv(serverSocket, buffer, messageSize, 0);
        buffer[messageSize] = '\0';
        string message(buffer);

        if (message == "ACCEPTED")
        {
            continue;
        }

        //  handle the /users command output
        if (message.find("Connected Users:") == 0)
        {
            cout << "\r\033[2K" << message << endl;
        }
        else
        {
            cout << "\r\033[2K" << message << endl;
            cout << colourPurple << "You: " << colourReset;
            cout.flush();
        }
    }
}



//  --- SECTION 4: Main Function ---

int main()
{
    //  startup message
    cout << "==============================" << endl;
    cout << "           PingLink           " << endl;
    cout << "       Chat Application       " << endl;
    cout << "==============================" << endl;
    cout << endl;

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

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

    //  step 3: set up the server address
    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(PORT);
    inet_pton(AF_INET, SERVER_IP, &serverAddr.sin_addr);

    //  connecting loading bar
    cout << "     Connecting to Server     " << endl;
    cout << "[";
    for (int i = 0; i < 28; i++)
    {
        cout << "#";
        cout.flush();
        this_thread::sleep_for(chrono::milliseconds(100));
    }
    cout << "]" << endl;
    cout << endl;

    //  step 4: connect to the server
    if (connect(serverSocket, (sockaddr*)&serverAddr, sizeof(serverAddr))
        == SOCKET_ERROR)
    {
        cout << "==============================" << endl;
        cout << " No Connection To The Server! " << endl;
        cout << "==============================" << endl;
        closesocket(serverSocket);
        WSACleanup();
        return -1;
    }

    cout << "==============================" << endl;
    cout << "      Status : Connected      " << endl;
    cout << "==============================" << endl;
    cout << endl;

    //  step 5: ask the user to enter their username
    string username;
    while (true)
    {
        cout << "   Enter Username: ";
        getline(cin, username);

        //  ANSI cursor codes to erase the username line after entry
        cout << "\033[1A";
        cout << "\033[2K";
        cout << "\033[1A";
        cout << "\033[2K";
        cout.flush();

        //  if the username is blank, show an error and ask again
        if (username.empty())
        {
            cout << "*  Username Cannot Be Blank  *" << endl;
            continue;
        }

        //  send the username to the server
        WriteSize(serverSocket, (int)username.length());
        Writen(serverSocket, username.c_str(), (int)username.length());

        //  step 6: wait for the server to respond
        int responseSize = 0;
        int recvCheck = recv(serverSocket, (char*)&responseSize, sizeof(int), 0);

        if (recvCheck <= 0)
        {
            cout << endl;
            cout << "==============================" << endl;
            cout << "     Server Disconnected!     " << endl;
            cout << "==============================" << endl;
            closesocket(serverSocket);
            WSACleanup();
            return -1;
        }

        responseSize = ntohl(responseSize);

        char responseBuffer[256];
        recv(serverSocket, responseBuffer, responseSize, 0);
        responseBuffer[responseSize] = '\0';
        string response(responseBuffer);

        //  if rejected, show the reason and ask again
        if (response.find("REJECTED:") == 0)
        {
            string reason = response.substr(9);
            cout << "*   " << reason << "   *" << endl;
            continue;
        }

        //  server accepts the username
        break;
    }

    //  show the welcome panel
    cout << endl;
    cout << "        Welcome, " << username << "!" << endl;
    cout << "------------------------------" << endl;
    cout << "     /users == View Users     " << endl;
    cout << "     /exit  == Disconnect     " << endl;
    cout << "------------------------------" << endl;
    cout << endl;

    //  step 7: start the receive thread
    thread receiveThread(ReceiveMessages, serverSocket);
    receiveThread.detach();

    //  step 8: main loop for sending messages
    while (running)
    {
        cout << colourPurple << "You: " << colourReset;
        string message;
        getline(cin, message);

        //  ANSI cursor codes to move up and erase the line
        if (message.empty())
        {
            cout << "\033[1A\033[2K";
            cout.flush();
            continue;
        }

        //  /exit command
        if (message == "/exit")
        {
            cout << endl;
            cout << "==============================" << endl;
            cout << "     Thank You For Using      " << endl;
            cout << "           PingLink           " << endl;
            cout << "==============================" << endl;
            running = false;
            break;
        }

        //  /users command
        if (message == "/users")
        {
            WriteSize(serverSocket, (int)message.length());
            Writen(serverSocket, message.c_str(), (int)message.length());
            this_thread::sleep_for(chrono::milliseconds(100));
            continue;
        }

        //  ANSI cursor codes to move up and erase the text line
        cout << "\033[1A\033[2K";
        cout << colourPurple << "You: " << message << colourReset << endl;

        //  send the message to the server
        WriteSize(serverSocket, (int)message.length());
        Writen(serverSocket, message.c_str(), (int)message.length());
    }

    //  step 9: clean up and close
    closesocket(serverSocket);
    WSACleanup();

    return 0;
}