#!/usr/bin/env python3
""" Rebuild AUR packages against newer dependencies
"""
import json
import fileinput
import sys
import os
import subprocess
import tempfile
import shutil

from pwd import getpwnam
from decimal import Decimal
#from pyalpm import vercmp
from srcinfo.parse import parse_srcinfo
ARGV0 = 'rebuild'

def xdg_cache_home(user=None):
    """Retrieve XDG_CACHE_HOME from the XDG Base Directory specification
    """
    if user is not None:
        user_home = os.path.expanduser("~" + user)
    else:
        user_home = os.path.expanduser("~")
    cache_home = os.path.join(user_home, '.cache')

    # Note: this only retrieves `XDG_CACHE_HOME` from the current user
    # environment regardless if `user` is specified.
    if 'XDG_CACHE_HOME' in os.environ:
        return os.getenv('XDG_CACHE_HOME')

    return cache_home


def run_readline(command, check=True, cwd=None):
    """Run the output from a command line-by-line.

    `aur` programs typically use newline delimited output. Here, this function
    is used with `aur repo` to read JSON objects, with each line representing
    one local repository package.

    """
    with subprocess.Popen(command, stdout=subprocess.PIPE, cwd=cwd) as process:
        while True:
            output = process.stdout.readline()
            if output == b'' and process.poll() is not None:
                break
            if output:
                yield output.strip()

        return_code = process.poll()
        if return_code > 0 and check:
            raise subprocess.CalledProcessError(return_code, command)


def srcinfo_get_version(srcinfo):
    """Return the full version string from a .SRCINFO file.

    The `epoch` key is optional, `pkgver` and `pkgrel` are assumed present.

    """
    with open(srcinfo, 'r', encoding='utf-8') as file:
        (data, errors) = parse_srcinfo(file.read())
        if errors:
            sys.exit(1)

        epoch  = data.get('epoch')
        pkgver = data['pkgver']
        pkgrel = data['pkgrel']

        if epoch is not None:
            return epoch + ':' + pkgver, pkgrel
        return pkgver, pkgrel


def increase_decimal(decimal_number, increment, n_digits=2):
    """Only increase the fractional part of a number.
    """
    # Convert the decimal number and increment to Decimal objects
    decimal_num = Decimal(str(decimal_number))
    inc = Decimal(str(increment))

    # Calculate the increased decimal
    increased_decimal = decimal_num + inc

    # Convert the increased decimal to a formatted string with fixed precision
    precision = '.' + str(n_digits) + 'f'
    increased_decimal_str = format(increased_decimal, precision)

    return increased_decimal_str


def update_pkgrel(buildscript, pkgrel=None, increment=0.1):
    """Update pkgrel in a PKGBUILD by a given increment.

    Modifications assume a single caller and are not thread-safe.
    """
    n_digits = sum(ch.isdigit() for ch in str(increment).strip('0'))
    new_pkgrel = None

    # Creates PKGBUILD.bak which is deleted when `finput` is closed
    with fileinput.input(buildscript, inplace=True) as finput:
        for line in finput:
            pkgrel_keyword = 'pkgrel='

            if line.startswith(pkgrel_keyword):
                # Extract and update the current pkgrel value
                if pkgrel is None:
                    pkgrel = float(line.split('=')[1])  # Only the last written pkgrel holds
                new_pkgrel = increase_decimal(pkgrel, increment, n_digits)

                # Replace the pkgrel value in the line
                line = f'{pkgrel_keyword}{new_pkgrel}\n'

            # Write the modified line to stdout (which redirects to the PKGBUILD file)
            print(line, end='')

    return new_pkgrel


