#!/usr/bin/env python

import argparse
import os
import shutil
import subprocess
import sys

from pathlib import Path


SRC_DIR = Path('/usr/src')
BOOT_DIR = Path('/boot')


def update_config(old_dir: Path, new_dir: Path, make_cmd: list[str]) -> None:
	if old_dir == new_dir:
		return
	old_config = old_dir/'.config'
	new_config = new_dir/'.config'
	while new_config.is_file():
		response = input('New config present. Overwrite? [y/N]').strip().lower()
		if response == 'y':
			break
		elif response == '' or response == 'n':
			return
		else:
			print("unrecognized option {}", response)
	print(f'Copying config from {old_config} to {new_config}')
	shutil.copy2(old_config, new_config)

	print(f'Setting symlink to {new_dir}')
	subprocess.run(['eselect', 'kernel', 'set', new_dir.name])
	print('Migrating config options')
	migrate = make_cmd + ['-C', new_dir.as_posix(), 'oldconfig']
	subprocess.run(migrate)
	menuconfig = make_cmd + ['-C', new_dir.as_posix(), 'menuconfig']
	while True:
		subprocess.run(menuconfig)
		response = input('Stop editing? [Y/n]').strip().lower()
		if response == '' or response == 'y':
			break
		elif response == 'n':
			continue
		else:
			print("unrecognized option {}", response)


def compile_kernel(new_dir: Path, make_cmd: list[str]) -> None:
	cc = make_cmd + ['-C', new_dir.as_posix()]
	subprocess.run(cc)


def install_kernel(kernel_dir: Path, make_cmd: list[str], kver: str) -> None:
	install_modules = make_cmd + ['-C', kernel_dir.as_posix(), 'modules_install']
	subprocess.run(install_modules)

	kver = kver[6:]
	# assumes proper dracut config in /etc
	dracut = ['dracut', f'--kver={kver}', '--force', '--no-machineid']
	res = subprocess.run(dracut)

	print(res)
	return
	# todo: sign uki
	uki_dir = res.split('')
	sign_uki = ['sbctl', 'bundle', '--save', '']
	subprocess.run(sign_uki)

def linux_folder(src_dir: Path, version: str) -> Path:
	revision = ''
	version = version.split('-')
	if len(version) > 1:
		revision = '-' + version[1]
	version = version[0]
	return (src_dir / (f'linux-{version}-gentoo{revision}'))


def update_kernel(boot_dir: Path, args: argparse.Namespace) -> None:
	old_dir = linux_folder(SRC_DIR, args.old_version)
	new_dir = linux_folder(SRC_DIR, args.new_version)
	new_version = new_dir.name
	make_cmd = ['make', f'-j{len(os.sched_getaffinity(0))}']
	# https://docs.kernel.org/kbuild/llvm.html

	if args.llvm:
		make_cmd.extend(['LLVM=1'])
	none_selected = not (args.backup or args.config or args.compile or args.install or args.rollback)
	if none_selected:
		update_config(old_dir, new_dir, make_cmd)
		compile_kernel(new_dir, make_cmd)
		install_kernel(new_dir, make_cmd, new_version)
	if args.config:
		update_config(old_dir, new_dir, make_cmd)
	if args.compile:
		compile_kernel(new_dir, make_cmd)
	if args.install:
		install_kernel(new_dir, make_cmd, new_version)

def main() -> None:
	parser = argparse.ArgumentParser(description='Convenience for manual kernel updates')
	subparsers = parser.add_subparsers()

	update = subparsers.add_parser('update',
		usage=f'{sys.argv[0]} update 5.15.12 5.16.3',
	)
	update.add_argument(
		'--llvm', '-l', action='store_true',
		help="Use clang/llvm to compile kernel")
	update.add_argument(
		'--backup', '-b', action='store_true',
		help="Backup old kernel files as .old")
	update.add_argument(
		'--rollback', '-r', action='store_true',
		help='Restore .old kernel files as main boot choice')
	update.add_argument(
		'--config', '-C', action='store_true',
		help='Migrate config from old kernel')
	update.add_argument(
		'--compile', '-c', action='store_true',
		help='Compile new kernel')
	update.add_argument(
		'--install', '-i', action='store_true',
		help='Install new kernel')
	update.add_argument('old_version', help='Old kernel version')
	update.add_argument('new_version', help='New kernel version')
	update.set_defaults(func=update_kernel)

	args = parser.parse_args()
	args.func(BOOT_DIR, args)


if __name__ == "__main__":
	main()