"""RAK4631 application-only UF2 validation and USB-bound destination lookup.

UF2 format: https://github.com/microsoft/uf2
Board identity: meshtastic/nRF52_Bootloader, wiscore_rak4631_board/board.h.
No caller-selected mount, filename, block device, erase, or bootloader image.
"""
import hashlib
import json
import os
from pathlib import Path
import re
import stat
import struct

MAX_IMAGE = 2 * 1024 * 1024
FAMILY_NRF52840_S140_V6 = 0xADA52840
APPLICATION_START = 0x26000
APPLICATION_END = 0xF4000


def require(value, message):
    if not value:
        raise ValueError(message)


def validate_image(raw):
    require(type(raw) is bytes and 512 <= len(raw) <= MAX_IMAGE and len(raw) % 512 == 0,
            'invalid RAK4631 UF2 image size')
    count = len(raw) // 512
    addresses, numbers = set(), set()
    for offset in range(0, len(raw), 512):
        block = raw[offset:offset + 512]
        first, second, flags, address, size, number, total, family = struct.unpack('<8I', block[:32])
        require((first, second, struct.unpack('<I', block[508:])[0]) ==
                (0x0A324655, 0x9E5D5157, 0x0AB16F30), 'invalid UF2 block magic')
        require(flags == 0x2000 and family == FAMILY_NRF52840_S140_V6,
                'RAK4631 requires application UF2 for nRF52840/S140 v6')
        require(size == 256 and address % 256 == 0 and
                APPLICATION_START <= address < APPLICATION_END,
                'UF2 would overwrite reserved firmware or settings')
        require(total == count and number < count and number not in numbers and address not in addresses,
                'missing, duplicate or overlapping UF2 blocks')
        numbers.add(number)
        addresses.add(address)
    require(addresses == set(range(APPLICATION_START, APPLICATION_START + count * 256, 256)),
            'UF2 application has a missing region or unexpected start address')
    return {'size': len(raw), 'blocks': count, 'sha256': hashlib.sha256(raw).hexdigest(),
            'family': FAMILY_NRF52840_S140_V6}


def validate_board_info(raw):
    require(len(raw) <= 4096, 'UF2 board information exceeds bound')
    fields = {}
    for line in raw.decode('utf-8', errors='strict').splitlines():
        if ':' in line:
            key, value = line.split(':', 1)
            require(key not in fields, 'ambiguous UF2 board information')
            fields[key] = value.strip()
    require(fields.get('Board-ID') == 'WisBlock-RAK4631-Board',
            'USB bootloader is not a RAK4631')
    softdevice = fields.get('SoftDevice')
    if softdevice is not None:
        require(re.fullmatch(r'S140\s+6\.[0-9]+\.[0-9]+', softdevice) is not None,
                'this RAK4631 firmware requires S140 v6; bootloader was not changed')


def usb_parent(path):
    path = Path(path).resolve(strict=True)
    for candidate in (path, *path.parents):
        if (candidate / 'idVendor').is_file() and (candidate / 'idProduct').is_file():
            return candidate
    raise ValueError('device has no physical USB parent')


def usb_identity(path):
    meta = path.stat()
    def read(name):
        with (path / name).open('rb') as stream:
            raw = stream.read(257)
        require(0 < len(raw) <= 256, 'USB identity exceeds bound')
        return raw.strip().decode('ascii')
    return (meta.st_dev, meta.st_ino, read('idVendor'), read('idProduct'), read('serial'))


def mount_path(value):
    # mountinfo escapes whitespace and backslashes with octal sequences.
    return re.sub(r'\\([0-7]{3})', lambda m: chr(int(m[1], 8)), value)


