import socket
import select
import time
import threading

class Socket:
    def __init__(self):
        self.sock = socket.socket()
    
    def connect(self, ip, port):
        self.sock.connect((ip, port))

    def create(self, port, maxUsers):
        self.sock.bind(('', port))
        self.sock.listen(maxUsers)
    
    def send(self, msg):
        self.sock.send(msg)            
        sent = self.sock.send(msg)
        if sent == 0:
            raise RuntimeError("socket connection broken")

    def forward(self, users):
        for user in users:
            print(user)

def manageClients(c, addr):
    while True:
        try:
            response = c.recv(1024)
            for user in users:
                    user.send(f"<{addr[0]}> {response.decode()}".encode())
            print(f"<{addr[0]}> {response.decode()}")
        except socket.error:
            users.pop(c)
            c.shutdown(2)
            c.close()
            for user in users:
                user.send(f"{addr[0]}:{addr[1]} lost connection".encode())
            print(f"{addr[0]}:{addr[1]} lost connection")
            break

print("1. Send\n2. Receive\n3. Server\n4. Connect\n")
opt = int(input("> "))

if opt == 1:
    try:
        clientSocket = Socket()
        targetIP = input("IP: ")
        targetPort = int(input("Port: "))
        clientSocket.connect(targetIP, targetPort)
        print(f"Connected to {targetIP}")
        while True:
            msg = input("> ")
            clientSocket.sock.send(msg.encode())
    except KeyboardInterrupt:
        clientSocket.shutdown(2)
        clientSocket.close()

elif opt == 2:
    recvPort = int(input("Port: "))
    maxUsers = int(input("Max users: "))
    srvSocket = Socket()
    srvSocket.create(recvPort, 10)
    c, addr = srvSocket.sock.accept()
    print(f"Connection from {addr[0]}")
    while True:
        response = c.recv(1024)
        if response == b'':
            raise RuntimeError("socket connection broken")
        else:
            print(f"<{addr[0]}> {response.decode()}")

elif opt == 3:
    users = {}
    recvPort = int(input("Port: "))
    srvSocket = Socket()
    srvSocket.create(recvPort, 10)
    while True:
        c, addr = srvSocket.sock.accept()
        print(f"Connection from {addr[0]}")
        users[c] = addr
        clientThread = threading.Thread(target=manageClients, args=(c, addr))
        clientThread.start()

elif opt == 4:
    try:
        clientSocket = Socket()
        targetIP = input("IP: ")
        targetPort = int(input("Port: "))
        clientSocket.connect(targetIP, targetPort)
        print(f"Connected to {targetIP}")
        while True:
            response = clientSocket.sock.recv(1024)
            if response == b'':
                raise RuntimeError("socket connection broken")
            else:
                print(response.decode())
    except KeyboardInterrupt:
        clientSocket.shutdown(2)
        clientSocket.close()

else:
    print("What the fuck?")