import asyncio
import logging
import re
import sys
import time
from pathlib import Path
from functools import wraps
from runner import get_output, iscsi_login, mount, run, sshfs

logging.basicConfig(level=logging.INFO, format='%(asctime)s.%(msecs)03d %(levelname)s %(filename) 12s:%(lineno)d %(message)s', datefmt='%H:%M:%S')
log = logging.getLogger()

WORK = Path('/tmp/pi-setup')
LITE_PREFIX = '2024-03-15-raspios-bookworm'

# These come from fdisk -l on the img:
IMG_BOOT_OFFSET = 512 * 8192
IMG_ROOT_OFFSET = 512 * 1056768

TFTP_SPEC = 'root@pipe:/opt/dnsmasq/tftp/'


def step(func):
    name = func.__name__

    @wraps(func)
    async def wrapper(*a, **kw):
        print("", file=sys.stderr)
        log.info(f'👣 step {name}')
        t1 = time.time()
        ret = await func(*a, **kw)
        dt = time.time() - t1
        log.info(f' -> step {name} took {dt:3} seconds')
        return ret

    return wrapper


@step
async def init_work_dir():
    WORK.mkdir(exist_ok=True)
    await run(
        'wget', '--continue', '--no-verbose', '-O', WORK / 'raspios.img.xz',
        f'https://downloads.raspberrypi.org/raspios_lite_arm64/images/raspios_lite_arm64-2024-03-15/{LITE_PREFIX}-arm64-lite.img.xz'
    )


@step
async def unpack_fresh_img_copy():
    await run('xz', '-dkf', WORK / 'raspios.img.xz')
    await run('qemu-img', 'resize', ['-f', 'raw'], WORK / 'raspios.img', '4G')


@step
async def extract_boot_files():
    async with mount(WORK, WORK / 'raspios.img', IMG_BOOT_OFFSET) as img_boot:
        await run('cp', img_boot / 'bcm2710-rpi-3-b-plus.dtb', WORK)
        await run('cp', img_boot / 'kernel8.img', WORK)


@step
async def setup_ssh_login():
    async with mount(WORK, WORK / 'raspios.img', IMG_ROOT_OFFSET) as img_root:
        root_pub_key = Path('/root/.ssh/id_ecdsa.pub').read_text()
        (img_root / 'root/.ssh/authorized_keys').write_text(root_pub_key)
        # (img_root / 'etc/iscsi/initiatorname.iscsi').write_text(f"InitiatorName=iqn.2024-03.com.bigasterisk:{PI_HOSTNAME}\n")

    async with mount(WORK, WORK / 'raspios.img', IMG_BOOT_OFFSET) as img_boot:
        await run('touch', img_boot / 'ssh')


@step
async def qemu():
    await run(
        'qemu-system-aarch64',
        ['-machine', 'raspi3b'],
        ['-cpu', 'cortex-a72'],
        ['-nographic'],
        ['-dtb', WORK / 'bcm2710-rpi-3-b-plus.dtb'],
        ['-m', '1G'],
        ['-smp', '4'],
        ['-kernel', WORK / 'kernel8.img'],
        ['-sd', WORK / 'raspios.img'],
        ['-append', "rw earlyprintk loglevel=8 console=ttyAMA0,115200 dwc_otg.lpm_enable=0 root=/dev/mmcblk0p2 rootdelay=1"],
        ['-device', 'usb-net,netdev=net0'],
        ['-netdev', 'user,id=net0,hostfwd=tcp::2222-:22'],
    )


@step
async def setup_pi_in_emulator():
    src_path = Path(__file__).parent / 'on_pi_setup.sh'
    path_on_pi = '/tmp/on_pi_setup.sh'

    ssh_opts = [
        '-oPort=2222',
        '-oConnectTimeout=2',
        '-oBatchMode=yes',
        '-oNoHostAuthenticationForLocalHost=yes',
    ]
    give_up = time.time() + 60
    final_err = ValueError("timed out")
    while time.time() < give_up:
        try:
            await run('scp', ssh_opts, src_path, f'root@localhost:{path_on_pi}')
        except ValueError as err:
            final_err = err
            log.info('waiting for qemu to boot...')
            await asyncio.sleep(8)
            continue
        break
    else:
        raise final_err

    cmd_on_pi = ['sh', path_on_pi, PI_HOSTNAME]
    await run('ssh', ssh_opts, 'root@localhost', *cmd_on_pi)


async def _get_iscsi_device(better_be_the_iscsi_device='sde') -> Path:
    # don't screw up- this device is about to get formatted!
    dev_path = Path(f'/dev/{better_be_the_iscsi_device}')

    iscsi = await get_output('iscsiadm', '-m', 'session', '-P3')
    for m in re.findall(r'Attached scsi disk (\S+)', iscsi):
        if f'/dev/{m}' != str(dev_path):
            raise ValueError(f'surprised by attached iscsi disk {m!r} (try `iscsiadm -m node --logoutall=all`)')

    fdisk = await get_output('fdisk', '-l', dev_path)
    for m in re.findall(r'Disk model: (\S+)', fdisk):
        if m != 'VIRTUAL-DISK':
            raise ValueError(f'surprised that {dev_path} is model {m!r} instead of "VIRTUAL-DISK"')

    return dev_path


