-module(server).
-export([server/2, manager/1, broadcast/2, broadcast/3, worker/2, create_workers/2]).

server(Port, NumWorkers) ->
    process_flag(trap_exit, true),
    {ok, ServerSock} = gen_tcp:listen(Port, [list]),
    spawn(?MODULE, create_workers, [NumWorkers, self()]),
    supervisor(ServerSock, []).

create_workers(0, _) -> ok;
create_workers(Num, Supervisor) ->
    Supervisor ! create,
    create_workers(Num - 1, Supervisor).

supervisor(ServerSock, Pids) ->
    receive
        create ->
            io:format("Creating new worker\n"),
            NewPid = spawn_link(?MODULE, worker, [ServerSock, Pids]),
            ping_all(Pids, NewPid),
            supervisor(ServerSock, [NewPid | Pids]);
        {'EXIT', From, Reason} ->
            io:format("~p exited due to ~p\n", [From, Reason]),
            self() ! create,
            supervisor(ServerSock, lists:delete(From, Pids))
    end.

ping_all([Pid|Pids], NewPid) ->
    Pid ! { ping, NewPid },
    ping_all(Pids, NewPid);
ping_all([], _) -> ok.

worker(ServerSock, Pids) ->
    process_flag(trap_exit, false),
    ManagerPid = self(),
    spawn_link(
        fun() ->
           process_flag(trap_exit, false),
           listener(ManagerPid, ServerSock)
        end),
    manager(Pids).

listener(ManagerPid, ServerSock) ->
    case gen_tcp:accept(ServerSock) of
        {ok, Sock} ->
            ok = gen_tcp:controlling_process(Sock, ManagerPid),
            ManagerPid ! {ping, Sock},
            listener(ManagerPid, ServerSock);
        Other ->
            exit(Other)
    end.

manager(Socks) ->
    receive
        {tcp, Sock, Data} ->
            broadcast(lists:delete(Sock, Socks), Data),
            manager(Socks);
        {tcp_closed, Sock} -> 
            manager(lists:delete(Sock, Socks));
        {ping, Sock} ->
            manager([Sock | Socks]);

        {repeat, Data} ->
            broadcast(Socks, Data, norepeat),
            manager(Socks);
        Other ->
            exit(Other)
    end.

broadcast([Pid|Socks], Data, norepeat) when is_pid(Pid) ->
    broadcast(Socks, Data, norepeat);
broadcast([Pid|Socks], Data, Opt) when is_pid(Pid) ->
    Pid ! {repeat, Data},
    broadcast(Socks, Data, Opt);

broadcast([Sock|Socks], Data, Opt)  ->
    gen_tcp:send(Sock, Data),
    broadcast(Socks, Data, Opt);
broadcast([], _, _) -> ok.

broadcast(Socks, Data) ->
    broadcast(Socks, Data, ok).