#!/usr/bin/env python3
"""Minimal chat server: TCP *and* UDP on the same port number.

Server usage : python3 server.py [PORT]         (default: 9000)
Clients      : nc    HOST PORT                  (TCP)
               nc -u HOST PORT                  (UDP)

Command: /who (list of connected TCP clients)

- A TCP client has a session: it receives everyone's messages.
- A UDP client has no session: its messages are relayed to the TCP clients,
  and the server only sends it the replies to its own commands.
"""
import socket
import sys
import threading

HELP_MESSAGE = """Available commands:
/who - list of connected TCP clients
/help - show this help message
"""

HOST = "0.0.0.0"
PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 9000

tcp_clients = {}   # socket -> name: one connection = one session, with a start and an end
lock = threading.Lock()

# Same port, two protocols: the kernel tells them apart, these are two different sockets.
udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
udp_sock.bind((HOST, PORT))

tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
tcp_sock.bind((HOST, PORT))
tcp_sock.listen()


def broadcast(text: str, sender=None):
    """Send text to every TCP client except the sender."""
    msg = text.encode()
    with lock:
        for c in list(tcp_clients):
            if c is not sender:
                try:
                    c.sendall(msg)
                except OSError:
                    tcp_clients.pop(c, None)


def who() -> str:
    with lock:
        names = ", ".join(tcp_clients.values()) or "-"
    return f"* TCP connected: {names}\n"


def handle_tcp(conn, addr):
    name = f"{addr[0]}:{addr[1]}/tcp"
    conn.sendall(f"* Welcome {name} — /who\n".encode())
    with lock:
        tcp_clients[conn] = name
    print(f"* {name} joined\n", end="")
    broadcast(f"* {name} joined\n", sender=conn)
    try:
        for raw in conn.makefile("rb"):
            line = raw.decode(errors="replace")
            print(f"[{name}] {line}", end="")
            cmd = line.strip()
            if cmd == "/who":
                conn.sendall(who().encode())
            elif cmd == "/help":
                conn.sendall(HELP_MESSAGE.encode())
            elif cmd:
                broadcast(f"[{name}] {line}", sender=conn)
    except OSError:
        pass
    finally:
        with lock:
            tcp_clients.pop(conn, None)
        conn.close()
        print(f"* {name} left\n", end="")
        broadcast(f"* {name} left\n")


def udp_loop():
    while True:
        data, addr = udp_sock.recvfrom(65535)
        name = f"{addr[0]}:{addr[1]}/udp"
        line = data.decode(errors="replace")
        print(f"[{name}] {line}", end="")
        cmd = line.strip()

        if cmd == "/who":
            udp_sock.sendto(who().encode(), addr)
        elif cmd:
            broadcast(f"[{name}] {line}")


threading.Thread(target=udp_loop, daemon=True).start()
print(f"Chat listening on port {PORT} (TCP + UDP) — Ctrl-C to stop")

try:
    while True:
        conn, addr = tcp_sock.accept() # Three-way TCP handshake completed
        threading.Thread(target=handle_tcp, args=(conn, addr), daemon=True).start()
except KeyboardInterrupt:
    print("\nStopped.")