#
# Copyright (C) 2023 Soleta Networks <info@soleta.eu>
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the
# Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.

import json
import logging
import os
import shlex
import subprocess

from src.utils.probe import OSFamily, get_os_family, get_linux_distro_id, os_probe
from src.utils.disk import get_partition_device, get_efi_partition
from src.utils.bios import *
from src.utils.uefi import *
from src.utils.fs import *


def _boot_bios_linux(disk, part, mountpoint):
    logging.info(f'Booting Linux system')

    if not get_linux_distro_id(mountpoint) == 'ubuntu':
        raise NotImplementedError(f'{os_probe(mountpoint)} detected, only Ubuntu is supported for legacy BIOS boot')

    kernel_path = get_vmlinuz_path(mountpoint)
    initrd_path = get_initrd_path(mountpoint)

    device = get_partition_device(disk, part)
    grub_boot_params = get_grub_boot_params(mountpoint, device)

    kexec_cmd = f'kexec -l {kernel_path} --append="{grub_boot_params}" --initrd="{initrd_path}"'
    kexec_reboot_cmd = 'kexec -e'

    logging.info(f'Booting with: {kexec_cmd}')
    try:
        subprocess.run(shlex.split(kexec_cmd), check=True, text=True)
        subprocess.run(shlex.split(kexec_reboot_cmd), check=True, text=True)
    except OSError as e:
        raise OSError(f'Error processing kexec: {e}') from e

def _boot_bios_windows(disk, part, mountpoint):
    logging.info(f'Booting Windows system')
    try:
        with open(f'{mountpoint}/ogboot.me', 'w') as f:
            f.write('\0' * (3072))
        with open(f'{mountpoint}/ogboot.firstboot', 'w') as f:
            f.write('iniciado')
            f.write('\0' * (3072 - 8))
        with open(f'{mountpoint}/ogboot.secondboot', 'w') as f:
            f.write('\0' * (3072))
    except OSError as e:
        raise OSError(f'Could not create ogboot files in Windows partition: {e}') from e

def _boot_uefi_windows(disk, part, mountpoint):
    logging.info(f'Booting windows system')
    bootlabel = f'Part-{disk:02d}-{part:02d}'
    esp, esp_disk, esp_part_number = get_efi_partition(disk, enforce_gpt=True)
    esp_mountpoint = esp.replace('dev', 'mnt')
    if not mount_mkdir(esp, esp_mountpoint):
        raise RuntimeError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')

    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/bootmgfw.efi',
                    f'{esp_mountpoint}/EFI/Microsoft/Boot/bootmgfw.efi']
    try:
        for efi_app in loader_paths:
            if os.path.exists(efi_app):
                loader = efi_app[len(esp_mountpoint):]
                logging.info(f'Found bootloader at ESP partition: {loader}')
                break
        else:
            raise RuntimeError(f'Unable to locate Windows EFI bootloader bootmgfw.efi')

        efibootmgr_delete_bootentry(bootlabel)
        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader, bootlabel)
        efibootmgr_bootnext(bootlabel)
    finally:
        umount(esp_mountpoint)

def _boot_uefi_linux(disk, part, mountpoint):
    logging.info(f'Booting Linux system')
    bootlabel = f'Part-{disk:02d}-{part:02d}'
    esp, esp_disk, esp_part_number = get_efi_partition(disk, enforce_gpt=False)
    esp_mountpoint = esp.replace('dev', 'mnt')
    if not mount_mkdir(esp, esp_mountpoint):
        raise RuntimeError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')

    loader_paths = [f'{esp_mountpoint}/EFI/{bootlabel}/Boot/shimx64.efi',
                    f'{esp_mountpoint}/EFI/ubuntu/shimx64.efi']
    try:
        for efi_app in loader_paths:
            if os.path.exists(efi_app):
                loader = efi_app[len(esp_mountpoint):]
                logging.info(f'Found bootloader at ESP partition: {loader}')
                break
        else:
            raise RuntimeError(f'Unable to locate Linux EFI bootloader shimx64.efi')

        efibootmgr_delete_bootentry(bootlabel)
        efibootmgr_create_bootentry(esp_disk, esp_part_number, loader, bootlabel)
        efibootmgr_bootnext(bootlabel)
    finally:
        umount(esp_mountpoint)

def boot_os_at(disk, part):
    logging.info(f'Booting disk={disk} partition={part}')
    device = get_partition_device(disk, part)
    mountpoint = device.replace('dev', 'mnt')
    if not mount_mkdir(device, mountpoint):
        raise RuntimeError(f'Cannot probe OS family. Unable to mount {device} at {esp_mountpoint}.')

    is_uefi = is_uefi_supported()
    if is_uefi:
        logging.info('UEFI support detected')
    else:
        logging.info('UEFI support not detected')

    os_family = get_os_family(mountpoint)
    logging.info(f'{os_family} detected at {device}.')

    try:
        if is_uefi and os_family == OSFamily.WINDOWS:
            _boot_uefi_windows(disk, part, mountpoint)
        elif is_uefi and os_family == OSFamily.LINUX:
            _boot_uefi_linux(disk, part, mountpoint)
        elif not is_uefi and os_family == OSFamily.WINDOWS:
            _boot_bios_windows(disk, part, mountpoint)
        elif not is_uefi and os_family == OSFamily.LINUX:
            _boot_bios_linux(disk, part, mountpoint)
        else:
            raise RuntimeError(f'Unknown OS family {os_family}')
    finally:
        umount(mountpoint)