#!python
#
# Display local desktop on remove host over LAN without VGA/HDMI cable.
# Copyright (c) 2019, Hiroyuki Ohsaki.
# All rights reserved.
#

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import os
import socket
import sys
import time
import zlib

from Xlib import X, display
from perlcompat import die, warn, getopts
import pygame
import pygame.gfxdraw
import pygame.transform

POINTER_COLOR = 255, 255, 0, 128
POINTER_SIZE = 10
FONT_NAME = '/usr/share/fonts/truetype/freefont/FreeSans.ttf'
FONT_SIZE = 20
LABEL_COLOR = 255, 255, 255, 128
BLANK_COLOR = 0, 0, 255, 255

def usage():
    die("""\
usage: {} [-vr] [-s host] [-p port] [-W width] [-H height]
  -v       verbose mode
  -r       receiver mode
  -s host  specify receiver hostname/address
  -p port  port number
  -W #     screen width (default: 800)
  -H #     screen height (default: 600)
""".format(sys.argv[0]))

def pointer_geometry(xscreen):
    """Obtain the current geometry of the pointer (i.e., mouse) on screen
    XSCREEN.  This code is contributed by Hal."""
    stat = xscreen.root.query_pointer()
    x, y = stat.root_x, stat.root_y
    return x, y

class Frame:
    def __init__(self, width=800, height=600, user=None, hostname=None):
        self.width = width
        self.height = height
        self.user = user
        self.hostname = hostname
        self.pnt_x = None
        self.pnt_y = None
        self.header = None
        self.img = None
        self.zbuf = None

    def compose_header(self):
        """Compose a frame header, which is composed of field values separated
        by spaces and terminated by the newline."""
        astr = '{} {} {} {} {} {} {}\n'.format('P0', self.user, self.hostname,
                                               self.width, self.height,
                                               self.pnt_x, self.pnt_y)
        return bytes(astr, encoding='utf-8')

    def dump_root_window(self):
        """Return the rectangle of the root window in RGBX format."""
        xscreen = display.Display().screen()
        img = xscreen.root.get_image(0, 0, self.width, self.height, X.ZPixmap,
                                     0xffffffff)
        # format conversion from BGRX to RGBX
        buf = bytearray(img.data)
        for i in range(0, self.width * self.height * 4, 4):
            buf[i], buf[i + 2] = buf[i + 2], buf[i]  # swap red and blue
        return buf

    def capture(self):
        """Capture the current desktop and store the header and the image in
        HEADER and IMG attributes, respectively.  Zlib-compressed data is also
        stored in ZBUF attribute."""
        self.header = self.compose_header()
        self.img = self.dump_root_window()
        self.zbuf = zlib.compress(self.header + self.img)

    def load(self, zbuf):
        """Decode zlib-compressed data ZBUF and update object attributes
        (PROTO, USER, HOSTNAME, WIDTH, HEIGHT, PNT_X, PNT_Y, HEADER, and
        IMG)."""
        try:
            buf = zlib.decompress(zbuf)
        except zlib.error:
            warn('zlib decompression failed.  skipping...')
            return None
        try:
            offset = buf.index(b'\n') + 1
        except ValueError:
            warn('invalid header.  skipping...')
            return None
        self.header = buf[:offset - 1]
        self.img = buf[offset:]
        self.zbuf = zbuf

        # parse frame header
        proto, user, hostname, width, height, pnt_x, pnt_y = self.header.split(
            b' ')
        self.proto = proto
        self.user = user
        self.hostname = hostname
        self.width = int(width)
        self.height = int(height)
        self.pnt_x = int(pnt_x)
        self.pnt_y = int(pnt_y)

