import atexit
import os
import subprocess
import click
from logger import logging
from chroot import create_chroot, create_chroot_user, get_chroot_path
from constants import DEVICES, FLAVOURS


def get_device_and_flavour() -> tuple[str, str]:
    if not os.path.exists('.device'):
        logging.fatal(f'Please set the device using \'kupferbootstrap image device ...\'')
        exit(1)
    if not os.path.exists('.flavour'):
        logging.fatal(f'Please set the flavour using \'kupferbootstrap image flavour ...\'')
        exit(1)

    with open('.device', 'r') as file:
        device = file.read()
    with open('.flavour', 'r') as file:
        flavour = file.read()

    return (device, flavour)


def get_image_name(device, flavour) -> str:
    return f'{device}-{flavour}-rootfs.img'


def mount_rootfs_image(image_path, mount_path):
    if not os.path.exists(mount_path):
        os.makedirs(mount_path)

    def umount():
        subprocess.run(
            [
                'umount',
                '-lc',
                mount_path,
            ],
            stderr=subprocess.DEVNULL,
        )

    atexit.register(umount)

    result = subprocess.run([
        'mount',
        '-o',
        'loop',
        image_path,
        mount_path,
    ])
    if result.returncode != 0:
        logging.fatal(f'Failed to loop mount {image_path} to {mount_path}')
        exit(1)


def dump_bootimg(image_name: str) -> str:
    path = '/tmp/boot.img'
    result = subprocess.run([
        'debugfs',
        image_name,
        '-R',
        f'dump /boot/boot.img {path}',
    ])
    if result.returncode != 0:
        logging.fatal(f'Faild to dump boot.img')
        exit(1)
    return path


def dump_lk2nd(image_name: str) -> str:
    """
    This doesn't append the image with the appended DTB which is needed for some devices, so it should get added in the future.
    """
    path = '/tmp/lk2nd.img'
    result = subprocess.run([
        'debugfs',
        image_name,
        '-R',
        f'dump /boot/lk2nd.img {path}',
    ])
    if result.returncode != 0:
        logging.fatal(f'Faild to dump lk2nd.img')
        exit(1)
    return path


def dump_qhypstub(image_name: str) -> str:
    path = '/tmp/qhypstub.bin'
    result = subprocess.run([
        'debugfs',
        image_name,
        '-R',
        f'dump /boot/qhypstub.bin {path}',
    ])
    if result.returncode != 0:
        logging.fatal(f'Faild to dump qhypstub.bin')
        exit(1)
    return path


@click.group(name='image')
def cmd_image():
    pass


@click.command(name='device')
@click.argument('device')
def cmd_device(device):
    for key in DEVICES.keys():
        if '-'.join(key.split('-')[1:]) == device:
            device = key
            break

    if device not in DEVICES:
        logging.fatal(f'Unknown device {device}. Pick one from:\n{", ".join(DEVICES.keys())}')
        exit(1)

    logging.info(f'Setting device to {device}')

    with open('.device', 'w') as file:
        file.write(device)


@click.command(name='flavour')
@click.argument('flavour')
def cmd_flavour(flavour):
    if flavour not in FLAVOURS:
        logging.fatal(f'Unknown flavour {flavour}. Pick one from:\n{", ".join(FLAVOURS.keys())}')
        exit(1)

    logging.info(f'Setting flavour to {flavour}')

    with open('.flavour', 'w') as file:
        file.write(flavour)


@click.command(name='build')
def cmd_build():
    device, flavour = get_device_and_flavour()
    image_name = get_image_name(device, flavour)

    if not os.path.exists(image_name):
        result = subprocess.run([
            'fallocate',
            '-l',
            '4G',
            image_name,
        ])
        if result.returncode != 0:
            logging.fatal(f'Failed to allocate {image_name}')
            exit(1)

        result = subprocess.run([
            'mkfs.ext4',
            '-L',
            'kupfer',
            image_name,
        ])
        if result.returncode != 0:
            logging.fatal(f'Failed to create ext4 filesystem on {image_name}')
            exit(1)
    chroot_name = f'rootfs_{device}-{flavour}'
    rootfs_mount = get_chroot_path(chroot_name)
    mount_rootfs_image(image_name, rootfs_mount)

    extra_repos = {
        'main': {
            'Server': 'https://gitlab.com/kupfer/packages/prebuilts/-/raw/main/$repo',
        },
        'device': {
            'Server': 'https://gitlab.com/kupfer/packages/prebuilts/-/raw/main/$repo',
        },
    }

    if os.path.exists('/prebuilts'):
        extra_repos = {
            'main': {
                'Server': 'file:///prebuilts/$repo',
            },
            'device': {
                'Server': 'file:///prebuilts/$repo',
            },
        }

    create_chroot(
        chroot_name,
        packages=['base', 'base-kupfer'] + DEVICES[device] + FLAVOURS[flavour],
        pacman_conf='/app/local/etc/pacman.conf',
        extra_repos=extra_repos,
    )
    create_chroot_user(chroot_name)


"""
This doesn't work, because the mount isn't passed through to the real host
"""
"""
@click.command(name='inspect')
def cmd_inspect():
    device, flavour = get_device_and_flavour()
    image_name = get_image_name(device, flavour)

    rootfs_mount = get_chroot_path(f'rootfs_{device}-flavour')
    mount_rootfs_image(image_name, rootfs_mount)

    logging.info(f'Inspect the rootfs image at {rootfs_mount}')

    signal.pause()
"""

cmd_image.add_command(cmd_device)
cmd_image.add_command(cmd_flavour)
cmd_image.add_command(cmd_build)
# cmd_image.add_command(cmd_inspect)