# SPDX-License-Identifier: GPL-2.0
# Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.

import gzip
import os
import os.path
import pytest

import u_boot_utils

def mkdir_cond(dirname):
    """Create a directory if it doesn't already exist

    Args:
        dirname: Name of directory to create
    """
    if not os.path.exists(dirname):
        os.mkdir(dirname)

def setup_bootflow_image(u_boot_console):
    """Create a 20MB disk image with a single FAT partition"""
    cons = u_boot_console
    fname = os.path.join(cons.config.source_dir, 'mmc1.img')
    mnt = os.path.join(cons.config.persistent_data_dir, 'mnt')
    mkdir_cond(mnt)

    u_boot_utils.run_and_log(cons, 'qemu-img create %s 20M' % fname)
    u_boot_utils.run_and_log(cons, 'sudo sfdisk %s' % fname,
                             stdin=b'type=c')

    loop = None
    mounted = False
    complete = False
    try:
        out = u_boot_utils.run_and_log(cons,
                                       'sudo losetup --show -f -P %s' % fname)
        loop = out.strip()
        fatpart = '%sp1' % loop
        u_boot_utils.run_and_log(cons, 'sudo mkfs.vfat %s' % fatpart)
        u_boot_utils.run_and_log(
            cons, 'sudo mount -o loop %s %s -o uid=%d,gid=%d' %
            (fatpart, mnt, os.getuid(), os.getgid()))
        mounted = True

        vmlinux = 'vmlinuz-5.3.7-301.fc31.armv7hl'
        initrd = 'initramfs-5.3.7-301.fc31.armv7hl.img'
        dtbdir = 'dtb-5.3.7-301.fc31.armv7hl'
        script = '''# extlinux.conf generated by appliance-creator
ui menu.c32
menu autoboot Welcome to Fedora-Workstation-armhfp-31-1.9. Automatic boot in # second{,s}. Press a key for options.
menu title Fedora-Workstation-armhfp-31-1.9 Boot Options.
menu hidden
timeout 20
totaltimeout 600

label Fedora-Workstation-armhfp-31-1.9 (5.3.7-301.fc31.armv7hl)
        kernel /%s
        append ro root=UUID=9732b35b-4cd5-458b-9b91-80f7047e0b8a rhgb quiet LANG=en_US.UTF-8 cma=192MB cma=256MB
        fdtdir /%s/
        initrd /%s''' % (vmlinux, dtbdir, initrd)
        ext = os.path.join(mnt, 'extlinux')
        mkdir_cond(ext)

        with open(os.path.join(ext, 'extlinux.conf'), 'w') as fd:
            print(script, file=fd)

        inf = os.path.join(cons.config.persistent_data_dir, 'inf')
        with open(inf, 'wb') as fd:
            fd.write(gzip.compress(b'vmlinux'))
        u_boot_utils.run_and_log(cons, 'mkimage -f auto -d %s %s' %
                                 (inf, os.path.join(mnt, vmlinux)))

        with open(os.path.join(mnt, initrd), 'w') as fd:
            print('initrd', file=fd)

        mkdir_cond(os.path.join(mnt, dtbdir))

        dtb_file = os.path.join(mnt, '%s/sandbox.dtb' % dtbdir)
        u_boot_utils.run_and_log(
            cons, 'dtc -o %s' % dtb_file, stdin=b'/dts-v1/; / {};')
        complete = True
    except ValueError as exc:
        print('Falled to create image, failing back to prepared copy: %s',
              str(exc))
    finally:
        if mounted:
            u_boot_utils.run_and_log(cons, 'sudo umount %s' % mnt)
        if loop:
            u_boot_utils.run_and_log(cons, 'sudo losetup -d %s' % loop)

    if not complete:
        # Use a prepared image since we cannot create one
        infname = os.path.join(cons.config.source_dir,
                               'test/py/tests/bootstd/mmc1.img.xz')
        u_boot_utils.run_and_log(
            cons,
            ['sh', '-c', 'xz -dc %s >%s' % (infname, fname)])


@pytest.mark.buildconfigspec('ut_dm')
def test_ut_dm_init(u_boot_console):
    """Initialize data for ut dm tests."""

    fn = u_boot_console.config.source_dir + '/testflash.bin'
    if not os.path.exists(fn):
        data = b'this is a test'
        data += b'\x00' * ((4 * 1024 * 1024) - len(data))
        with open(fn, 'wb') as fh:
            fh.write(data)

    fn = u_boot_console.config.source_dir + '/spi.bin'
    if not os.path.exists(fn):
        data = b'\x00' * (2 * 1024 * 1024)
        with open(fn, 'wb') as fh:
            fh.write(data)

    # Create a file with a single partition
    fn = u_boot_console.config.source_dir + '/scsi.img'
    if not os.path.exists(fn):
        data = b'\x00' * (2 * 1024 * 1024)
        with open(fn, 'wb') as fh:
            fh.write(data)
        u_boot_utils.run_and_log(
            u_boot_console, f'sfdisk {fn}', stdin=b'type=83')

@pytest.mark.buildconfigspec('cmd_bootflow')
def test_ut_dm_init_bootstd(u_boot_console):
    """Initialise data for bootflow tests"""

    setup_bootflow_image(u_boot_console)

    # Restart so that the new mmc1.img is picked up
    u_boot_console.restart_uboot()


def test_ut(u_boot_console, ut_subtest):
    """Execute a "ut" subtest.

    The subtests are collected in function generate_ut_subtest() from linker
    generated lists by applying a regular expression to the lines of file
    u-boot.sym. The list entries are created using the C macro UNIT_TEST().

    Strict naming conventions have to be followed to match the regular
    expression. Use UNIT_TEST(foo_test_bar, _flags, foo_test) for a test bar in
    test suite foo that can be executed via command 'ut foo bar' and is
    implemented in C function foo_test_bar().

    Args:
        u_boot_console (ConsoleBase): U-Boot console
        ut_subtest (str): test to be executed via command ut, e.g 'foo bar' to
            execute command 'ut foo bar'
    """

    output = u_boot_console.run_command('ut ' + ut_subtest)
    assert output.endswith('Failures: 0')