def open_boot_volume(serial_path):
    """Return a retained directory FD belonging to this physical USB radio.

    The volume must already be mounted by the OS. Never mount an arbitrary
    block device or search for a matching drive label elsewhere on the host.
    """
    require(re.fullmatch(r'/dev/ttyACM[0-9]+', serial_path) is not None,
            'RAK4631 bootloader requires its native USB serial port')
    radio = usb_parent(Path('/sys/class/tty') / Path(serial_path).name / 'device')
    identity = usb_identity(radio)
    require(identity[2:4] == ('239a', '0029'), 'RAK4631 is not in UF2 bootloader mode')
    devices = set()
    blocks = list(Path('/sys/class/block').iterdir())
    require(len(blocks) <= 4096, 'USB block inventory exceeds bound')
    for block in blocks:
        try:
            if usb_parent(block) == radio:
                with (block / 'dev').open() as stream:
                    number = stream.read(65).strip()
                require(re.fullmatch('[0-9]+:[0-9]+', number) is not None, 'invalid USB block identity')
                devices.add(number)
        except (FileNotFoundError, ValueError):
            continue
    candidates = []
    with Path('/proc/self/mountinfo').open() as stream:
        mounts = stream.read(1024 * 1024 + 1)
    require(len(mounts) <= 1024 * 1024, 'mount inventory exceeds bound')
    for line in mounts.splitlines():
        fields = line.split()
        separator = fields.index('-')
        if fields[2] in devices and fields[separator + 1] in ('vfat', 'msdos') and 'rw' in fields[5].split(','):
            candidates.append((mount_path(fields[4]), fields[2]))
    require(len(candidates) == 1, 'mount the RAK4631 bootloader drive before starting its firmware update')
    path, number = candidates[0]
    directory = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC)
    try:
        require(os.fstat(directory).st_dev == os.makedev(*map(int, number.split(':'))),
                'mounted UF2 device changed')
        info = os.open('INFO_UF2.TXT', os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK | os.O_CLOEXEC, dir_fd=directory)
        try:
            require(stat.S_ISREG(os.fstat(info).st_mode), 'UF2 board information is not a regular file')
            validate_board_info(os.read(info, 4097))
        finally:
            os.close(info)
        require(usb_identity(radio) == identity, 'USB radio changed before firmware update')
        return directory
    except BaseException:
        os.close(directory)
        raise


def open_destination(serial_path, nonce):
    require(re.fullmatch('[a-f0-9]{32}', nonce) is not None, 'invalid firmware job nonce')
    directory = open_boot_volume(serial_path)
    try:
        # Exclusive creation avoids replacing any file on the mounted volume.
        return os.open('archy-' + nonce + '.uf2',
                       os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC,
                       0o600, dir_fd=directory)
    finally:
        os.close(directory)


def validate_fd(image_fd):
    meta = os.fstat(image_fd)
    require(stat.S_ISREG(meta.st_mode) and 512 <= meta.st_size <= MAX_IMAGE,
            'invalid UF2 image descriptor')
    raw = os.pread(image_fd, MAX_IMAGE + 1, 0)
    return validate_image(raw)


def main(arguments):
    # Protected owner supplies an immutable image and a single newly-created
    # destination FD. The helper never resolves host paths in this mode.
    if len(arguments) == 2 and arguments[0] == '--archy-uf2-owned':
        nonce = arguments[1]
        require(re.fullmatch('[a-f0-9]{32}', nonce) is not None, 'invalid firmware job nonce')
        receipt = copy_image(5, 6)
    elif len(arguments) == 4 and arguments[0] == '--archy-uf2':
        _, port, image, nonce = arguments
        source = os.open(image, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC)
        try:
            validate_fd(source)  # Do not touch the volume for a bad image.
            destination = open_destination(port, nonce)
            try:
                receipt = copy_image(source, destination)
            finally:
                os.close(destination)
        finally:
            os.close(source)
    else:
        raise ValueError('invalid fixed UF2 operation')
    print(json.dumps({**receipt, 'nonce': nonce}, separators=(',', ':')))


def copy_image(image_fd, destination_fd):
    """Copy a validated complete image to the owner's already-open USB file.

    This receipt means bytes reached the bootloader volume. Post-boot firmware
    verification is a separate step and must never be inferred from this copy.
    """
    require(stat.S_ISREG(os.fstat(image_fd).st_mode) and os.fstat(image_fd).st_size <= MAX_IMAGE,
            'UF2 source is not a bounded regular file')
    os.lseek(image_fd, 0, os.SEEK_SET)
    raw = b''
    while len(raw) <= MAX_IMAGE:
        block = os.read(image_fd, min(65536, MAX_IMAGE + 1 - len(raw)))
        if not block:
            break
        raw += block
    receipt = validate_image(raw)
    require(stat.S_ISREG(os.fstat(destination_fd).st_mode) and os.fstat(destination_fd).st_size == 0,
            'UF2 destination is not a new regular file')
    for start in range(0, len(raw), 512):
        block = memoryview(raw)[start:start + 512]
        while block:
            written = os.write(destination_fd, block)
            require(written > 0, 'UF2 transfer stopped before completion')
            block = block[written:]
    os.fsync(destination_fd)
    return {**receipt, 'contract': 1, 'board': 'rak4631', 'state': 'written-awaiting-device-verification'}