#!/usr/bin/env python3

import io, os, sys, time, pygame

try:
    import picamera

    class Camera:
        def __init__(self, rect):
            self.rect = rect
            self.camera = picamera.PiCamera()
            self.camera.led = False
            self.camera.resolution = (self.rect.width, self.rect.height)
            self.rgb = bytearray(self.rect.width * self.rect.height * 4)

        def close(self):
            self.camera.close()

        def take(self, name):
            name = name + '.jpg'
            self.camera.resolution = (3280, 2464)
            self.camera.capture(name)
            self.camera.resolution = (self.rect.width, self.rect.height)
            return name

        def stream(self):
            stream = io.BytesIO()
            self.camera.capture(stream, use_video_port=True, format='rgba', resize=(self.rect.width, self.rect.height))
            stream.seek(0)
            stream.readinto(self.rgb)
            stream.close()
            img = pygame.image.frombuffer(self.rgb, (self.rect.width, self.rect.height), 'RGBA')

            return pygame.transform.flip(img, True, False)

        def draw_on(self, screen):
            screen.blit(self.stream(), self.rect)


except:
    class Camera:
        def __init__(self, rect):
            self.rect = rect
            self.img = pygame.transform.scale(pygame.image.load('tournesol.jpg'), (self.rect.width, self.rect.height))

        def close(self):
            pass

        def take(self, name):
            print('Take picture! ' + name)
            return 'tournesol.jpg'

        def draw_on(self, screen):
            screen.blit(self.img, self.rect)



class Button:
    def __init__(self, rect):
        self.rect = rect
        self.pressed = False
        self.enabled = True

    def draw_on(self, screen):
        pass

    def is_pressed(self, pos):
        return self.rect.collidepoint(pos)

    def press(self):
        self.pressed = True

    def unpress(self):
        self.pressed = False
        self.run()

    def run(self):
        pass

class InfoButton(Button):
    def __init__(self, rect, photoboite):
        self.count = 0
        self.timestamp = 0
        self.photoboite = photoboite

        super(InfoButton, self).__init__(rect)

        self.enabled = False

        self.font = pygame.font.Font(None, 25)
        self.logo = pygame.image.load(os.path.dirname(os.path.realpath(__file__)) + '/raspberrypi.png').convert_alpha()

    def draw_on(self, screen):
        if self.enabled:
            screen.blit(self.logo, self.logo.get_rect())

            top = self.logo.get_rect().height + 10

            text = "COUNT: " + str(self.photoboite.capture_button.count)
            size = self.font.size(text)
            ren = self.font.render(text, True, (255, 255, 255))
            screen.blit(ren, (0, top))

            text = "FPS: " + str(self.photoboite.clock.get_fps())
            size = self.font.size(text)
            ren = self.font.render(text, True, (255, 255, 255))
            screen.blit(ren, (0, top + 20))

    def run(self):
        if time.time() - self.timestamp > 0.5:
            self.count = 1
        else:
            self.count += 1

        self.timestamp = time.time()

        if self.count >= 2:
            self.enabled = not self.enabled

class PowerButton(Button):
    def __init__(self, rect):
        self.count = 0
        self.timestamp = 0

        super(PowerButton, self).__init__(rect)

    def run(self):
        if time.time() - self.timestamp > 0.5:
            self.count = 1
            self.timestamp = time.time()
        else:
            self.count += 1
            self.timestamp = time.time()

        if self.count >= 5:
            print('POWEROFF')

class CaptureButton(Button):
    countdown = 3

    def __init__(self, size, topleft, photoboite):
        self.up_image = pygame.image.load(os.path.dirname(os.path.realpath(__file__)) + '/button-up.png').convert_alpha()
        self.up_image = pygame.transform.smoothscale(self.up_image, (size, size))

        self.down_image = pygame.image.load(os.path.dirname(os.path.realpath(__file__)) + '/button-down.png').convert_alpha()
        self.down_image = pygame.transform.smoothscale(self.down_image, (size, size))

        rect = self.up_image.get_rect()
        rect.topleft = topleft

        super(CaptureButton, self).__init__(rect)
        self.photoboite = photoboite
        self.elapsed = 0
        self.capture_mode = False
        self.count = 0

        self.background_rect = pygame.Rect(0, self.photoboite.screen.get_height() - 80, self.photoboite.screen.get_width(), 80)
        self.background = pygame.Surface((self.background_rect.width, self.background_rect.height))
        self.background.fill((0, 0, 0))
        self.background.set_alpha(150)

        self.font = pygame.font.Font(None, 80)

        self.photo_film = pygame.image.load(os.path.dirname(os.path.realpath(__file__)) + '/film.png').convert_alpha()

    def draw_on(self, screen):
        if self.enabled:
            if self.pressed:
                screen.blit(self.down_image, self.rect)
            else:
                screen.blit(self.up_image, self.rect)
        if self.capture_mode and self.elapsed < CaptureButton.countdown:
            screen.blit(self.background, self.background_rect.topleft)

            text = str(CaptureButton.countdown - self.elapsed)
            size = self.font.size(text)
            ren = self.font.render(text, True, (255, 255, 255))
            ren.set_alpha(150)

            topleft = (self.rect.topleft[0] + self.rect.width / 2 - size[0] / 2, self.rect.topleft[1] + self.rect.height / 2 - size[1] / 2)
            top = self.background_rect.top + self.background_rect.height / 2 - size[1] / 2
            left = self.background_rect.left + self.background_rect.width / 2 - size[0] / 2
            screen.blit(ren, (left, top))

    def run(self):
        pygame.event.set_blocked((pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP))
        self.capture_mode = True
        self.enabled = False

        capture = Capture(self.count)
        self.count += 1

        for i in range(3):
            wait = True
            start = time.time()
            while wait:
                self.elapsed = int(time.time() - start)
                if self.elapsed >= CaptureButton.countdown:
                    wait = False

                self.photoboite.event()
                self.photoboite.draw()
                pygame.display.flip()

            capture.take(self.photoboite.camera)
            pygame.time.wait(600)

        for photo in capture.photos:
            self.photoboite.screen.blit(photo.capture, (81, 0))
            self.photoboite.screen.blit(self.photo_film, (0, 0))
            pygame.display.flip()
            pygame.time.wait(800)
            self.photoboite.event()

        pygame.time.wait(200)

        pygame.event.set_allowed((pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP))
        self.capture_mode = False
        self.enabled = True
        self.photoboite.sleep.move()