# TODO: use vercmp to ensure rebuilds, abort reverse depends when depends fails (sync--ninja)
def rebuild_packages(repo_targets, db_name, start_dir, pkgver=False,
                     fail_fast=False, user=None, *build_args):
    """Rebuild a series of packages in successive order.
    """
    build_cmd  = ['aur', 'build'] + list(*build_args)
    srcver_cmd = ['aur', 'srcver']

    if db_name is not None:
        build_cmd.extend(('--database', db_name))

    if user is not None:
        srcver_cmd = ['runuser', '-u', user, '--'] + srcver_cmd

    # Check that `pkgver` is consistent between local repository and .SRCINFO
    rebuilds = {}

    for pkgname, pkg in repo_targets.items():
        # Only run once per pkgbase
        if pkgname in rebuilds:
            continue

        # Retrieve metdata from local repository entry
        pkgbase = pkg['PackageBase']
        pkgver, pkgrel = pkg['Version'].rsplit('-', 1)
        src_dir = os.path.join(start_dir, pkgbase)

        # Run pkgver() function for VCS packages
        if pkgver:
            print(f'{ARGV0}: {pkgname}: updating pkgver with aur-srcver', file=sys.stderr)
            for n, pkg_str in enumerate(run_readline(srcver_cmd, cwd=src_dir)):
                if n > 0:
                    raise RuntimeError('ambiguous aur-srcver output')
                src_pkgver, _ = pkg_str.decode('utf-8').split('\t')[1].rsplit('-', 1)

        # Use .SRCINFO for other packages (faster)
        else:
            src_pkgver, _ = srcinfo_get_version(os.path.join(src_dir, '.SRCINFO'))

        buildscript = os.path.join(src_dir, 'PKGBUILD')
        buildscript_backup = None

        # Increase subrelease level to avoid conflicts with intermediate PKGBUILD updates
        if src_pkgver == pkgver:
            # Set backup file for PKGBUILD
            buildscript_backup = buildscript + '.tmp'

            # Preserve permissions of PKGBUILD
            bst = os.stat(buildscript)
            shutil.copy2(buildscript, buildscript_backup)
            shutil.chown(buildscript_backup, user=bst.st_uid)

            new_pkgrel = update_pkgrel(buildscript, pkgrel=float(pkgrel), increment=0.1)

            # Print bumped pkgrel to standard error
            print(f'{ARGV0}: {pkgbase}: {pkgver}-{pkgrel} -> {pkgver}-{new_pkgrel}',
                  file=sys.stderr)
        else:
            print(f'{ARGV0}: {pkgbase}: source and local repository version differ', file=sys.stderr)
            print(f'{ARGV0}: {pkgbase}: using existing pkgver', file=sys.stderr)

        failed = []

        # Build package with modified pkgrel
        try:
            if user is None:
                subprocess.run(build_cmd, check=True, cwd=src_dir)
            else:
                # Drop privileges when running as root, see `examples/sync-rebuild`
                asroot_env = {
                    'AUR_ASROOT'       : '1',
                    'AUR_MAKEPKG'      : f'runuser -u {user} -- makepkg',
                    'AUR_GPG'          : f'runuser -u {user} -- gpg',
                    'AUR_REPO_ADD'     : f'runuser -u {user} -- repo-add',
                    'AUR_BUILD_PKGLIST': f'runuser -u {user} -- aur build--pkglist'
                }
                subprocess.run([*build_cmd, '--user', user], check=True, cwd=src_dir, 
                               env=dict(os.environ, **asroot_env))

            # Build process completed successfully, remove backup PKGBUILD if it
            # was created above
            if buildscript_backup is not None:
                os.remove(buildscript_backup)

        except subprocess.CalledProcessError:
            # Build process failed, revert to unmodified PKGBUILD
            if buildscript_backup is not None:
                print(f'{ARGV0}: {pkgbase}: build failed, reverting PKGBUILD', file=sys.stderr)
                os.replace(buildscript_backup, buildscript)

            # --fail-fast: if a package failed to build, consider remaining targets as failed
            if fail_fast:
                print(f'{ARGV0}: {pkgbase}: build failed, exiting', file=sys.stderr)
                # XXX: Preserve original order of inputs
                return rebuilds, list(set(repo_targets) - set(rebuilds))

            # Mark rebuild as failure for later reporting to the user
            failed.append(pkgbase)

        rebuilds[pkgname] = pkgbase

    return rebuilds, failed


def print_cached_packages(pkgnames):
    """Print cached packages in `vercmp` order.
    """
    name_args = ['--name=' + item for item in pkgnames]
    pacsift   = ['pacsift', *name_args, '--exact', '--cache']

    with subprocess.Popen(pacsift, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) as p1:
        with subprocess.Popen(['pacsort'], stdin=p1.stdout, stderr=subprocess.PIPE) as p2:
            p2.communicate()


