import socket
import cv2
import struct
import pickle

SERVER_IP = "localhost"   # Change to server IP
PORT = 7777

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((SERVER_IP, PORT))
cap = cv2.VideoCapture(0)  # webcam

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Compress frame as JPEG
    _, buffer = cv2.imencode(".jpg", frame)
    data = pickle.dumps(buffer, protocol=pickle.HIGHEST_PROTOCOL)
    message = struct.pack(">L", len(data)) + data
    sock.sendall(message)

    # Receive inference frame
    payload_size = struct.calcsize(">L")
    data = b""
    while len(data) < payload_size:
        data += sock.recv(4096)
    packed_msg_size = data[:payload_size]
    data = data[payload_size:]
    msg_size = struct.unpack(">L", packed_msg_size)[0]

    while len(data) < msg_size:
        data += sock.recv(4096)
    frame_data = data[:msg_size]
    data = data[msg_size:]

    annotated = pickle.loads(frame_data)
    annotated = cv2.imdecode(annotated, cv2.IMREAD_COLOR)

    cv2.imshow("YOLO Client - Inference", annotated)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

cap.release()
sock.close()
cv2.destroyAllWindows()