#!/usr/bin/python3

import sys
import socket
import select
import string


#!/usr/bin/python3

import sys
import socket
import select
import string
import argparse

def shutdown():
    sd.close()
    exit(0)

def print_msg(msg):
        msg = ''.join(filter(lambda x: x in string.printable, msg)) if not msg.isprintable() else msg
        msg = '\033[0;36m{}:{} \033[0;31m>>\033[0m {}'.format(address[0], address[1], msg)
        print(msg)

def address_to_str(addr):
    if addr == 'Server':
        return 'Server'
    elif ':' in addr[0]:
        return f'[{addr[0]}]:{addr[1]}'
    else:
        return f'{addr[0]}:{addr[1]}'

parser = argparse.ArgumentParser(
        description="Simple chat server using TCP sockets"
    )
parser.add_argument(
        "lport",
        type=int,
        help="The local port to listen on for incoming chat messages",
    )
parser.add_argument(
        "--local",
        type=str,
        help="The local address to send outgoing chat messages to",
        default=""
    )
args = parser.parse_args()
lport = args.lport
local = args.local

sd = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
sd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sd.bind((local, lport))
sd.listen(25)

rfds = {sd, sys.stdin}

clients = set()

while True:
    address = None
    msg = None

    try:
        rfd, _, _ = select.select(rfds, [], [])
    except KeyboardInterrupt:
        shutdown()

    if sys.stdin in rfd:
        # We have typed something
        msg = sys.stdin.readline().rstrip('\n')
        address = 'Server'
    elif sd in rfd:
        # TODO start code: accept client connection request
        address = 'Server'
        try:
            csock, address = sd.accept()
        except Exception as e:
            print(f'Failed to accept connection request: {str(e)}')
            continue
        clients.add(csock)
        rfds.add(csock)
        msg = f'{address_to_str(address)} joined the chatroom.'
        print_msg(msg)
        # end code
    else:
        csock = rfd[0]
        # TODO start code: handle client message
        data = csock.recv(4096)
        if not data:
            # Client disconnected/EOF
            clients.remove(csock)
            rfds.remove(csock)
        address = csock.getpeername()
        if csock not in clients:
            clients.add(csock)
        # end code
        try:
            msg = data.decode('utf8')
        except Exception as e:
            # Ignore invalid message
            print(f'An error occured decoding data to send: {str(e)}', file=sys.stderr)
            continue
        msg = ''.join(filter(lambda x: x in string.printable, msg)).rstrip('\n')
        print_msg(msg)
    msg = '{} > {}\n'.format(address_to_str(address), msg)
    data = bytearray(msg, 'utf8')
    # TODO start code: send message to clients
    for csock in clients:
        if csock == rfd:
            continue
        csock.send(data)
    # end code