#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p grub2 rsync utillinux shadow utillinux e2fsprogs
from subprocess import run
from pathlib import Path
from tempfile import TemporaryDirectory
import argparse
import json
import os
import sys
import time

root_device = Path('/dev/disk/by-id/ata-Crucial_CT250MX200SSD1_1537108FC44F')
bios_boot_part_type = '21686148-6449-6E6F-744E-656564454649'
config_dir = Path(__file__).parent.parent.absolute()

if os.getlogin() != 'root':
    print("ERROR: must be run as root", file=sys.stderr)
    sys.exit(1)

if not root_device.exists():
    print(f"ERROR: {root_device} not found", file=sys.stderr)
    sys.exit(1)

print(f"WARNING: this script will WIPE all data on {root_device}")
if input('Write "Yes" to continue:') != 'Yes':
    sys.exit(1)

with TemporaryDirectory() as tmp_path:
    tmp = Path(tmp_path)
    print(f"Created temporary directory {tmp}")

    parts = (
        'label: gpt\n'
        f'name=grub start=2MiB size=10MiB type={bios_boot_part_type}\n'
        'name=root size=100GiB\n'
    )
    run(['sfdisk', root_device], input=parts.encode())

    parts_info = json.loads(run(['sfdisk', '--json', root_device], capture_output=True, check=True).stdout.decode())
    root_part = Path(parts_info["partitiontable"]["partitions"][1]['node']).resolve()

    for i in range(40):
        if root_part.exists():
            break
        time.sleep(0.2)
    else:
        print(f"ERROR: create partition not exists: {root_part}", file=sys.stderr)
        sys.exit(1)

    run(['mkfs.ext4', root_part])

    root = tmp.joinpath('root')
    root.mkdir()

    try:
        run(['mount', root_part, root], check=True)

        run(['mkdir', '-p', root.joinpath('etc', 'nixos')], check=True)
        run(['rsync', '-r', '--progress', f'{config_dir!s}/', root.joinpath('etc', 'nixos')], check=True)

        root_uuid = parts_info["partitiontable"]["partitions"][1]['uuid'].lower()
        root.joinpath('etc', 'nixos', 'hw.json').write_text(json.dumps({
            "rootUUID": f'{root_uuid}',
        }))

        run(['nixos-install', '--no-root-passwd', '--root', root], check=True)
        run(['grub-install', f'--root-directory={root!s}', f'--boot-directory={root.joinpath("boot")!s}', root_device], check=False)
        run(['chpasswd', '--root', root], input=b'root:toor')

    finally:
        run(['umount', root])