#!/usr/bin/env python3
# Copyright (C) 2023 Rudra Saraswat
#
# This file is part of akshara.
#
# akshara is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# akshara is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with akshara.  If not, see <http://www.gnu.org/licenses/>.


import os
import sys
import time
import yaml
import shutil
import filecmp
import argparse
import requests
import platform
import fasteners
from threading import Event
import subprocess

__version = '1.0.0'


# Colors


class colors:
    reset = '\033[0m'
    bold = '\033[01m'
    disable = '\033[02m'
    underline = '\033[04m'
    reverse = '\033[07m'
    strikethrough = '\033[09m'
    invisible = '\033[08m'

    class fg:
        black = '\033[30m'
        red = '\033[31m'
        green = '\033[32m'
        orange = '\033[33m'
        blue = '\033[34m'
        purple = '\033[35m'
        cyan = '\033[36m'
        lightgrey = '\033[37m'
        darkgrey = '\033[90m'
        lightred = '\033[91m'
        lightgreen = '\033[92m'
        yellow = '\033[93m'
        lightblue = '\033[94m'
        pink = '\033[95m'
        lightcyan = '\033[96m'

    class bg:
        black = '\033[40m'
        red = '\033[41m'
        green = '\033[42m'
        orange = '\033[43m'
        blue = '\033[44m'
        purple = '\033[45m'
        cyan = '\033[46m'
        lightgrey = '\033[47m'


def exec(*cmd, **kwargs):
    return subprocess.call(cmd, shell=False, stdout=sys.stdout, stderr=sys.stderr, **kwargs)


def exec_chroot(*cmd, **kwargs):
    exec('mount', '--bind', '/.new_rootfs', '/.new_rootfs')
    ret = exec('arch-chroot', '/.new_rootfs', *cmd, **kwargs)
    exec('umount', '-l', '/.new_rootfs')
    return ret


fg = colors.fg()


def info(msg):
    print(colors.bold + fg.cyan + '[INFO] ' +
          colors.reset + msg + colors.reset)


def warn(warning):
    print(colors.bold + fg.yellow + '[WARNING] ' +
          colors.reset + warning + colors.reset)


def error(err):
    print(colors.bold + fg.red + '[ERROR] ' +
          colors.reset + err + colors.reset)


def interpret_track(blend_release):
    result = yaml.safe_load(requests.get(blend_release.get(
        'impl') + '/' + blend_release.get('track') + '.yaml', allow_redirects=True).content.decode())

    if (type(result.get('impl')) == str and
            type(result.get('track')) != 'custom'):
        res = interpret_track(result)

        for i in res.keys():
            if type(res[i]) is list:
                if type(result.get(i)) is list:
                    result[i] = res[i] + result[i]

    return result


