#
# 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.disk import get_efi_partition

import fdisk

EFIBOOTMGR_BIN='/opt/opengnsys/bin/efibootmgr'

def _find_bootentry(entries, label):
    for entry in entries:
        if entry['description'] == label:
            return entry
    else:
        raise NameError('Boot entry {label} not found')


def _strip_boot_prefix(entry):
    try:
        num = entry['name'][4:]
    except:
        raise KeyError('Unable to strip "Boot" prefix from boot entry')
    return num


def _check_efibootmgr_json():
    """
    Checks if efibootmgr --json option is supported.
    Returns True if supported, False otherwise.
    """
    supported = True
    try:
        subprocess.run([EFIBOOTMGR_BIN, '--json'], stdout=subprocess.DEVNULL,
                       stderr=subprocess.DEVNULL, check=True)
    except subprocess.CalledProcessError:
        supported = False
    return supported


def is_uefi_supported(disknum):
    is_supported = os.path.exists("/sys/firmware/efi")

    if is_supported:
        try:
            get_efi_partition(disknum)
        except Exception as e:
            logging.info(e)
            is_supported = False

    return is_supported


def run_efibootmgr_json():
    if _check_efibootmgr_json() is False:
        raise RuntimeError(f'{EFIBOOTMGR_BIN} not available')

    proc = subprocess.run([EFIBOOTMGR_BIN, '--json'], capture_output=True, text=True)
    dict_json = json.loads(proc.stdout)

    return dict_json


def efibootmgr_bootnext(description):
    logging.info(f'Setting BootNext to value from boot entry with label {description}')
    bootnext_cmd = '{efibootmgr} -n {bootnum}'
    boot_entries = run_efibootmgr_json().get('vars', [])

    entry = _find_bootentry(boot_entries, description)
    num = _strip_boot_prefix(entry)  # efibootmgr output uses BootXXXX for each entry, remove the "Boot" prefix.
    bootnext_cmd = bootnext_cmd.format(bootnum=num, efibootmgr=EFIBOOTMGR_BIN)
    subprocess.run(shlex.split(bootnext_cmd), check=True,
                   stdout=subprocess.DEVNULL)


def efibootmgr_delete_bootentry(label):
    dict_json = run_efibootmgr_json()
    efibootmgr_cmd = '{efibootmgr} -q -b {bootnum} -B'

    for entry in dict_json.get('vars', []):
        if entry['description'] == label:
            num = entry['name'][4:]  # Remove "Boot" prefix to extract num
            efibootmgr_cmd = efibootmgr_cmd.format(bootnum=num, efibootmgr=EFIBOOTMGR_BIN)
            subprocess.run(shlex.split(efibootmgr_cmd), check=True)
            break
    else:
        logging.info(f'Cannot delete boot entry {label} because it was not found.')


def efibootmgr_create_bootentry(disk, part, loader, label, add_to_bootorder=True):
    entries = run_efibootmgr_json().get('vars', [])
    create_opt = '-c' if add_to_bootorder else '-C'
    index_opt = f'-I {len(entries)}'  # Requires efibootmgr version 18
    efibootmgr_cmd = f'{EFIBOOTMGR_BIN} {create_opt} {index_opt} -d {disk} -p {part} -L {label} -l {loader}'
    logging.info(f'{EFIBOOTMGR_BIN} command creating boot entry: {efibootmgr_cmd}')
    try:
        proc = subprocess.run(shlex.split(efibootmgr_cmd), check=True, text=True)
    except subprocess.CalledProcessError as e:
        logging.error(f'Unexpected error adding boot entry to nvram. UEFI firmware might be buggy?.')
        raise e