#!/usr/bin/env python2
#-*- coding: utf-8 -*-
#
#

import pygame
import pygame.camera

import os, sys, socket
import struct, base64, StringIO

__author__ = '@laszlokuehl'

HELP_TEXT = """\n
\tUsage:
\t------
\t* Server:
\t\t$ python2 camera.py -s 0.0.0.0 8000 800x600

\t* Client:
\t\t$ python2 camera.py -c 127.0.0.1 8000\n\n"""

class __Base:
    def send(self, conn, data):
        data = base64.b64encode(data)
        data = struct.pack('>I', len(data)) + data

        conn.sendall(data)

    def __recv(self, conn, buff):
        data = ''

        while len(data) < buff:
            packet = conn.recv(buff - len(data))

            if not packet:
                break

            data += packet

        return data

    def recv(self, conn):
        data_len = self.__recv(conn, 4)

        if not data_len:
            return None

        data_len = struct.unpack('>I', data_len)[0]
        return base64.b64decode(self.__recv(conn, data_len))

class Server(__Base):
    def __init__(self, address=('0.0.0.0', 8000), geometry=(800, 600)):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind(address)
        self.socket.listen(1)

        self.geometry = geometry

    def handle_cli(self, conn):
        self.surface = pygame.display.set_mode(self.geometry, pygame.RESIZABLE)
        pygame.display.set_caption('Raspberry Pi Web Cam')

        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    conn.close()
                    return

                if event.type == pygame.VIDEORESIZE:
                    self.geometry = (event.w, event.h)
                    self.surface = pygame.display.set_mode((event.w, event.h),
                                                           pygame.RESIZABLE)

            img = pygame.image.load(StringIO.StringIO(self.recv(conn)))
            img = pygame.transform.scale(img, self.geometry)

            self.surface.blit(img, (0, 0))
            pygame.display.update()

    def start(self):
        while True:
            conn, addr = self.socket.accept()
            self.handle_cli(conn)

class Client(__Base):
    def __init__(self, address):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect(address)

        pygame.camera.init()
        self.cam = pygame.camera.Camera("/dev/video0", (640, 480))
        self.cam.start()

    def shot(self):
        pygame.image.save(self.cam.get_image(), 'image.jpg')
        return open('image.jpg', 'rb').read()

    def start(self):
        while True:
            self.send(self.socket, self.shot())

if __name__ == '__main__':
    if len(sys.argv[1:]) == 4 and sys.argv[1] in ['-s', '-server', '--server']:
        pygame.init()

        server = Server((sys.argv[2], int(sys.argv[3])), map(int, sys.argv[4].split('x')))
        server.start()

    elif len(sys.argv[1:]) == 3:
        client = Client((sys.argv[2], int(sys.argv[3])))
        client.start()
    else:
        print HELP_TEXT