#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Author: KaratekHD <karatekhd@opensuse.org>
#
# Copyright (c) 2021 KaratekHD.
# is 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
# (at your option) 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 <http://www.gnu.org/licenses/>.
#

import argparse
import os
import sys
import tarfile

from termcolor import colored

parser = argparse.ArgumentParser(
    description="Deploy openSUSE Tumbleweed to the Mars MA3.")
parser.add_argument("-r", "--rootfs", help="Path to rottfs.tar", required=True)
parser.add_argument(
    "-d",
    "--device",
    help="Deviceprefix to deploy to (e.g. /dev/mmcblk0p)",
    required=True)
parser.add_argument(
    "-i",
    "--input_dir",
    help="Build files generated by build.sh",
    required=True)
args = parser.parse_args()
rootfs_path = args.rootfs
block = args.device


# Check if script is run as root
if os.geteuid() != 0:
    print(colored("Script must be run as root, relaunching...", "cyan"))
    # Relaunch script as root
    # Get all args to one string
    os.system("sudo python3 " + " ".join(sys.argv))
    # Exit script
    sys.exit(1)

# Check if args.input-dir is a directory
if not os.path.isdir(args.input_dir):
    print(colored("Input directory is invalid!", "red"))
    sys.exit(1)
os.chdir(args.input_dir)
# Check if block device exists
if not os.path.exists(block):
    print(colored("Block device does not exist", "red"))
    sys.exit(1)

# Check if something is mounted on /mnt abort if so
if os.path.ismount("/mnt"):
    print(colored("Please unmount /mnt before running this script", "red"))
    sys.exit(1)

# Check if /mnt/boot exists and create if not
if not os.path.exists("/mnt/boot"):
    os.makedirs("/mnt/boot")
# Check if /mnt/rootfs exists and create if not
if not os.path.exists("/mnt/rootfs"):
    os.makedirs("/mnt/rootfs")

# Umount all partitions
# Check if /mnt/rootfs is mounted
if os.path.ismount("/mnt/rootfs"):
    print(colored("Unmounting /mnt/rootfs", "yellow"))
    os.system("umount -l /mnt/rootfs")
# Check if /mnt/boot is mounted
if os.path.ismount("/mnt/boot"):
    print(colored("Unmounting /mnt/boot"))
    os.system("umount -l /mnt/boot")

# TODO: Partitioning of the SD card

# Format the SD card
print(colored("Formatting partition 1 with fat32...", "yellow"))
os.system("mkfs.vfat -F 32 -n BOOT {}1".format(block))
print(colored("Formatted partition 1", "green"))
print(colored("Formatting partition 3 with ext2...", "yellow"))
os.system("mkfs.ext2 -L rootfs {}3".format(block))
print(colored("Formatted partition 3", "green"))

# Flash preloader
print(colored("Flashing preloader...", "yellow"))
os.system("dd if=preloader-mkpimage.bin of={}2 status=progress".format(block))
print(colored("Preloader flashed!", "green"))
print(colored("Syncing...", "yellow"))
os.system("sync")
print(colored("Synced!", "green"))

# Check if device is mounted
if not os.path.ismount("/mnt/rootfs"):
    print(colored("Mounting /mnt/rootfs", "yellow"))
    os.system("sudo mount {}3 /mnt/rootfs".format(block))
if not os.path.ismount("/mnt/boot"):
    print(colored("Mounting /mnt/boot", "yellow"))
    os.system("sudo mount {}1 /mnt/boot".format(block))

print(colored("Copying files to /mnt/boot", "yellow"))

# List of files to copy
files = ["uImage", "devicetree.dtb", "fpga.rbf", "u-boot.img", "uboot.scr"]

# Copy files to /mnt/boot with progress bar
for file in files:
    print(colored("- " + file, "green"))
    with open(file, "rb") as f:
        with open("/mnt/boot/" + file, "wb") as f2:
            while True:
                data = f.read(1024)
                if not data:
                    break
                f2.write(data)

# Remove everything from /mnt/rootfs
print(colored("Removing everything from /mnt/rootfs...", "yellow"))
os.system("rm -rf /mnt/rootfs/*")

# Extract rootfs from tarfile to /mnt/rootfs
# Check wether tarfile is .gz or .xz
compression = ""
if rootfs_path.endswith(".gz") or rootfs_path.endswith(".xz"):
    if rootfs_path.endswith(".gz"):
        compression = "gz"
    elif rootfs_path.endswith(".xz"):
        compression = "xz"
    print(
        colored(
            "Extracting rootfs from {} to /mnt/rootfs...".format(rootfs_path),
            "yellow"))
    with tarfile.open(rootfs_path, "r:" + compression) as tar:
        tar.extractall("/mnt/rootfs")
    print(colored("Extracted {} to /mnt/rootfs".format(rootfs_path), "green"))
else:
    # Extract .tar
    print(colored("Extracting {} to /mnt/rootfs...".format(rootfs_path), "yellow"))
    with tarfile.open(rootfs_path) as tar:
        tar.extractall("/mnt/rootfs")
    print(colored("Extracted {} to /mnt/rootfs".format(rootfs_path), "green"))
print(colored("Syncing...", "yellow"))
os.system("sync")
print(colored("Umounting devices...", "yellow"))
os.system("umount /mnt/boot")
os.system("umount /mnt/rootfs")

print(colored("Done", "green"))