#!/usr/bin/python3

import argparse
import sys
import socket
import select
import string


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


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)

lport = args.lport
rport = args.rport
raddr = args.remote
family, ty, proto, canonname, sockaddr = socket.getaddrinfo(raddr, rport, proto=socket.IPPROTO_UDP)[0]

# TODO start code here: create and bind socket
sd = socket.socket(family, ty)
sd.bind(('', lport))
# TODO end code here

rfds = [sd, sys.stdin]

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

    if sys.stdin in rfd:
        # TODO start code here
        # We have typed something
        msg = sys.stdin.readline()
        if len(msg) == 0:
            rfds.remove(sys.stdin)
            continue
        data = bytearray(msg, 'utf8')
        sd.sendto(data, sockaddr)
        # end code here
    elif sd in rfd:
        # TODO start code here
        # We received something
        data, address = sd.recvfrom(4096)
        try:
            msg = data.decode('utf8').rstrip('\n')
        except:
            # Ignore invalid message
            continue
        print_msg(msg)
        # end code here
    else:
        # Should not happen
        print(f'Got file we did not expect! Got {rfd}', file=sys.stderr)
        continue