#!/usr/bin/env python3
"""
p2p_chat.py – minimal P2P UDP chat with STUN + HTTP signalling
--------------------------------------------------------------
• Discovers public address/port via STUN
• Registers with a super-simple signalling server
• UDP hole-punches the remote peer
• Provides a tiny chat interface (type /quit to exit)

Dependencies:
    pip install pystun3 requests
"""

import logging
import socket
import threading
import time
from datetime import datetime

import requests
import stun  # pystun3

SIGNAL_URL = "http://77.247.92.168:8000"
N_PROBES = 10
PROBE_INTERVAL = 0.2

# ---------------------------------------------------------------------------

def timestamp() -> str:
    return datetime.utcnow().strftime("%H:%M:%S")

def setup_logging() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )

# ---------------------------------------------------------------------------

def discover_public_addr() -> tuple[str, int]:
    """Use STUN to learn our public (IP, port)."""
    nat_type, public_ip, public_port = stun.get_ip_info()
    logging.info(f"STUN: NAT type = {nat_type}")
    logging.info(f"STUN: public address = {public_ip}:{public_port}")
    return public_ip, public_port

def register_with_server(my_id: str, ip: str, port: int) -> None:
    """POST /register until the server replies 200 OK."""
    payload = {"id": my_id, "ip": ip, "port": port}
    while True:
        try:
            r = requests.post(f"{SIGNAL_URL}/register", json=payload, timeout=5)
            if r.status_code == 200:
                logging.info(f"Registered as '{my_id}'")
                return
            logging.warning(f"Register failed: HTTP {r.status_code}")
        except Exception as e:
            logging.warning(f"Register exception: {e}")
        time.sleep(1)

def fetch_peer_info(peer_id: str) -> tuple[str, int]:
    """Poll /peer/<id> until it appears; return (ip, port)."""
    logging.info(f"Waiting for peer '{peer_id}' to register …")
    while True:
        try:
            r = requests.get(f"{SIGNAL_URL}/peer/{peer_id}", timeout=5)
            if r.status_code == 200:
                info = r.json()
                ip, port = info["ip"], int(info["port"])
                logging.info(f"Got peer '{peer_id}': {ip}:{port}")
                return ip, port
            elif r.status_code != 404:
                logging.warning(f"Lookup error: HTTP {r.status_code}")
        except Exception as e:
            logging.warning(f"Lookup exception: {e}")
        time.sleep(1)

# ---------------------------------------------------------------------------

def recv_loop(sock: socket.socket):
    """Background thread to print incoming UDP messages."""
    while True:
        try:
            data, addr = sock.recvfrom(4096)
            print(f"\n{timestamp()} <{addr[0]}:{addr[1]}> {data.decode()}")
        except OSError:
            break          # socket was closed
        except Exception as e:
            logging.debug(f"recv_loop error: {e}")

def punch_hole(sock: socket.socket, peer_addr: tuple[str, int]) -> None:
    """Send rapid dummy packets to create NAT mappings on both sides."""
    logging.info("Punching a hole through the NAT …")
    for _ in range(N_PROBES):
        sock.sendto(b"probe", peer_addr)
        time.sleep(PROBE_INTERVAL)
    logging.info("Hole-punch phase complete – start chatting!")

# ---------------------------------------------------------------------------

def main() -> None:
    setup_logging()

    my_id = input("Your ID   : ").strip()
    peer_id = input("Peer's ID : ").strip()

    public_ip, public_port = discover_public_addr()
    register_with_server(my_id, public_ip, public_port)
    peer_ip, peer_port = fetch_peer_info(peer_id)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("", 0))
    logging.info(f"Local UDP socket bound on {sock.getsockname()}")

    threading.Thread(target=recv_loop, args=(sock,), daemon=True).start()

    peer_addr = (peer_ip, peer_port)
    punch_hole(sock, peer_addr)

    try:
        while True:
            msg = input()
            if msg.lower() in ("/quit", "/exit"):
                break
            if msg:
                sock.sendto(msg.encode(), peer_addr)
    finally:
        sock.close()
        logging.info("Goodbye!")

# ---------------------------------------------------------------------------

if __name__ == "__main__":
    main()