import socket
import sys


mode = 'quit'
while True:
    resp = input('Run as server or client (Ss/Cc): ')
    if resp in ['S', 's']:
        mode = 'server'
        break
    if resp in ['C', 'c']:
        mode = 'client'
        break
    if resp in ['Q', 'q']:
        mode = 'quit'
        break

if mode == 'quit':
    sys.exit(0)

if mode == 'server':
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(('0.0.0.0', 8081))
    server_socket.listen()
    print('Server: Waiting for connection')
    connection, address = server_socket.accept()
    print('Server: Client connected, address:', address)
else:
    address = input('Client: What is the servers IP address? (default: 127.0.0.1): ')
    if not address:
        address = '127.0.0.1'
    connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    connection.connect((address, 8081))
    print('Client: Connected to server at', address)
    name = input('Client: enter your chat name: ')
    if not name:
        name = 'anonymous'
    connection.send(bytes(name + " connected to the server", 'utf-8'))

if mode == 'server':
    prompt = 'Server: '
else:
    prompt = 'Client: '

while True:
    message = connection.recv(1024)
    if not message:
        break
    print(prompt, message.decode())
    new_message = input(prompt + 'type a message: ')
    if new_message == 'quit':
        break
    connection.send(new_message.encode())

print(prompt + 'stopped')
connection.close()
if mode == "server":
    server_socket.close()