class Sender:
    def __init__(self, server='localhost', port=5000, width=800, height=600):
        self.server = server
        self.port = port
        self.width = width
        self.height = height

    def mainloop(self):
        """Repeatedly send the rectangle of the current desktop to remote host
        using UDP."""
        user = os.getenv('USER')
        hostname = socket.gethostname()
        xscreen = display.Display().screen()
        self.sk = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
        while True:
            frame = Frame(self.width, self.height, user, hostname)
            frame.pnt_x, frame.pnt_y = pointer_geometry(xscreen)
            frame.capture()
            warn('sending frame ({:,} bytes) to {}:{}...'.format(
                len(frame.zbuf), self.server, self.port))
            addr = self.server, self.port
            self.sk.sendto(frame.zbuf, addr)

class Receiver:
    def __init__(self, port=5000, width=800, height=600):
        self.port = port
        self.width = width
        self.height = height
        self.ndrawn = 0

    def init_display(self):
        """Initialize display settings and create a window for display."""
        warn('initializing display...')
        # disable screen saver and display power management
        # FIXME: restore power management settings
        os.system('xset s off')
        os.system('xset -dpms')

        # initialize pygame and create a window
        pygame.display.init()
        pygame.font.init()
        self.screen = pygame.display.set_mode((self.width, self.height))
        self.font = pygame.font.Font(FONT_NAME, FONT_SIZE)

    def recv(self):
        """Wait incoming frame data and decode the frame.  Return the frame
        decoded"""
        self.sk.settimeout(3.)
        try:
            warn('waiting frame...')
            zbuf, addr = self.sk.recvfrom(10_000_000)
        except socket.timeout:
            return None
        host, port = addr
        warn('frame received ({:,} bytes) from {}:{}...'.format(
            len(zbuf), host, port))
        frame = Frame()
        frame.load(zbuf)
        return frame

    def draw_img(self, screen, img, width, height):
        """Draw an image IMG with width WIDTH and height HEIGHT at the origin
        (0, 0) on screen SCREEN."""
        try:
            # convert RGBX image to pygame.Surface object
            img = pygame.image.fromstring(img, (width, height), 'RGBX')
        except ValueError:
            return
        screen.blit(img, (0, 0))

    def draw_pointer(self, screen, x, y):
        """Draw a virtual pointer at (X, Y) on screen SCREEN."""
        pygame.gfxdraw.filled_ellipse(screen, x, y, POINTER_SIZE, POINTER_SIZE,
                                      POINTER_COLOR)

    def draw_label(self, screen, font, label):
        """Draw a text LABEL using font FONT around the upper-right corner of
        the screen SCREEN."""
        text = font.render(label, 1, LABEL_COLOR)
        x = screen.get_width() - text.get_width() - FONT_SIZE // 2
        y = FONT_SIZE // 2
        screen.blit(text, (x, y))

    def draw(self, frame):
        """Draw frame FRAME at the origin (0, 0) of the current window."""
        self.screen.fill((0, 0, 0))
        self.draw_img(self.screen, frame.img, frame.width, frame.height)
        # blink the pointer as keep-alive indicator
        if self.ndrawn % 2 == 0:
            self.draw_pointer(self.screen, frame.pnt_x, frame.pnt_y)
        label = frame.user + b'@' + frame.hostname
        self.draw_label(self.screen, self.font, label)
        pygame.display.update()
        self.ndrawn += 1

    def mainloop(self):
        """Repeatedly receive frames from a client, and display the frame in
        the window."""
        self.init_display()
        self.sk = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sk.bind(('', self.port))  # all available interfaces
        while True:
            frame = self.recv()
            if frame and frame.img:
                self.draw(frame)
            else:
                self.screen.fill(BLANK_COLOR)
                pygame.display.update()

def main():
    opt = getopts('vrs:p:w:h:') or usage()
    verbose = opt.v
    receiver_mode = opt.r
    server = opt.s if opt.s else 'localhost'
    port = int(opt.p) if opt.p else 5000
    width = int(opt.w) if opt.w else 800
    height = int(opt.h) if opt.h else 600
    if receiver_mode:
        receiver = Receiver(port, width, height)
        receiver.mainloop()
    else:
        sender = Sender(server, port, width, height)
        sender.mainloop()

if __name__ == "__main__":
    main()
