#
# Copyright (C) 2024 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 subprocess
import logging
import shutil
import shlex
import os

from enum import Enum
from src.utils.probe import *
from src.utils.bios import *
from src.utils.uefi import *
from src.utils.disk import *
from src.log import OgError

GRUB_TIMEOUT = 5

class GrubConfig:
    def __init__(self, timeout):
        self.timeout = timeout
        self.entries = []

    def add_entry(self, entry):
        self.entries.append(entry)

    def __str__(self):
        res = f'set timeout={self.timeout}\n'
        res += f'set default=0\n'
        res += '\n'
        for entry in self.entries:
            res += str(entry)
            res += '\n'
        return res

class MenuEntry:
    class PartType(Enum):
        MBR = 'part_msdos'
        GPT = 'part_gpt'

    class FsType(Enum):
        EXT = 'ext2'
        FAT = 'fat'

    def __init__(self, name, disk_num, part_num, part_type, fs):
        self.name = self._escape_menuentry_name(name)
        self.disk_num = disk_num
        self.part_num = part_num
        self.part_type = part_type
        self.fs = fs
        self.efipath = None
        self.initrd = None
        self.vmlinuz = None
        self.vmlinuz_root = None
        self.params = None

    def set_chainload_data(self, efipath):
        self.efipath = efipath

    def set_linux_data(self, initrd, vmlinuz, vmlinuz_root, params):
        self.initrd = initrd
        self.vmlinuz = vmlinuz
        self.vmlinuz_root = vmlinuz_root
        self.params = params

    def _set_root_str(self):
        if self.part_type == MenuEntry.PartType.GPT:
            return f'set root=(hd{self.disk_num - 1},gpt{self.part_num})'
        else:
            return f'set root=(hd{self.disk_num - 1},{self.part_num})'

    def _get_chainload_str(self):
        res =  f'menuentry "{self.name}" {{\n'
        res += f'    insmod {self.part_type.value}\n'
        res +=  '    insmod chain\n'
        res += f'    insmod {self.fs.value}\n'
        res += f'    {self._set_root_str()}\n'
        res += f'    chainloader {self.efipath}\n'
        res += '}\n'
        return res

    def _get_linux_str(self):
        res =  f'menuentry "{self.name}" {{\n'
        res += f'    insmod {self.part_type.value}\n'
        res +=  '    insmod linux\n'
        res += f'    insmod {self.fs.value}\n'
        res += f'    {self._set_root_str()}\n'
        res += f'    linux {self.vmlinuz} root={self.vmlinuz_root} {self.params}\n'
        res += f'    initrd {self.initrd}\n'
        res += '}\n'
        return res

    def __str__(self):
        if self.efipath:
            return self._get_chainload_str()
        else:
            return self._get_linux_str()

    @staticmethod
    def _escape_menuentry_name(entry_name):
        entry_name = entry_name.replace('"', r'\"')
        entry_name = entry_name.replace('\\', r'\\')
        entry_name = entry_name.replace('$', r'\$')
        entry_name = entry_name.replace('*', r'\*')
        entry_name = entry_name.replace('!', r'\!')
        return entry_name


def get_grub_boot_params(mountpoint, device):
    grub_conf = f'{mountpoint}/etc/default/grub'
    res = []

    with open(grub_conf, 'r') as f:
        for line in f:
            if line.find('=') == -1:
                continue
            key, value = line.split('=', 1)
            if key == 'GRUB_CMDLINE_LINUX' or key == 'GRUB_CMDLINE_LINUX_DEFAULT':
                value = value.replace('\n', '')
                value = value.strip('"')
                res.append(value)
    res.append(f'root={device}')
    return " ".join(res)


def _get_linux_data(disk_num, part_num, mountpoint):
    os_entry = {}
    os_entry['name'] = f'{os_probe(mountpoint)} ({disk_num}, {part_num})'
    os_entry['device'] = get_partition_device(disk_num, part_num)
    os_entry['part_type'] = get_disk_part_type(disk_num)

    kernel_path = get_vmlinuz_path(mountpoint)
    os_entry['vmlinuz'] = '/' + os.path.relpath(kernel_path, mountpoint)
    initrd_path = get_initrd_path(mountpoint)
    os_entry['initrd'] = '/' + os.path.relpath(initrd_path, mountpoint)
    return os_entry

def _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg):
    grub_config = GrubConfig(timeout=0)

    os_entry = _get_linux_data(disk_num, part_num, mountpoint)
    part_type = get_disk_part_type(disk_num)

    menu_entry = MenuEntry(name=os_entry['name'],
                           disk_num=disk_num,
                           part_num=part_num,
                           part_type=part_type,
                           fs=MenuEntry.FsType.EXT)
    menu_entry.set_linux_data(initrd=os_entry['initrd'],
                              vmlinuz=os_entry['vmlinuz'],
                              vmlinuz_root=os_entry['device'],
                              params='ro quiet splash')
    grub_config.add_entry(menu_entry)

    with open(grub_cfg, 'w') as f:
        f.write(str(grub_config))