def update_system():
    if os.path.isdir('/.update_rootfs'):
        error('update already downloaded, you must reboot first')
        sys.exit(75)

    os.chdir('/')

    if exec('rm', '-rf', '/.new_rootfs', '/.new.etc', '/.new.var.lib') != 0:
        exec('umount', '-l', '/.new_rootfs')
        exec('rm', '-rf', '/.new_rootfs', '/.new.etc', '/.new.var.lib')

    # Check if update is available
    if not os.path.isfile('/system.yaml'):
        error('system.yaml does not exist in /')
        sys.exit(100)

    with open('/system.yaml') as blend_release_file:
        blend_release = yaml.load(
            blend_release_file, Loader=yaml.FullLoader)

        # TODO: Add check that all packages actually exist

        info('downloading Arch tarball...')

        # The mirror to use for downloading the bootstrap image
        # For example, for the Arch mirror at mirrors.acm.wpi.edu, you'd use https://mirrors.acm.wpi.edu/archlinux
        arch_repo = blend_release.get("arch-repo")
        if type(arch_repo) != str:
            # default arch and bootstrap repo
            arch_repo = "geo.mirror.pkgbuild.com"

        bootstrap_repo = blend_release.get("bootstrap-repo")
        if type(bootstrap_repo) != str:
            bootstrap_repo = arch_repo

        # TODO: default to https if http/https isn't listed
        # keeping disabled for consistency because i can't find `repo` to add it (and to add a fallback)
        if not (arch_repo.startswith("https://") or arch_repo.startswith("http://")):
            arch_repo = "https://" + arch_repo
        if not (bootstrap_repo.startswith("https://") or bootstrap_repo.startswith("http://")):
            bootstrap_repo = "https://" + bootstrap_repo

        # same (fallback and https) for repo
        repo = blend_release.get("repo")
        if type(repo) != str:
            repo = "https://pkg-repo.blendos.co"
        elif not (repo.startswith("https://") or repo.startswith("http://")):
            repo = "https://" + repo

        if not os.path.isfile('/.update.tar.zst'):
            if exec('wget', '-q', '--show-progress', f'{bootstrap_repo}/iso/latest/archlinux-bootstrap-x86_64.tar.zst', '-O', '/.update.tar.zst') != 0:
                warn('failed download')
                print()
                info('trying download again...')
                print()
                exec('rm', '-f', '/.update.tar.zst')
                if exec('wget', '-q', '--show-progress', f'{bootstrap_repo}/iso/latest/archlinux-bootstrap-x86_64.tar.zst', '-O', '/.update.tar.zst') != 0:
                    error('failed download')
                    print()
                    error('update failed')
                    sys.exit(50)

        if exec('bash', '-c', f'sha256sum -c --ignore-missing <(wget -qO- {bootstrap_repo}/iso/latest/sha256sums.txt | sed "s/archlinux-bootstrap-x86_64\\.tar\\.zst/.update.tar.zst/g") 2>/dev/null') != 0:
            error('failed checksum verification')
            print()
            info('trying download again...')
            exec('rm', '-f', '/.update.tar.zst')
            if exec('wget', '-q', '--show-progress', f'{bootstrap_repo}/iso/latest/archlinux-bootstrap-x86_64.tar.zst', '-O', '/.update.tar.zst') != 0:
                error('failed download')
                print()
                error('update failed')
                sys.exit(50)
                return
            if exec('bash', '-c', f'sha256sum -c --ignore-missing <(wget -qO- {bootstrap_repo}/iso/latest/sha256sums.txt | sed "s/archlinux-bootstrap-x86_64\\.tar\\.zst/.update.tar.zst/g") 2>/dev/null') != 0:
                error('failed checksum verification')
                print()
                error('update failed')
                sys.exit(25)
                return

        info('checksum verification was successful')

        print()

        info('generating new system...')

        exec('tar', '--acls', '--xattrs', '-xf', '.update.tar.zst')
        exec('mv', 'root.x86_64', '.new_rootfs')
        exec('rm', '-f', '/.new_rootfs/pkglist.x86_64.txt')
        exec('rm', '-f', '/.new_rootfs/version')

        packages = [
            'akshara',
            'blend'
        ]

        aur_packages = []

        services = [
            'akshara'
        ]

        user_services = [
            'blend-files'
        ]

        keep = [
            '/usr/share'
        ]

        if (type(blend_release.get('impl')) == str and
                type(blend_release.get('track')) != 'custom'):
            res = interpret_track(blend_release)

            for i in res.keys():
                if type(res[i]) is list:
                    if type(blend_release.get(i)) is list:
                        blend_release[i] += res[i]
                    else:
                        blend_release[i] = res[i]

        if type(blend_release.get('packages')) == list:
            packages += blend_release.get('packages')

        if type(blend_release.get('aur-packages')) == list:
            packages += ['fakeroot', 'chaotic-aur/paru']
            aur_packages += blend_release.get('aur-packages')

        if type(blend_release.get('services')) == list:
            services += blend_release.get('services')

        if type(blend_release.get('user-services')) == list:
            user_services += blend_release.get('user-services')

        exec_chroot('rm', '-f', '/.new_rootfs/etc/resolv.conf')

        with open('/.new_rootfs/etc/resolv.conf', 'w') as pacman_mirrorlist_conf:
            pacman_mirrorlist_conf.write('nameserver 1.1.1.1\n')

        with open('/.new_rootfs/etc/pacman.d/mirrorlist', 'w') as pacman_mirrorlist_conf:
            pacman_mirrorlist_conf.write(
                    f'Server = {arch_repo}/$repo/os/$arch\n')

        exec_chroot('mkdir', '-p', '/var/cache/pacman/pkg')
        exec_chroot('rm', '-rf', '/var/cache/pacman/pkg')
        exec('cp', '-r', '/var/cache/pacman/pkg',
             '/.new_rootfs/var/cache/pacman')

        # update packages
        exec_chroot('pacman-key', '--init')
        exec_chroot('pacman-key', '--populate')

        counter = 0
        while True:
            return_val = exec_chroot('pacman', '-Sy', '--ask=4', 'reflector')
            counter += 1
            if counter > 30:
                error('failed to download packages')
                exit(50)
            if return_val == 0:
                break

        exec_chroot('reflector', '--latest', '5', '--protocol', 'https',
                    '--sort', 'rate', '--save', '/etc/pacman.d/mirrorlist')

        # exec_chroot('sed', 's/#//g', '-i', '/etc/pacman.d/mirrorlist')
        # exec_chroot('bash', '-c', 'grep "^Server =" /etc/pacman.d/mirrorlist > /etc/pacman.d/mirrorlist.tmp; mv /etc/pacman.d/mirrorlist.tmp /etc/pacman.d/mirrorlist')

        with open('/.new_rootfs/etc/pacman.conf', 'r') as original:
            data = original.read()
        with open('/.new_rootfs/etc/pacman.conf', 'w') as modified:
            modified.write(data.replace(
                "[options]", "[options]\nParallelDownloads = 32\n"))
        with open('/.new_rootfs/etc/pacman.conf', 'w') as modified:
            modified.write(data.replace(
                "#[multilib]\n#Include = /etc/pacman.d/mirrorlist", "[multilib]\nInclude = /etc/pacman.d/mirrorlist"))

        with open('/.new_rootfs/etc/pacman.conf', 'a') as pacman_conf:
            pacman_conf.write(f'''
[breakfast]
SigLevel = Never
Server = {repo}
''')
            pacman_conf.write('''
[chaotic-aur]
SigLevel = Never
Server = https://cdn-mirror.chaotic.cx/$repo/$arch
''')

            if type(blend_release.get('package-repos')) == list:
                for package_repo in blend_release.get('package-repos'):
                    if (type(package_repo.get('name')) == str and
                            type(package_repo.get('repo-url')) == str):
                        pacman_conf.write(f'''
[{package_repo["name"]}]
SigLevel = Never
Server = {package_repo["repo-url"]}
''')

        counter = 0
        while True:
            return_val = exec_chroot('pacman', '-Syu', '--noconfirm')
            counter += 1
            if counter > 30:
                error('failed to download packages')
                exit(50)
            if return_val == 0:
                break

        exec('cp', '/etc/mkinitcpio.conf', '/.new_rootfs/etc/mkinitcpio.conf')

        counter = 0
        while True:
            return_val = exec_chroot('pacman', '-S', '--ask=4', *packages)
            counter += 1
            if counter > 30:
                error('failed to download packages')
                exit(50)
            if return_val == 0:
                break

        counter = 0
        if aur_packages != []:
            while True:
                exec_chroot('useradd', '-m', '-G', 'wheel',
                            '-s', '/bin/bash', 'aur')
                exec_chroot(
                    'bash', '-c', 'echo "aur ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/aur')
                return_val = exec_chroot(
                    'runuser', '-u', 'aur', '--', 'paru', '-Sy', '--noconfirm', '--needed',
                    '--noprogressbar', '--removemake', '--skipreview', '--cleanafter', '--ask=4',
                    *aur_packages)
                exec_chroot('userdel', '-r', 'aur')
                exec_chroot('rm', '-f', '/etc/sudoers.d/aur')
                counter += 1
                if counter > 30:
                    error('failed to download AUR packages')
                    exit(50)
                if return_val == 0:
                    break

        for service in services:
            if type(service) is str:
                exec_chroot('systemctl', 'enable', service)

        for user_service in user_services:
            if type(user_service) is str:
                exec_chroot('systemctl', 'enable', '--global', user_service)

        if type(blend_release.get('commands')) == list:
            for command in blend_release.get('commands'):
                if type(command) == str:
                    exec_chroot('bash', '-c', command)
                elif type(command) == list:
                    exec_chroot(*command)

        kernel_exists = False

        for f in os.listdir('/.new_rootfs/boot'):
            if f.startswith('vmlinuz'):
                kernel_exists = True
                break

        if not kernel_exists:
            error('no Linux kernel found in new system')
            error('cancelling update so as not to render the system unbootable')
            sys.exit(10)

        exec_chroot('pacman', '-S', '--noconfirm', 'shadow')

        exec_chroot('mkinitcpio', '-P')

        exec('cp', '-ax', '/etc/locale.gen', '/.new_rootfs/etc/locale.gen')
        exec_chroot('locale-gen')

        exec_chroot('rm', '-f', '/.new_rootfs/etc/resolv.conf')

        exec('cp', '-ax', '/.new_rootfs/etc', '/.new.etc')

        if os.environ.get('AKSHARA_INSTALL') == '1':
            exec('rm', '-rf', '/usr/etc')
            exec('cp', '-ax', '/.new_rootfs/etc', '/usr/etc')

        etc_diff = filecmp.dircmp('/etc/', '/usr/etc/')

        def get_diff_etc_files(dcmp):
            dir_name = dcmp.left.replace('/etc/', '/.new.etc/', 1)
            for name in dcmp.left_only:
                exec('mkdir', '-p', dir_name)
                exec('cp', '-ax', os.path.join(dcmp.left, name), dir_name)
            for name in dcmp.diff_files:
                exec('cp', '-ax', os.path.join(dcmp.left, name), dir_name)
            for sub_dcmp in dcmp.subdirs.values():
                get_diff_etc_files(sub_dcmp)

        get_diff_etc_files(etc_diff)

        exec('cp', '-ax', '/var/lib', '/.new.var.lib')

        var_lib_diff = filecmp.dircmp(
            '/.new_rootfs/var/lib/', '/.new.var.lib/')

        dir_name = '/.new.var.lib/'
        for name in var_lib_diff.left_only:
            if os.path.isdir(os.path.join(var_lib_diff.left, name)):
                exec('cp', '-ax', os.path.join(var_lib_diff.left, name), dir_name)

        exec('cp', '/etc/passwd', '/.new_rootfs/etc')
        exec('cp', '/etc/group', '/.new_rootfs/etc')
        exec('cp', '/etc/shadow', '/.new_rootfs/etc')
        exec('cp', '/etc/gshadow', '/.new_rootfs/etc')

        exec_chroot('systemd-sysusers')

        exec('cp', '/.new_rootfs/etc/passwd', '/.new.etc')
        exec('cp', '/.new_rootfs/etc/group', '/.new.etc')
        exec('cp', '/.new_rootfs/etc/shadow', '/.new.etc')
        exec('cp', '/.new_rootfs/etc/gshadow', '/.new.etc')

        exec('cp', '/.new_rootfs/etc/pacman.conf', '/.new.etc')
        exec('rm', '-rf', '/.new.etc/systemd/system')
        exec('cp', '-ax', '/.new_rootfs/etc/systemd/system', '/.new.etc/systemd')
        exec('rm', '-rf', '/.new.var.lib/pacman')
        exec('cp', '-ax', '/.new_rootfs/var/lib/pacman', '/.new.var.lib/pacman')

        exec('mv', '.new_rootfs', '.update_rootfs')
        exec('cp', '-ax', '/.update_rootfs/etc', '/.update_rootfs/usr/etc')

        new_boot_files = []

        for f in os.listdir('/.update_rootfs/boot'):
            if not os.path.isdir(f'/.update_rootfs/boot/{f}'):
                exec('mv', f'/.update_rootfs/boot/{f}', '/boot')
                new_boot_files.append(f)

        for f in os.listdir('/boot'):
            if not os.path.isdir(f'/boot/{f}'):
                if f not in new_boot_files:
                    exec('rm', '-f', f'/boot/{f}')

        exec('grub-mkconfig', '-o', '/boot/grub/grub.cfg')

        exec('touch', '/.update')

        info('downloaded update and generated new rootfs')
        info('you may reboot now')


