-module(server).

-export([start_server/0]).

-include_lib("./defs.hrl").

-spec start_server() -> _.
-spec loop(_State) -> _.
-spec do_join(_ChatName, _ClientPID, _Ref, _State) -> _.
-spec do_leave(_ChatName, _ClientPID, _Ref, _State) -> _.
-spec do_new_nick(_State, _Ref, _ClientPID, _NewNick) -> _.
-spec do_client_quit(_State, _Ref, _ClientPID) -> _NewState.

start_server() ->
    catch(unregister(server)),
    register(server, self()),
    case whereis(testsuite) of
	undefined -> ok;
	TestSuitePID -> TestSuitePID!{server_up, self()}
    end,
    loop(
      #serv_st{
	 nicks = maps:new(), %% nickname map. client_pid => "nickname"
	 registrations = maps:new(), %% registration map. "chat_name" => [client_pids]
	 chatrooms = maps:new() %% chatroom map. "chat_name" => chat_pid
	}
     ).

loop(State) ->
    receive 
	%% initial connection
	{ClientPID, connect, ClientNick} ->
	    NewState =
		#serv_st{
		   nicks = maps:put(ClientPID, ClientNick, State#serv_st.nicks),
		   registrations = State#serv_st.registrations,
		   chatrooms = State#serv_st.chatrooms
		  },
	    loop(NewState);
	%% client requests to join a chat
	{ClientPID, Ref, join, ChatName} ->
	    NewState = do_join(ChatName, ClientPID, Ref, State),
	    loop(NewState);
	%% client requests to join a chat
	{ClientPID, Ref, leave, ChatName} ->
	    NewState = do_leave(ChatName, ClientPID, Ref, State),
	    loop(NewState);
	%% client requests to register a new nickname
	{ClientPID, Ref, nick, NewNick} ->
	    NewState = do_new_nick(State, Ref, ClientPID, NewNick),
	    loop(NewState);
	%% client requests to quit
	{ClientPID, Ref, quit} ->
	    NewState = do_client_quit(State, Ref, ClientPID),
	    loop(NewState);
	{TEST_PID, get_state} ->
	    TEST_PID!{get_state, State},
	    loop(State)
    end.

%% executes join protocol from server perspective
do_join(ChatName, ClientPID, Ref, State) ->
    Chatrooms = State#serv_st.chatrooms,
    Registrations = State#serv_st.registrations,
    Nicks = State#serv_st.nicks,

    case maps:get(ChatName, Chatrooms, undefined) of
        undefined ->
            % Chatroom doesn't exist, make a new one 
            NewChatPid = spawn(chatroom, start_chatroom, [ChatName]),
            {ok, ClientNick} = maps:find(ClientPID, Nicks),
            NewChatPid!{self(), Ref, register, ClientPID, ClientNick},

            % Update state with the new chatroom
            UpdatedState = State#serv_st{
                chatrooms = maps:put(ChatName, NewChatPid, Chatrooms),
                registrations = maps:put(ChatName, [ClientPID], Registrations)
            },
            UpdatedState;

        ExistingChatPid ->
            % Chatroom exists, register the client to it
            {ok, ClientNick} = maps:find(ClientPID, Nicks),
            ExistingChatPid!{self(), Ref, register, ClientPID, ClientNick},

            % Update state with the new client added to the chatroom
            {ok, ExistingRegs} = maps:find(ChatName, Registrations),
            UpdatedState = State#serv_st{
                registrations = maps:put(ChatName, lists:append(ExistingRegs, [ClientPID]), Registrations)
            },
            UpdatedState
    end.


%% executes leave protocol from server perspective
do_leave(ChatName, ClientPID, Ref, State) ->
    case maps:is_key(ChatName, State#serv_st.chatrooms) of
        false ->
            ClientPID!{self(), Ref, err},
            State;
        true ->
            Clients = maps:get(ChatName, State#serv_st.registrations, []),
            case lists:member(ClientPID, Clients) of
                false ->
                    ClientPID!{self(), Ref, err},
                    State;
                true ->
                    UpdatedClients = lists:delete(ClientPID, Clients),
                    UpdatedRegistrations = maps:put(ChatName, UpdatedClients, State#serv_st.registrations),
                    ChatroomPID = maps:get(ChatName, State#serv_st.chatrooms),
                    ChatroomPID!{self(), Ref, unregister, ClientPID},
                    ClientPID!{self(), Ref, ack_leave},
                    State#serv_st{registrations = UpdatedRegistrations}
            end
    end.

%% executes new nickname protocol from server perspective
do_new_nick(State, Ref, ClientPID, NewNick) ->
    CurrentNick = maps:get(ClientPID, State#serv_st.nicks),

    case CurrentNick of
        NewNick ->
            ClientPID!{self(), Ref, err_same},
            State;
        _ ->
            ExistingNicks = maps:values(State#serv_st.nicks),
            case lists:member(NewNick, ExistingNicks) of
                true ->
                    ClientPID!{self(), Ref, err_nick_used},
                    State;
                false ->
                    UpdatedNicks = maps:put(ClientPID, NewNick, State#serv_st.nicks),
                    Registrations = State#serv_st.registrations,
                    lists:foreach(
                        fun(ChatName) ->
                            ChatroomPID = maps:get(ChatName, State#serv_st.chatrooms),
                            ChatroomPID!{self(), Ref, update_nick, ClientPID, NewNick}
                        end,
                        maps:keys(Registrations)
                    ),

                    ClientPID!{self(), Ref, ok_nick},
                    State#serv_st{nicks = UpdatedNicks}
            end
    end.

%% executes client quit protocol from server perspective
do_client_quit(State, Ref, ClientPID) ->
    _ClientNick = maps:get(ClientPID, State#serv_st.nicks, undefined),
    Registrations = State#serv_st.registrations,
    RegisteredChatrooms = [ChatName || {ChatName, Clients} <- maps:to_list(Registrations), lists:member(ClientPID, Clients)],

    lists:foreach(
        fun(ChatName) ->
            ChatroomPID = maps:get(ChatName, State#serv_st.chatrooms),
            ChatroomPID!{self(), Ref, unregister, ClientPID}
        end,
        RegisteredChatrooms
    ),

    UpdatedRegistrations =
        lists:foldl(
            fun(ChatName, Acc) ->
                CurrentClients = maps:get(ChatName, Acc),
                UpdatedClients = lists:delete(ClientPID, CurrentClients),
                maps:put(ChatName, UpdatedClients, Acc)
            end,
            Registrations,
            RegisteredChatrooms
        ),

    UpdatedNicks = maps:remove(ClientPID, State#serv_st.nicks),
    ClientPID!{self(), Ref, ack_quit},

    State#serv_st{
        nicks = UpdatedNicks,
        registrations = UpdatedRegistrations
    }.