def _mangle_efi_folder(entry_dir, boot_dir):
    efi_boot_src = f'{entry_dir}/EFI/BOOT'
    for file_name in os.listdir(efi_boot_src):
        shutil.move(f'{efi_boot_src}/{file_name}', boot_dir)

    shutil.rmtree(f'{entry_dir}/EFI')

    shutil.copyfile("/usr/lib/shim/shimx64.efi.signed",
                    f'{boot_dir}/shimx64.efi')

    shutil.copyfile(f'{boot_dir}/grubx64.efi',
                    f'{boot_dir}/ogloader.efi')

def _install_linux_grub_efi(disk_num, part_num, device, mountpoint):
    if interpreter_is64bit():
        arch = 'x86_64-efi'
    else:
        logging.warning("Old 32-bit UEFI system found here")
        arch = 'i386-efi'
    _esp_disk_num = 1

    esp, _esp_disk, _esp_part_number = get_efi_partition(_esp_disk_num, enforce_gpt=False)
    esp_mountpoint = esp.replace('dev', 'mnt')

    if not mount_mkdir(esp, esp_mountpoint):
        raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')

    _bootlabel = f'Part-{disk_num:02d}-{part_num:02d}'
    entry_dir = f'{esp_mountpoint}/EFI/{_bootlabel}'
    boot_dir = f'{entry_dir}/Boot'
    if os.path.exists(entry_dir):
        shutil.rmtree(entry_dir)
    os.makedirs(boot_dir)

    logging.info(f'Calling grub-install with target {arch} in {entry_dir}')
    grub_install_cmd = (f'grub-install --removable --no-nvram --uefi-secure-boot '
                        f'--target={arch} --efi-directory={entry_dir} '
                        f'--root-directory={entry_dir} --recheck')

    try:
        subprocess.run(shlex.split(grub_install_cmd), check=True)
    except subprocess.CalledProcessError as e:
        umount(esp_mountpoint)
        raise OgError(f"Error during GRUB install: {e}") from e

    _mangle_efi_folder(entry_dir, boot_dir)

    logging.info(f'Generating grub.cfg in {boot_dir}')
    grub_cfg = f'{boot_dir}/grub.cfg'
    try:
        _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg)
    except Exception as e:
        raise OgError(f'Error generating {grub_cfg}: {e}') from e
    finally:
        umount(esp_mountpoint)

def _install_linux_grub_bios(disk_num, part_num, device, mountpoint):
    arch = 'i386-pc'

    entry_dir = f'{mountpoint}/boot'
    grub_dir = f'{entry_dir}/grub'
    if os.path.exists(grub_dir):
        shutil.rmtree(grub_dir)
    os.makedirs(grub_dir)

    logging.info(f'Calling grub-install with target {arch} in {entry_dir}')
    grub_install_cmd = f'grub-install --force --target={arch} --boot-directory={entry_dir} --recheck {device}'

    try:
        subprocess.run(shlex.split(grub_install_cmd), check=True)
    except subprocess.CalledProcessError as e:
        umount(mountpoint)
        raise OgError(f"Error during GRUB install: {e}") from e

    grub_cfg = f'{grub_dir}/grub.cfg'

    logging.info(f'Generating grub.cfg in {grub_dir}')
    try:
        _generate_linux_grub_config(disk_num, part_num, mountpoint, grub_cfg)
    except Exception as e:
        raise OgError(f'Error generating {grub_cfg}: {e}') from e
    finally:
        umount(mountpoint)

def install_linux_grub(disk_num, part_num):
    device = get_partition_device(disk_num, part_num)

    mountpoint = device.replace('dev', 'mnt')
    if not mount_mkdir(device, mountpoint):
        raise OgError(f'Unable to mount {device} into {mountpoint}')

    try:
        if is_uefi_supported():
            logging.info(f'Installing GRUB for UEFI Linux at {device}...')
            _install_linux_grub_efi(disk_num, part_num, device, mountpoint)
        else:
            logging.info(f'Installing GRUB for BIOS Linux at {device}...')
            _install_linux_grub_bios(disk_num, part_num, device, mountpoint)
    finally:
        umount(mountpoint)