def daemon():
    for dir in os.listdir('/'):
        if dir.startswith('.old.'):
            exec('rm', '-rf', '/' + dir)


description = f'''
{colors.bold}{colors.fg.cyan}usage:{colors.reset}
  {os.path.basename(sys.argv[0])} [command] [options] [arguments]

{colors.bold}{colors.fg.cyan}version:{colors.reset} {__version}{colors.bold}

{colors.bold}{colors.fg.cyan}available commands{colors.reset}:
  {colors.bold}help{colors.reset}                  Show this help message and exit.
  {colors.bold}update{colors.reset}                Update your blendOS system.
  {colors.bold}version{colors.reset}               Show version information and exit.

{colors.bold}{colors.fg.cyan}options for commands{colors.reset}:
  {colors.bold}-v, --version{colors.reset}         show version information and exit
'''

epilog = f'''
{colors.bold}Made with {colors.fg.red}\u2764{colors.reset}{colors.bold} by Rudra Saraswat.{colors.reset}
'''

parser = argparse.ArgumentParser(description=description, usage=argparse.SUPPRESS,
                                 epilog=epilog, formatter_class=argparse.RawTextHelpFormatter)
command_map = {'help': 'help',
               'version': 'version',
               'update': update_system,
               'daemon': daemon}
parser.add_argument('command', choices=command_map.keys(),
                    help=argparse.SUPPRESS)
