#!/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 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'

# END


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


def exec_chroot(*cmd, **kwargs):
    return exec('systemd-nspawn', '-D', '/.new_rootfs', '--', *cmd, **kwargs)


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 os.path.isdir('/.new_rootfs'):
        exec('rm', '-rf', '/.new_rootfs')

    # 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)

        info('downloading Arch tarball...')

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

        if exec('bash', '-c', 'sha256sum -c --ignore-missing <(wget -qO- https://geo.mirror.pkgbuild.com/iso/latest/sha256sums.txt | sed "s/archlinux-bootstrap-x86_64\.tar\.gz/.update.tar.gz/g") 2>/dev/null') != 0:
            error('failed checksum verification')
            print()
            info('trying download again...')
            exec('rm', '-f', '/.update.tar.gz')
            if exec('wget', '-q', '--show-progress', 'https://geo.mirror.pkgbuild.com/iso/latest/archlinux-bootstrap-x86_64.tar.gz', '-O', '/.update.tar.gz') != 0:
                error('failed download')
                print()
                error('update failed')
                sys.exit(50)
                return
            if exec('bash', '-c', 'sha256sum -c --ignore-missing <(wget -qO- https://geo.mirror.pkgbuild.com/iso/latest/sha256sums.txt | sed "s/archlinux-bootstrap-x86_64\.tar\.gz/.update.tar.gz/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.gz')
        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'
        ]

        services = [
            'akshara'
        ]

        user_services = [
            'blend-files'
        ]

        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('services')) == list:
            services += blend_release.get('services')

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

        print(blend_release)

        exec_chroot('cp', '/.new_rootfs/etc/pacman.d/mirrorlist', '/.new_rootfs/etc/pacman.d/mirrorlist.bkp')

        with open('/.new_rootfs/etc/pacman.d/mirrorlist', 'w') as pacman_mirrorlist_conf:
            pacman_mirrorlist_conf.write('Server = https://cloudflaremirrors.com/archlinux/$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('mv', '/.new_rootfs/etc/pacman.d/mirrorlist.bkp', '/.new_rootfs/etc/pacman.d/mirrorlist')

        exec_chroot('reflector', '--latest', '20', '--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', 'a') as pacman_conf:
            pacman_conf.write(f'''
[breakfast]
SigLevel = Never
Server = {blend_release['repo']}
''')

            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

        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('mv', '.new_rootfs', '.update_rootfs')

        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')

        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)

    Event.wait()


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('pkg', action='store', type=str,
                    nargs='*', help=argparse.SUPPRESS)
parser.add_argument('--headless',
                    action='store_true', help=argparse.SUPPRESS)
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 KeyboardInterrupt:
    error('aborting')