class Capture:
    def __init__(self, uid):
        self.uid = uid
        self.photos = []

    def take(self, camera):
        self.photos.append(Photo(self.uid, len(self.photos)).take(camera))


class Photo:
    def __init__(self, cid, uid):
        self.cid = cid
        self.uid = uid

    def take(self, camera):
        self.name = "photos/%s-%04d-%04d" % (time.strftime('%Y%m%d-%H%M%S'), self.cid, self.uid)
        self.capture = pygame.transform.smoothscale(pygame.image.load(camera.take(self.name)), (638, 480))

        return self


class SleepMode:
    def __init__(self, seconds, screen):
        self.enabled = True
        self.drawn = False
        self.seconds = seconds

        self.screen = screen

        self.clicker = pygame.image.load(os.path.dirname(os.path.realpath(__file__)) + '/clicker.png').convert_alpha()
        self.clicker_rect = self.clicker.get_rect()
        self.clicker_rect.topleft = ((self.screen.get_width() - self.clicker_rect.width) / 2, (self.screen.get_height() - self.clicker_rect.height) / 2)

    def move(self):
        self.last_action = time.time()

    def event(self):
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONUP or event.type == pygame.KEYUP:
                self.move()
                self.enabled = False

    def draw(self):
        if not self.drawn:
            self.screen.fill((255,255,255))
            self.screen.blit(self.clicker, self.clicker_rect)
            pygame.display.flip()
            self.drawn = True

    def sleep(self):
        if time.time() - self.last_action > self.seconds:
            self.enabled = True
            self.drawn = False

class Photoboite:
    def __init__(self):
        self.screen = self.screen()
        pygame.font.init()

        pygame.mouse.set_visible(False)

        self.buttons = []

        button_size = int(self.screen.get_height() / 4)
        topleft = (self.screen.get_width() - button_size) / 2, self.screen.get_height() - button_size - 10

        self.capture_button = CaptureButton(button_size, topleft, self)
        self.buttons.append(self.capture_button)
        self.buttons.append(PowerButton(pygame.Rect(self.screen.get_width() - 50, 0, 50, 50)))
        self.buttons.append(InfoButton(pygame.Rect(0, 0, 50, 50), self))

        camera_rect = pygame.Rect(0, 0, 800, 480)
        self.camera = Camera(camera_rect)

        self.clock = pygame.time.Clock()

        self.sleep = SleepMode(30, self.screen)


    def screen(self):
        if os.getenv('DISPLAY'):
            pygame.display.init()
            display = pygame.display.Info()
            return pygame.display.set_mode((800, 480))
        else:
            drivers = ['directfb', 'fbcon', 'svgalib']
            found = False
            for driver in drivers:
                if not os.getenv('SDL_VIDEODRIVER'):
                    os.putenv('SDL_VIDEODRIVER', driver)
                try:
                    pygame.display.init()
                except pygame.error as e:
                    print('Driver: {0} failed.'.format(driver))
                    continue
                found = True
                break

            if not found:
               raise Exception('No suitable video driver found!')

            return pygame.display.set_mode()

    def draw(self):
        self.clock.tick(60)

        self.camera.draw_on(self.screen)

        for button in self.buttons:
            button.draw_on(self.screen)

    def event(self):
        for event in pygame.event.get():
            self.sleep.move()

            if event.type == pygame.QUIT or event.type == pygame.KEYDOWN and event.unicode == 'q':
                self.quit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                for button in self.buttons:
                    if button.is_pressed(event.pos):
                        button.press()
            elif event.type == pygame.MOUSEBUTTONUP:
                for button in self.buttons:
                    if button.is_pressed(event.pos):
                        button.unpress()

    def run(self):
        pressed = False
        while True:
            if self.sleep.enabled:
                self.sleep.event()
                self.sleep.draw()
                self.clock.tick(60)
            else:
                self.event()
                self.draw()
                self.sleep.sleep()
                pygame.display.flip()

    def quit(self):
        self.camera.close()
        sys.exit()

Photoboite().run()