def _get_os_entries(esp_mountpoint):
    os_entries =  []
    available_disks = get_disks()
    for disk_num, diskname in enumerate(available_disks, start=1):
        disk_device = f'/dev/{diskname}'
        try:
            partitions_data = get_partition_data(device=disk_device)
        except OgError as e:
            continue

        for p in partitions_data:
            part_num = p.partno + 1
            mountpoint = p.padev.replace('dev', 'mnt')

            if mountpoint == esp_mountpoint:
                continue

            if not mount_mkdir(p.padev, mountpoint):
                raise OgError(f'Unable to mount {p.padev} into {mountpoint}')

            try:
                os_family = get_os_family(mountpoint)
                system_name = os_probe(mountpoint)
            except Exception as e:
                umount(mountpoint)
                raise

            if os_family == OSFamily.UNKNOWN:
                umount(mountpoint)
                continue

            os_entry = {}
            os_entry['name'] = f'{system_name} ({disk_num}, {part_num})'

            _bootlabel = f'Part-{disk_num:02d}-{part_num:02d}'
            try:
                if os_family == OSFamily.WINDOWS:
                    efi_loader = find_windows_efi_loader(esp_mountpoint, _bootlabel)
                elif os_family == OSFamily.LINUX:
                    efi_loader = find_linux_efi_loader(esp_mountpoint, _bootlabel)
                os_entry['efipath'] = '/' + os.path.relpath(efi_loader, esp_mountpoint)
            finally:
                umount(mountpoint)

            os_entries.append(os_entry)
    return os_entries

def _generate_main_grub_config(grub_cfg, esp_disk_num, esp_part_num, esp_mountpoint):
    os_entries = _get_os_entries(esp_mountpoint)

    esp_part_type = get_disk_part_type(disk_num=1)

    grub_config = GrubConfig(timeout=GRUB_TIMEOUT)
    for os_entry in os_entries:
        menu_entry = MenuEntry(name=os_entry['name'],
                               disk_num=esp_disk_num,
                               part_num=esp_part_num,
                               part_type=esp_part_type,
                               fs=MenuEntry.FsType.FAT)
        menu_entry.set_chainload_data(efipath=os_entry['efipath'])
        grub_config.add_entry(menu_entry)

    with open(grub_cfg, 'w') as f:
        f.write(str(grub_config))

def get_disk_part_type(disk_num):
    device = get_disks()[disk_num - 1]
    cxt = fdisk.Context(f'/dev/{device}')
    return MenuEntry.PartType.MBR if cxt.label.name == 'dos' else MenuEntry.PartType.GPT

def _update_nvram(esp_disk, esp_part_number):
    loader_path = '/EFI/grub/Boot/shimx64.efi'
    bootlabel = 'grub'
    efibootmgr_delete_bootentry(bootlabel)
    efibootmgr_create_bootentry(esp_disk, esp_part_number, loader_path, bootlabel, add_to_bootorder=False)
    efibootmgr_set_entry_order(bootlabel, 1)

def install_main_grub():
    disk_device = f'/dev/{get_disks()[0]}'
    is_uefi = is_uefi_supported()
    if not is_uefi:
        logging.info(f'Global GRUB install not supported in legacy BIOS')
        return

    logging.info(f'Installing GRUB at {disk_device}')
    _esp_disk_num = 1
    esp, _esp_disk, _esp_part_number = get_efi_partition(_esp_disk_num, enforce_gpt=False)
    esp_mountpoint = esp.replace('dev', 'mnt')

    if not mount_mkdir(esp, esp_mountpoint):
        raise OgError(f'Unable to mount detected EFI System Partition at {esp} into {esp_mountpoint}')

    entry_dir = f'{esp_mountpoint}/EFI/grub'
    boot_dir = f'{entry_dir}/Boot'
    if os.path.exists(entry_dir):
        shutil.rmtree(entry_dir)
    os.makedirs(boot_dir)

    if interpreter_is64bit():
        arch = 'x86_64-efi'
    else:
        logging.warning("Old 32-bit UEFI system found here")
        arch = 'i386-efi'
    logging.info(f'Calling grub-install with target {arch} in {entry_dir}')
    grub_install_cmd = (f'grub-install --removable --no-nvram --uefi-secure-boot '
                        f'--target={arch} --efi-directory={entry_dir} '
                        f'--root-directory={entry_dir} --recheck')

    try:
        subprocess.run(shlex.split(grub_install_cmd), check=True)
    except subprocess.CalledProcessError as e:
        umount(esp_mountpoint)
        raise OgError(f"Error during GRUB install: {e}") from e

    _mangle_efi_folder(entry_dir, boot_dir)

    logging.info(f'Generating grub.cfg in {boot_dir}')
    grub_cfg = f'{boot_dir}/grub.cfg'
    try:
        _generate_main_grub_config(grub_cfg, _esp_disk_num, _esp_part_number, esp_mountpoint)
    except Exception as e:
        umount(esp_mountpoint)
        raise OgError(f'Error generating {grub_cfg}: {e}') from e

    logging.info('Updating grub UEFI NVRAM entry')
    try:
        _update_nvram(_esp_disk, _esp_part_number)
    except Exception as e:
        logging.info(f'Error updating NVRAM: {e}')

    umount(esp_mountpoint)