def main(targets, db_name, start_dir, pkgver, fail_fast, run_fetch, chroot, user):
    # Ensure all sources are available. Only packages are cloned that are
    # already available in the local repository.
    # XXX: Does not retrieve or handle new dependencies.
    fetch_cmd = ['aur', 'fetch', '--existing']
    repo_cmd  = ['aur', 'repo', '--jsonl']

    if user is not None:
        fetch_cmd = ['runuser', '-u', user, '--'] + fetch_cmd
        repo_cmd  = ['runuser', '-u', user, '--'] + repo_cmd

    if db_name is not None:
        repo_cmd.extend(('--database', db_name))

    if chroot:
        build_args = ['--chroot']
    else:
        build_args = ['--syncdeps', '--rmdeps', '--noconfirm']

    repo_targets_tmp = {}

    # Read repository contents line by line to handle potentially large databases
    for pkg_str in run_readline(repo_cmd):
        pkg = json.loads(pkg_str)
        pkgname = pkg['Name']

        # Restrict to packages specified on the command-line
        if pkgname in targets:
            repo_targets_tmp[pkgname] = {
                'PackageBase': pkg['PackageBase'], 'Version': pkg['Version']
            }

    # Restore order of command-line targets
    repo_targets = {}
    for target in targets:
        if target in repo_targets_tmp:
            repo_targets[target] = repo_targets_tmp[target]

    # Clone targets that are part of the local repository
    if len(repo_targets) > 0:
        fetch_cmd.extend(list(repo_targets.keys()))

        if run_fetch:
            repo_targets_ordered = {}  # `dict` preserves order since python >=3.6

            # Temporary file for dependency order
            with tempfile.NamedTemporaryFile() as fetch_results:
                # Read access to build user
                if user is not None:
                    shutil.chown(fetch_results.name, user=user)

                # Clone AUR targets. Dependency order is taken from the command-line.
                subprocess.run([*fetch_cmd, '--results', fetch_results.name],
                               cwd=start_dir, check=True)

                # Retrieve names of cloned/fetched packages.
                with open(fetch_results.name, 'r') as f:
                    for line in f.readlines():
                        name = os.path.basename(line.split(':')[-1].rstrip())
                        repo_targets_ordered[name] = repo_targets[name]

            # Local repository targets not retrieved by `aur-fetch` are missing from AUR
            # XXX: append to queue if target directories are available
            not_aur = list(set(repo_targets.keys()) - set(repo_targets_ordered.keys()))

            # Build in dependency order
            rebuilds, failed = rebuild_packages(repo_targets_ordered, db_name, start_dir, 
                                                pkgver, fail_fast, user, build_args)
        else:
            not_aur = []

            # Build in sequential (argument) order
            rebuilds, failed = rebuild_packages(repo_targets, db_name, start_dir, 
                                                pkgver, fail_fast, user, build_args)

        if len(not_aur) > 0:
            print(f'{ARGV0}: the following targets are not in AUR:', file=sys.stderr)
            print(' '.join(not_aur), file=sys.stderr)

        if len(failed) > 0:
            print(f'{ARGV0}: the following targets failed to build:', end=' ', file=sys.stderr)
            print(' '.join(failed), file=sys.stderr)

        rest = list(set(targets) - set(rebuilds.keys()) - set(failed) - set(not_aur))
    else:
        rest = list(targets)

    if len(rest) > 0:
        print(f'{ARGV0}: the following targets are unavailable in the local repository',
              file=sys.stderr)
        print(' '.join(rest), file=sys.stderr)

        # Print any stale cached packages
        print(f'{ARGV0}: with cached entries:', file=sys.stderr)
        print_cached_packages(rest)


# Parse user arguments when run directly
if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(prog=f'{ARGV0}', description='rebuild packages')
    parser.add_argument('-d', '--database')
    parser.add_argument('-c', '--chroot', action='store_true')
    parser.add_argument('-U', '--user')
    parser.add_argument('--pkgver', action='store_true')
    parser.add_argument('--fail-fast', action='store_true')
    parser.add_argument('--no-fetch', action='store_false')
    parser.add_argument('targets', nargs='+')
    args = parser.parse_args()

    # Verify options
    if os.geteuid() == 0 and (args.user is None or getpwnam(args.user).pw_uid == 0):
        print(f'{ARGV0}: unprivileged user required (--user)', file=sys.stderr)
        sys.exit(1)

    elif os.getuid() != 0 and args.user is not None:
        print(f'{ARGV0}: --user requires root', file=sys.stderr)
        sys.exit(1)

    elif os.geteuid() == 0 and not args.chroot:
        raise NotImplementedError('--user requires --chroot')

    # Get the path to user-specific cache files
    # Note: this only retrieves `AURDEST` from the current user environment.
    if 'AURDEST' in os.environ:
        start_dir = os.getenv('AURDEST')
    else:
        start_dir = os.path.join(xdg_cache_home(args.user), 'aurutils/sync')

    main({i:1 for i in args.targets}, args.database, start_dir,
         args.pkgver, args.fail_fast, args.no_fetch, args.chroot, args.user)
