#!/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 UDP 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_DGRAM)
sd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sd.bind((local, lport))

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()
        address = 'Server'
    elif sd in rfd:
        # TODO start code: receive data from clients
        # We received something
        raise NotImplementedError()
        # 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)
    else:
        # Should not happen
        print(f'Got file we did not expect! Got {rfd}', file=sys.stderr)
        continue
    msg = '{} > {}'.format(address_to_str(address), msg)
    data = bytearray(msg, 'utf8')
    # TODO start code: send message to clients
    # end code