import signal
import socket
import struct
import time

import cv2
import numpy as np
from mss import mss


RECEIVER_IP = "172.16.11.81"   # Your friend's computer IP
RECEIVER_PORT = 5000

FPS = 15
JPEG_QUALITY = 55
MONITOR_NUMBER = 1

running = True


def stop_stream(signum=None, frame=None):
    global running
    running = False
    print("\nStopping screen stream safely...")

 
signal.signal(signal.SIGINT, stop_stream)
signal.signal(signal.SIGTERM, stop_stream)


def send_all(sock, data):
    """Send the complete byte sequence."""
    view = memoryview(data)

    while view and running:
        bytes_sent = sock.send(view)

        if bytes_sent == 0:
            raise ConnectionError("TCP connection closed")

        view = view[bytes_sent:]


def main():
    print(f"Connecting to {RECEIVER_IP}:{RECEIVER_PORT}...")

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        sock.connect((RECEIVER_IP, RECEIVER_PORT))

        print("Connected.")
        print("Press Ctrl+C once to stop.")

        with mss() as sct:
            monitor = sct.monitors[MONITOR_NUMBER]

            frame_interval = 1.0 / FPS
            next_frame_time = time.perf_counter()

            while running:
                screenshot = sct.grab(monitor)

                frame = np.asarray(screenshot)
                frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)

                # Reduce resolution to lower bandwidth.
                frame = cv2.resize(
                    frame,
                    None,
                    fx=0.6,
                    fy=0.6,
                    interpolation=cv2.INTER_AREA,
                )

                success, encoded_frame = cv2.imencode(
                    ".jpg",
                    frame,
                    [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY],
                )

                if not success:
                    continue

                frame_data = encoded_frame.tobytes()
                frame_size = len(frame_data)

                # Network byte order, unsigned 4-byte integer.
                header = struct.pack("!I", frame_size)

                send_all(sock, header)
                send_all(sock, frame_data)

                print(
                    f"\rSent: {frame_size / 1024:.1f} KB",
                    end="",
                    flush=True,
                )

                next_frame_time += frame_interval
                sleep_time = next_frame_time - time.perf_counter()

                if sleep_time > 0:
                    time.sleep(sleep_time)
                else:
                    next_frame_time = time.perf_counter()

    print("\nScreen streaming stopped.")


if __name__ == "__main__":
    try:
        main()
    except ConnectionRefusedError:
        print("Connection refused. Start the receiver first.")
    except (ConnectionError, BrokenPipeError) as error:
        print(f"\nConnection lost: {error}")