#include "core.h"
#define DJSTD_BASIC_ENTRY

#include "core.c"
#include "signal.h"

Server *openServer = NULL;
SocketList *openSockets = NULL;

void handleSigint(int dummy) {
    if (openServer) {
        println("");
        println("Closing server socket.");
        serverClose(openServer);
        println("Success.");
    }
    if (openSockets && openSockets->length) {
        println("");
        println("Closing open sockets.");
        for (EachEl(*openSockets, Socket, socket)) {
            socketClose(socket);
        }
        println("Success.");
    }
    signal(SIGINT, SIG_DFL);
    raise(SIGINT);
}

typedef struct ChatClient ChatClient;
struct ChatClient {
    Socket *socket;
    string nickname;
};
DefineList(ChatClient, ChatClient);

void startServer(Arena *arena, int32 port) {
    println("Starting server...");
    Server server = serverInit((ServerInitInfo){
        .concurrentClients=2,
        .port=port,
        .memory=Megabytes(64),
        .maxEvents=64,
    });
    openServer = &server;

    serverListen(&server);
    if (server.listening) {
        println("Listening on port %d", port);
    }

    Arena *serverLoopArena = arenaAlloc(Megabytes(64));
    ChatClientList chatClients = PushListZero(arena, ChatClientList, 256);

    ServerEvent *nextEvent;

    Forever {
        nextEvent = serverGetNextEvent(&server);
        switch (nextEvent->type) {
            case ServerEventType_AcceptClient: {
                Socket *client = serverAccept(&server);
                if (client != NULL) {
                    println("New client connected from %d", client->address);
                }
                break;
            };
            case ServerEventType_ClientMessage: {
                StringResult clientMsg = socketReadStr(serverLoopArena, nextEvent->tClientMessage.client);
                ChatClient *chatClient = NULL;
                if (clientMsg.valid) {
                    if (strStartsWith(clientMsg.result, s("hello-"))) {
                        StringList nickSplit = strSplit(serverLoopArena, s("-"), clientMsg.result);
                        if (nickSplit.length == 2 && nickSplit.data[1].length > 0) {
                            string newNick = PushString(arena, nickSplit.data[1].length);
                            newNick.length = nickSplit.data[1].length;
                            memcpy(newNick.str, nickSplit.data[1].str, nickSplit.data[1].length);
                            ChatClient newChatClient = (ChatClient){
                                .socket=nextEvent->tClientMessage.client,
                                .nickname=newNick,
                            };
                            AppendList(&chatClients, newChatClient);
                            println("Client from %d calls themselves \"%S\"", newChatClient.socket->address, newChatClient.nickname);
                        }
                    } else {
                        for (EachEl(chatClients, ChatClient, maybeChatClient)) {
                            if (maybeChatClient->socket->handle == nextEvent->tClientMessage.client->handle) {
                                chatClient = maybeChatClient;
                            }
                        }
                        if (chatClient != NULL) {
                            if (strStartsWith(clientMsg.result, s("say-"))) {
                                StringList saySplit = strSplit(arena, s("-"), clientMsg.result);
                                if (saySplit.length == 2 && saySplit.data[1].length > 0) {
                                    string broadcast = strPrintf(serverLoopArena, "%S says:\n%S", chatClient->nickname, saySplit.data[1]);
                                    for (EachEl(server.clients, Socket, client)) {
                                        socketWriteStr(client, broadcast);
                                    }
                                }
                            } else {
                                // Invalid client message
                            }
                        }
                    }
                }
                break;
            };
            case ServerEventType_None:
            default:
                break;
        }

        arenaFreeFrom(serverLoopArena, 0);
    }

    println("Shutting down chat.");
    serverClose(&server);
}

void clearStdInLn() {
    print("\r");
    print(ANSI_INSTRUCTION(J));
}

void clearStdInLnAfterInput() {
    print("\r");
    print(ANSI_INSTRUCTION(A));
    print(ANSI_INSTRUCTION(J));
}

void startClient(Arena *arena, string addr, int32 port, string nickname) {
    fcntl(0, F_SETFL, fcntl(0, F_GETFL) | O_NONBLOCK);

    println("Connecting to server at [%S]:%d with nickname \"%S\"", addr, port, nickname);
    Socket server = socketConnect(arena, (SocketConnectInfo){ .address=addr, .port=port });
    if (server.closed) {
        println("Connection error. Closing.");
    } else {
        println("Connected successfully");
        string message = strPrintf(arena, "hello-%S", nickname);
        CharList inputBuf = PushList(arena, CharList, 512);
        socketWriteStr(&server, message);

        print("(you)> ");
        Forever {
            Scratch scratch = scratchStart(&arena, 1);

            int32 numRead = read(0, inputBuf.data + inputBuf.length, inputBuf.capacity - inputBuf.length);
            if (numRead >= 0) {
                inputBuf.length += numRead;
                if (inputBuf.data[inputBuf.length - 1] == '\n') {
                    clearStdInLnAfterInput();
                    socketWriteStr(&server, strPrintf(scratch.arena, "say-%S", (string){.str=inputBuf.data,.length=inputBuf.length}));
                    inputBuf.length = 0;
                }
            }

            StringResult serverMsg = socketReadStr(scratch.arena, &server);
            if (serverMsg.valid && serverMsg.result.length > 0) {
                clearStdInLn();
                println("%S", serverMsg.result);
                print("(you)> %S", (string){.str=inputBuf.data, inputBuf.length});
            }

            scratchEnd(scratch);
        }
    }
    socketClose(&server);
}

int djstd_entry(Arena *arena, StringList args) {
    signal(SIGINT, &handleSigint);

    bool argumentErr = true;
    bool isServer = strEql(args.data[0], s("server"));
    bool isClient = strEql(args.data[0], s("client"));

    if (isServer) {
        Int32Result portParsed = parsePositiveInt(args.data[1]);
        if (portParsed.valid) {
            startServer(arena, portParsed.result);
            argumentErr = false;
        }
    } else if (isClient) {
        if (args.length == 3) {
            StringList split = strSplit(arena, s("]:"), args.data[1]);
            if (split.length == 2) {
                Int32Result portParsed = parsePositiveInt(split.data[1]);
                string addr = strSlice(split.data[0], 1, split.data[0].length);
                string nickname = args.data[2];
                if (portParsed.valid && addr.length > 0 && nickname.length > 0) {
                    startClient(arena, addr, portParsed.result, nickname);
                    argumentErr = false;
                }
            }
        }
    }

    if (argumentErr) {
        println("Usage:");
        println("server [PORT]");
        println("OR");
        println("client [REMOTE_ADDRESS:PORT] [NICKNAME]");
    }

    return 0;
}