parser.add_argument('--keep-files-on-error',
                    action='store_true', help="keep working files on error")
parser.add_argument('-v', '--version', action='version',
                    version=f'%(prog)s {__version}', help=argparse.SUPPRESS)

if len(sys.argv) == 1:
    parser.print_help()
    exit()

if os.geteuid() != 0 and not sys.argv[1] in ('help', 'version', '-v', '--version'):
    error('requires root')
    exit(1)

args = parser.parse_intermixed_args()

command = command_map[args.command]

try:
    if command == 'help':
        parser.print_help()
    elif command == 'version':
        parser.parse_args(['--version'])
    elif command == update_system:
        exec('touch', '/var/lib/.akshara-system-lock')
        system_lock = fasteners.InterProcessLock(
            '/var/lib/.akshara-system-lock')
        info('attempting to acquire system lock')
        with system_lock:
            command()
    else:
        command()
except:
    error('aborting')
    # remove update and akshara stuff if the program errors (either exited by the user or an error) and is updating
    if command == update_system and not args.keep_files_on_error:
        exec('umount', '-rf', '/.new_rootfs/')
        exec('rm', '-rf', '/.new_rootfs/')
        exec('rm', '-rf', '/.update_rootfs')
        exec('rm', '-f', '/.update')
    else:
        print("--keep-files-on-error specified, not deleting files (/.new_rootfs/, /.update_rootfs/, /.update)")