@step
async def fs_to_iscsi():
    async with mount(WORK, WORK / 'raspios.img', IMG_ROOT_OFFSET) as img_root:
        async with iscsi_login('-m', 'node', '-T', f'iqn.2024-03.com.bigasterisk:{PI_HOSTNAME}.target', '-p', '10.2.0.133'):
            dev = await _get_iscsi_device()
            await run('mkfs.ext4', '-F', dev)
            async with mount(WORK, dev, 0) as iscsi_root:
                await run(
                    'rsync',
                    '-r', # recurs
                    '-l', # symlinks as symlinks
                    '-p', # perms
                    '-t', # mtimes
                    '-g', # group
                    '-o', # owner
                    '-D', # devices + specials
                    '-H', # hard links
                    ['--exclude', '/boot'],
                    ['--exclude', '/dev'],
                    ['--exclude', '/mnt'],
                    ['--exclude', '/proc'],
                    ['--exclude', '/sys'],
                    str(img_root) + '/',
                    str(iscsi_root) + '/',
                )
                for d in [
                        'boot',
                        'dev',
                        'mnt',
                        'proc',
                        'sys',
                ]:
                    (iscsi_root / d).mkdir()
                (iscsi_root / 'etc/fstab').write_text('''
proc            /proc           proc    defaults          0       0
/dev/sda        /               ext4    defaults,noatime  0       1
''')
                # don't open cursesui for making first user
                (iscsi_root / 'etc/systemd/system/multi-user.target.wants/userconfig.service').unlink()
                log.info("there may be a delay here for flushing all the new files")


@step
async def setup_tftp():
    async with sshfs(WORK, TFTP_SPEC) as tftp:
        tftp_host_dir = tftp / f'{PI_HOSTNAME}-boot'
        async with mount(WORK, WORK / 'raspios.img', IMG_ROOT_OFFSET) as img_root:
            await run(
                'rsync',
                    '-r', # recurs
                    '-l', # symlinks as symlinks
                    '-p', # perms
                    #'-t', # mtimes
                    '-g', # group
                    '-o', # owner
                    '-D', # devices + specials
                    #'-H', # hard links
                '--delete',
                str(img_root) + '/boot/',
                str(tftp_host_dir) + '/',
            )
        async with mount(WORK, WORK / 'raspios.img', IMG_BOOT_OFFSET) as img_boot:
            await run(
                'rsync',
                    '-r', # recurs
                    '-l', # symlinks as symlinks
                    '-p', # perms
                    '-t', # mtimes
                    '-g', # group
                    '-o', # owner
                    '-D', # devices + specials
                    '-H', # hard links
                # (no delete)
                str(img_boot) + '/',
                str(tftp_host_dir) + '/firmware/',
            )
        await run('ln', '-sf', f'{PI_HOSTNAME}-boot/firmware/', str(tftp / PI_SERIAL))
        kernel_cmdline = [
            "console=serial0,115200",
            "console=tty1",
            f"ip=::::{PI_HOSTNAME}:eth0:dhcp",
            "root=/dev/sda",
            f"ISCSI_INITIATOR=iqn.2024-03.com.bigasterisk:{PI_HOSTNAME}.initiator",
            f"ISCSI_TARGET_NAME=iqn.2024-03.com.bigasterisk:{PI_HOSTNAME}.target",
            "ISCSI_TARGET_IP=10.2.0.133",
            "ISCSI_TARGET_PORT=3260",
            "rw",
            "rootfstype=ext4",
            "fsck.repair=yes",
            "rootwait",
            "cgroup_enable=cpuset",
            "cgroup_memory=1",
            "cgroup_enable=memory",
        ]
        (tftp_host_dir/'firmware/cmdline.txt').write_text(' '.join(kernel_cmdline))


async def main():
    global PI_HOSTNAME, PI_SERIAL

    if sys.argv[1:] == ['--init']:
        await init_work_dir()
    elif sys.argv[1] == '--iscsi':
        [PI_HOSTNAME] = sys.argv[2:]
        async with iscsi_login('-m', 'node', '-T', f'iqn.2024-03.com.bigasterisk:{PI_HOSTNAME}.target', '-p', '10.2.0.133'):
            dev = await _get_iscsi_device()
            async with mount(WORK, dev, 0) as iscsi_root:
                input(f"mounted {PI_HOSTNAME}'s iscsi drive at {iscsi_root}: ")
    else:
        PI_HOSTNAME, PI_SERIAL = sys.argv[1:]
        # todo: dd and add iscsi volume (handled by pyinfra)
        #   dd if=/dev/zero of={out} count=0 bs=1 seek=5G conv=excl || true
        #   edit /etc/tgt/conf.d/{pi_hostname}.conf etc

        await unpack_fresh_img_copy()
        await extract_boot_files()
        await setup_ssh_login()

        qemu_task = asyncio.create_task(qemu())
        await asyncio.sleep(20)  # inital qemu startup delay
        await setup_pi_in_emulator()  # finishes with a poweroff
        log.info('waiting for qemu to exit')
        await qemu_task

        await fs_to_iscsi()
        await setup_tftp()

        log.info(f'🎉 {PI_HOSTNAME} is ready for net boot')


asyncio.run(main())