#!/usr/bin/env python3
"""TCP / UDP chat client, API - Teleinfo, TP 04 (STUDENT VERSION).

Usage: python3 client.py HOST PORT tcp|udp

Two parts of this file must be completed.
Everything else is provided and already works.
"""
import os
import re
import readline
import signal
import socket
import sys
import threading
import zlib
from datetime import datetime

# ---------------------------------------------------------------------------
# Colour output (ANSI codes)
# ---------------------------------------------------------------------------
RESET, BOLD, DIM = "\033[0m", "\033[1m", "\033[2m"
RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN = (f"\033[{n}m" for n in (31, 32, 33, 34, 35, 36))
SENDER_COLORS = [BLUE, GREEN, YELLOW, MAGENTA, CYAN, RED]
BADGE = {"tcp": "\033[30;42m TCP \033[0m", "udp": "\033[30;45m UDP \033[0m"}
PROMPT = "> "

MSG_RE = re.compile(r"^\[(?P<who>[^\]]+?)/(?P<proto>tcp|udp)\] ?(?P<text>.*)$")

typing = False                  # True while the user is typing a line
screen_lock = threading.Lock()


def show(text):
    """Print a line above the input line, without erasing what is being typed."""
    stamp = f"{DIM}{datetime.now():%H:%M:%S}{RESET}"
    with screen_lock:
        if typing:
            sys.stdout.write(f"\r\033[K{stamp} {text}\n{PROMPT}{readline.get_line_buffer()}")
        else:
            sys.stdout.write(f"{stamp} {text}\n")
        sys.stdout.flush()


def format_line(line):
    """Format a line received from the server."""
    m = MSG_RE.match(line)
    if m:
        color = SENDER_COLORS[zlib.crc32(m["who"].encode()) % len(SENDER_COLORS)]
        return f"{BADGE[m['proto']]} {BOLD}{color}{m['who']}{RESET} › {m['text']}"
    if line.startswith("* ") and line.endswith("joined"):
        return f"{GREEN}-> {line[2:]}{RESET}"
    if line.startswith("* ") and line.endswith("left"):
        return f"{RED}<- {line[2:]}{RESET}"
    if line.startswith("* "):
        return f"{YELLOW}- {line[2:]}{RESET}"
    return line

def receive_loop(sock, proto, server):
    """Loop that receives messages from the server and displays them."""
    buffer = b""
    while True:
        try:
            if proto == "tcp":
                data = sock.recv(4096)
                if not data:                    # 0 bytes received = the server sent a FIN
                    show(f"{RED}! the server closed the connection (FIN received){RESET}")
                    os.kill(os.getpid(), signal.SIGINT)
                    return
                # TCP is a stream: we split it into lines ourselves.
                buffer += data
                *lines, buffer = buffer.split(b"\n")
            else:
                data, addr = sock.recvfrom(65535)   # UDP: one whole datagram at a time
                if addr != server:                  # anyone can write to our port
                    show(f"{RED}! datagram ignored, coming from {addr[0]}:{addr[1]}{RESET}")
                    continue
                lines = data.split(b"\n")
        except OSError:
            return
        for line in lines:
            if line:
                show(format_line(line.decode(errors="replace")))


# ---------------------------------------------------------------------------
# TODO 1 — The socket type
#
# A "socket" is a program's doorway to the network.
# When it is created, its type is chosen among two options:
#
#     socket.SOCK_DGRAM
#     socket.SOCK_STREAM
#
# Replace each None with the right type (hint : https://www.ibm.com/docs/en/aix/7.1.0?topic=protocols-socket-types)
# ---------------------------------------------------------------------------
TYPE_TCP = None
TYPE_UDP = None


def create_socket(proto, server):
    if proto == "tcp":
        sock = socket.socket(socket.AF_INET, TYPE_TCP)
        sock.connect(server)                      # SYN / SYN-ACK / ACK handshake
    else:
        sock = socket.socket(socket.AF_INET, TYPE_UDP)
        # No connect(): there is no connection in UDP.
    return sock


# ---------------------------------------------------------------------------
# TODO 2 — Sending a message over UDP
#
# Replace each "???" with one of the appropriate values: "proto", "server" or "sock"
# ---------------------------------------------------------------------------
def send_line(sock, proto, server, text):
    data = text.encode()
    if ??? == "tcp":
        sock.sendall(data)
    else:
        ???.sendto(data, ???)


# ---------------------------------------------------------------------------
# Main program
# ---------------------------------------------------------------------------
def fail(text):
    sys.exit(f"{BOLD}{RED}✗ {text}{RESET}")


def main():
    global typing
    if len(sys.argv) != 4 or sys.argv[3] not in ("tcp", "udp"):
        sys.exit("Usage: python3 client.py HOST PORT tcp|udp")
    proto = sys.argv[3]
    server = (socket.gethostbyname(sys.argv[1]), int(sys.argv[2]))

    if TYPE_TCP is None or TYPE_UDP is None:
        fail("TODO 1 not completed yet: TYPE_TCP and TYPE_UDP are still None")
    if TYPE_TCP == TYPE_UDP:
        fail("TODO 1: TCP and UDP cannot use the same socket type")
    if TYPE_TCP != socket.SOCK_STREAM:
        fail("TODO 1: the types seem swapped. Which of the two sets up a connection?")

    try:
        sock = create_socket(proto, server)
    except ConnectionRefusedError:
        fail(f"Connection refused: no TCP service on {server[0]}:{server[1]} (RST received)")
    except OSError as e:
        fail(f"Network error: {e}")

    port = sock.getsockname()[1]
    print(f"\n{BADGE[proto]} {BOLD}Chat {proto.upper()}{RESET}  server {server[0]}:{server[1]}"
          f"  -  local port {port if port else 'not assigned yet'}")
    print(f"{DIM}Ctrl-D: quit{RESET}")
    print()

    threading.Thread(target=receive_loop, args=(sock, proto, server), daemon=True).start()

    try:
        while True:
            typing = True
            text = input(PROMPT)
            typing = False
            sys.stdout.write("\033[A\r\033[K")    # erase the typed line, shown again in colour
            if not text.strip():
                continue
            send_line(sock, proto, server, text + "\n")
            show(f"{BADGE[proto]} {BOLD}{CYAN}me{RESET} › {text}")
            if not port:
                port = sock.getsockname()[1]
    except (EOFError, KeyboardInterrupt):
        pass
    finally:
        sock.close()      # TCP: sends a FIN. UDP: sends nothing.
    print("\nSession ended.")


if __name__ == "__main__":
    main()