client.py

"""
Tiny peer client with auto-discovery.

Run:
    python client.py --id alice   # --relay-host optional to skip discovery
"""

import socket
import struct
import pickle
import threading
import time
import math
import argparse
import queue
import cv2
import numpy as np

try:
    import sounddevice as sd
    HAS_AUDIO = True
except Exception:
    HAS_AUDIO = False

RELAY_PORT = 9000
BROADCAST_PORT = 54545
BROADCAST_TOKEN = b"VIDEO_RELAY"
DISCOVERY_TIMEOUT = 5.0
QUALITY = 60
TILE_W, TILE_H = 320, 240
AUDIO_RATE = 16000
AUDIO_BLOCK = 1024
AUDIO_CH = 1
CAM_INDEX = 0


def find_relay(timeout=DISCOVERY_TIMEOUT):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("", BROADCAST_PORT))
    sock.settimeout(timeout)
    try:
        while True:
            data, addr = sock.recvfrom(1024)
            if data.startswith(BROADCAST_TOKEN):
                parts = data.split()
                if len(parts) == 2:
                    try:
                        port = int(parts[1])
                    except ValueError:
                        continue
                    return addr[0], port
    except socket.timeout:
        return None
    finally:
        sock.close()


def recv_exact(sock, size):
    buf = b""
    while len(buf) < size:
        chunk = sock.recv(size - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


class Peer:
    def __init__(self, host, port, peer_id, quality=QUALITY, use_audio=True, cam_index=CAM_INDEX):
        self.host = host
        self.port = port
        self.id = peer_id
        self.quality = quality
        self.sock = None
        self.running = False
        self.frames = {}  # peer_id -> (ts, frame)
        self.frames_lock = threading.Lock()
        self.local_frame = None
        self.audio_send_q = queue.Queue(maxsize=50)
        self.audio_play_q = queue.Queue(maxsize=50)
        self.audio_in = None
        self.audio_out = None
        self.debug = True
        self.use_audio = use_audio and HAS_AUDIO
        self.cam_index = cam_index

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, self.port))
        print(f"Connected to relay at {self.host}:{self.port}")

    def start_camera(self):
        cam = cv2.VideoCapture(self.cam_index)
        if not cam.isOpened():
            raise RuntimeError("Could not open camera (index {})".format(self.cam_index))
        cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cam.set(cv2.CAP_PROP_FPS, 30)
        return cam

    def start_audio(self):
        if not self.use_audio:
            print("Audio disabled or sounddevice unavailable.")
            return
        def _in_callback(indata, frames, time_info, status):
            try:
                self.audio_send_q.put_nowait(bytes(indata))
            except queue.Full:
                pass

        try:
            self.audio_out = sd.RawOutputStream(
                samplerate=AUDIO_RATE,
                blocksize=AUDIO_BLOCK,
                channels=AUDIO_CH,
                dtype="int16",
            )
            self.audio_in = sd.RawInputStream(
                samplerate=AUDIO_RATE,
                blocksize=AUDIO_BLOCK,
                channels=AUDIO_CH,
                dtype="int16",
                callback=_in_callback,
            )
            self.audio_out.start()
            self.audio_in.start()
        except Exception as e:
            print(f"Audio init failed ({e}); continuing without audio.")
            self.use_audio = False

    def send_video_loop(self, cam):
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), self.quality]
        while self.running:
            ok, frame = cam.read()
            if not ok:
                continue
            ok, enc = cv2.imencode(".jpg", frame, encode_param)
            if not ok:
                continue
            payload = pickle.dumps({"kind": "video", "id": self.id, "frame": enc})
            try:
                self.sock.sendall(struct.pack("!I", len(payload)) + payload)
            except Exception as e:
                if self.debug:
                    print(f"[video] send error: {e}")
                break
            with self.frames_lock:
                self.local_frame = frame.copy()
            time.sleep(0.05)
        self.running = False

    def send_audio_loop(self):
        while self.running:
            try:
                chunk = self.audio_send_q.get(timeout=0.5)
            except queue.Empty:
                continue
            payload = pickle.dumps({"kind": "audio", "id": self.id, "audio": chunk})
            try:
                self.sock.sendall(struct.pack("!I", len(payload)) + payload)
            except Exception as e:
                if self.debug:
                    print(f"[audio] send error: {e}")
                break
        self.running = False

    def recv_loop(self):
        while self.running:
            raw = recv_exact(self.sock, 4)
            if not raw:
                if self.debug:
                    print("[recv] connection closed or timed out")
                break
            size = struct.unpack("!I", raw)[0]
            data = recv_exact(self.sock, size)
            if not data:
                if self.debug:
                    print("[recv] failed to read payload")
                break
            try:
                pkt = pickle.loads(data)
            except Exception:
                continue
            kind = pkt.get("kind", "video")
            sender = pkt.get("id")
            if kind == "audio":
                audio = pkt.get("audio")
                if audio:
                    try:
                        self.audio_play_q.put_nowait(audio)
                    except queue.Full:
                        pass
            else:
                enc = pkt.get("frame")
                frame = cv2.imdecode(enc, cv2.IMREAD_COLOR) if enc is not None else None
                if sender and frame is not None:
                    with self.frames_lock:
                        self.frames[sender] = (time.time(), frame)
        self.running = False

    def play_audio_loop(self):
        while self.running:
            try:
                chunk = self.audio_play_q.get(timeout=0.5)
            except queue.Empty:
                continue
            try:
                self.audio_out.write(chunk)
            except Exception:
                break
        self.running = False

    def build_grid(self):
        now = time.time()
        with self.frames_lock:
            items = [(pid, f) for pid, (ts, f) in self.frames.items() if now - ts < 5]
            if self.local_frame is not None:
                items = [(self.id, self.local_frame)] + items
        if not items:
            return None
        count = len(items)
        cols = math.ceil(math.sqrt(count))
        rows = math.ceil(count / cols)
        tiles = []
        for pid, frame in items:
            img = cv2.resize(frame, (TILE_W, TILE_H))
            cv2.rectangle(img, (0, 0), (TILE_W, 25), (0, 0, 0), -1)
            cv2.putText(img, pid, (10, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
            tiles.append(img)
        blank = np.zeros((TILE_H, TILE_W, 3), dtype=np.uint8)
        while len(tiles) < rows * cols:
            tiles.append(blank)
        rows_img = []
        for r in range(rows):
            rows_img.append(np.hstack(tiles[r * cols:(r + 1) * cols]))
        return np.vstack(rows_img)

    def start(self):
        try:
            self.connect()
            cam = self.start_camera()
            self.start_audio()
        except Exception as e:
            print(f"Startup failed: {e}")
            return
        self.running = True
        threading.Thread(target=self.send_video_loop, args=(cam,), daemon=True).start()
        if self.use_audio:
            threading.Thread(target=self.send_audio_loop, daemon=True).start()
        threading.Thread(target=self.recv_loop, daemon=True).start()
        if self.use_audio:
            threading.Thread(target=self.play_audio_loop, daemon=True).start()
        try:
            while self.running:
                grid = self.build_grid()
                if grid is not None:
                    cv2.imshow("Group Video", grid)
                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break
        except KeyboardInterrupt:
            pass
        self.running = False
        cam.release()
        if self.audio_in:
            try:
                self.audio_in.stop()
                self.audio_in.close()
            except:
                pass
        if self.audio_out:
            try:
                self.audio_out.stop()
                self.audio_out.close()
            except:
                pass
        if self.sock:
            try:
                self.sock.close()
            except:
                pass
        cv2.destroyAllWindows()
        print("Peer stopped.")


def main():
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument("--id", default="peer", help="Name label to show in the grid")
    parser.add_argument("--relay-host", default=None, help="Skip discovery and use this host")
    parser.add_argument("--relay-port", type=int, default=RELAY_PORT, help="Relay port (when using --relay-host)")
    parser.add_argument("--timeout", type=float, default=DISCOVERY_TIMEOUT, help="Discovery timeout seconds")
    parser.add_argument("--no-audio", action="store_true", help="Disable audio capture/playback")
    parser.add_argument("--camera", type=int, default=CAM_INDEX, help="Camera index (default 0)")
    args = parser.parse_args()

    host, port = args.relay_host, args.relay_port
    if host is None:
        print(f"Listening for relay beacon (timeout {args.timeout}s)...")
        found = find_relay(timeout=args.timeout)
        if not found:
            print("No relay found. Start `python server.py` on the LAN or pass --relay-host manually.")
            return
        host, port = found
        print(f"Found relay at {host}:{port}")
    else:
        print(f"Using relay {host}:{port}")

    peer = Peer(host, port, args.id, use_audio=not args.no_audio, cam_index=args.camera)
    peer.start()


if __name__ == "__main__":
    main()