#!/usr/bin/python3

import glob
import os
import shutil
import subprocess
import sys
import tempfile
import time


def join(item0, *items):
    return os.path.join(item0, *[p.lstrip('/') for p in items])


def reexec_root():
    # run as root
    if os.geteuid() != 0:
        os.execvp("sudo", ["sudo"] + sys.argv)


def run(*l, **d):
    print('run:', *l, d)
    return subprocess.run(*l, **d)


def out_of_date(path, timeout):
    'true if path mtime is greater than timeout in seconds'
    return (time.time() - os.path.getmtime(path)) > timeout


def get_fresh_firmware(firmware_dir):
    firmware_timeout = 24*3600 # 1d timeout
    if (not os.path.exists(firmware_dir)
        or out_of_date(firmware_dir, firmware_timeout)):
        shutil.rmtree(firmware_dir, ignore_errors=True)
        run(['git', 'clone', '--depth=1',
             'https://github.com/raspberrypi/firmware',
             firmware_dir], check=True)


def copy_to_boot(firmware_dir, boot_mountpoint):
    run(['rsync', '-aHX', join(firmware_dir, 'boot/.'),
         join(boot_mountpoint, '.')], check=True)
    shutil.copy('files/cmdline.txt',        boot_mountpoint)
    shutil.copy('files/config.txt',         boot_mountpoint)
    shutil.copy('debootstrapdir/vmlinuz',   join(boot_mountpoint, 'vmlinuz'),
                follow_symlinks=True)
    shutil.copy('debootstrapdir/initrd.img',join(boot_mountpoint, 'initrd.img'),
                follow_symlinks=True)
    content = run(['ls', '-lah',
                   'debootstrapdir/vmlinuz', 'debootstrapdir/initrd.img'],
                  check=True, stdout=subprocess.PIPE).stdout
    with open(join(boot_mountpoint, 'versions.txt'), 'wb') as f:
        f.write(content)


def copy_to_partitions(firmware_dir, boot_mountpoint, root_mountpoint):
    copy_to_boot(firmware_dir, boot_mountpoint)
    run(['rsync', '-aHX', 'debootstrapdir/.', root_mountpoint + '/.'],
        check=True)


def main():
    if len(sys.argv) < 2:
        print('Usage: {} <device>'.format(sys.argv[0]), file=sys.stderr)
        raise SystemExit(1)
    device = sys.argv[1]
    reexec_root()
    firmware_dir = 'rpi-firmware'
    boot_mountpoint = tempfile.TemporaryDirectory()
    root_mountpoint = tempfile.TemporaryDirectory()
    try:
        run(['mount', device + '1', boot_mountpoint.name], check=True)
        run(['mount', device + '2', root_mountpoint.name], check=True)
        get_fresh_firmware(firmware_dir)
        copy_to_partitions(firmware_dir,
                           boot_mountpoint.name,
                           root_mountpoint.name)
    finally:
        run(['umount', boot_mountpoint.name])
        run(['umount', root_mountpoint.name])


if __name__ == '__main__':
    main()