#!/usr/bin/python3

import argparse
import os
os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
import pygame

def bytes_from_file(filename, chunksize=8192):
    with open(filename, "rb") as f:
        while True:
            chunk = f.read(chunksize)
            if chunk:
                for b in chunk:
                    yield b
            else:
                break

def work(raw, width, height, png):
    pygame.init()

    screen = pygame.display.set_mode((width, height))
    screen.fill((0, 0, 0))

    pos = 0
    for byte in bytes_from_file(raw):
        x = pos % width
        y = int(pos / height)
        pos += 1

        if y % 2 == 0:
            if x % 2 == 0:
                color = (byte, 0, 0)
            else:
                color = (0, byte, 0)
        else:
            if x % 2 == 0:
                color = (0, byte, 0)
            else:
                color = (0, 0, byte)

        screen.set_at((x, y), color)

    pygame.display.update()

    if (png):
        pygame.image.save(screen, png)

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

def main():
    parser = argparse.ArgumentParser(description='View SRGGB8 and optionally save to PNG')
    parser.add_argument('-w', '--write', dest='output', help='Write png file')
    parser.add_argument('input')
    args = parser.parse_args()

    work(args.input, 640, 480, args.output)

if __name__ == '__main__':
    main()