import socket
import struct

import cv2
import numpy as np


LISTEN_IP = "0.0.0.0"
LISTEN_PORT = 5000

HEADER_SIZE = 4
MAX_FRAME_SIZE = 20 * 1024 * 1024  # 20 MB safety limit


def receive_exact(sock, size):
    """
    Receive exactly 'size' bytes.

    TCP recv() may return fewer bytes than requested,
    so we keep receiving until the full frame arrives.
    """
    data = bytearray()

    while len(data) < size:
        packet = sock.recv(size - len(data))

        if not packet:
            return None

        data.extend(packet)

    return bytes(data)


def main():
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((LISTEN_IP, LISTEN_PORT))
    server.listen(1)

    print(f"Waiting for sender on TCP port {LISTEN_PORT}...")

    try:
        connection, address = server.accept()

        with connection:
            connection.setsockopt(
                socket.IPPROTO_TCP,
                socket.TCP_NODELAY,
                1,
            )

            print(f"Sender connected from {address[0]}:{address[1]}")
            print("Press Q in the video window to stop.")

            while True:
                header = receive_exact(connection, HEADER_SIZE)

                if header is None:
                    print("Sender disconnected.")
                    break

                frame_size = struct.unpack("!I", header)[0]

                if frame_size <= 0 or frame_size > MAX_FRAME_SIZE:
                    print(f"Invalid frame size: {frame_size}")
                    break

                frame_data = receive_exact(connection, frame_size)

                if frame_data is None:
                    print("Connection closed during frame reception.")
                    break

                compressed_array = np.frombuffer(
                    frame_data,
                    dtype=np.uint8,
                )

                frame = cv2.imdecode(
                    compressed_array,
                    cv2.IMREAD_COLOR,
                )

                if frame is None:
                    print("Failed to decode frame.")
                    continue

                cv2.imshow("Remote Screen - Press Q to exit", frame)

                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break

    finally:
        server.close()
        cv2.destroyAllWindows()
        print("Receiver stopped.")


if __name__ == "__main__":
    main()