#!/usr/bin/env python3

# Author: Xyne

'''
ABOUT

  Display information about upgradable packages in Conky.

DEPENDENCIES

  * pyalpm
  * python3-aur
  * python-pyxdg

EXAMPLE CONKY CONFIGURATION FILE

  alignment top_left
  gap_x 5
  gap_y 0
  maximum_width 200
  minimum_size 200,1
  own_window yes
  own_window_transparent yes
  own_window_type override
  own_window_hints below

  update_interval 3600
  total_run_times 0
  double_buffer yes

  use_xft yes
  xftfont lime:pixelsize=10
  xftalpha 0.9

  default_color ff0000
  default_outline_color black
  default_shade_color black

  uppercase no
  override_utf8_locale no

  text_buffer_size 4096

  color1 444444
  color2 cccccc
  color3 777777


  TEXT
  ${execp /path/to/paconky --aur}
'''

import argparse
import contextlib
import io
import logging
import pathlib
import shutil
import sys
import time
import urllib.error

from collections import OrderedDict, deque

from xdg.BaseDirectory import save_cache_path

import pyalpm
from pycman import config, action_sync, transaction

try:
    import AUR.RPC as AUR
except ImportError:
    AUR = None

LOGGER = logging.getLogger(__name__)
PACMAN_CONFIG = '/etc/pacman.conf'
AUR_NAME = 'AUR'


class PaconkyError(Exception):
    '''
    Common exceptions raised by this module.
    '''


def get_all_requiredby(pkg, explicit=False):
    '''
    Generator over all packages that require the given package, directly or
    indirectly.

    Args:
        pkg:
            A pyalpm.Package from the local database.

        explicit:
            If True, limit returned packages to those that were explicitly
            installed.
    '''
    ldb = pkg.db
    seen = set()
    seen.add(pkg.name)
    queue = deque(pkg.compute_requiredby())
    while queue:
        name = queue.popleft()
        if name in seen:
            continue
        seen.add(name)
        pkg = ldb.get_pkg(name)
        if not explicit or pkg.reason == 0:
            yield name
        queue.extend(pkg.compute_requiredby())


class Paconky():
    '''
    Print lists of upgradable packages for Conky.
    '''
    def __init__(self, config_path=PACMAN_CONFIG, cache_dir=None, aur=False):
        self.config_path = config_path
        self.config = config.init_with_config(config_path)
        if cache_dir is None:
            cache_dir = save_cache_path('paconky')
        self.cache_dir = pathlib.Path(cache_dir)
        self.aur = aur
        self.initialize_cache()

    def initialize_cache(self):
        '''
        Symlink the local database to the cache directory. The local database is
        only read, never written.
        '''
        cache_dir = self.cache_dir
        db_dir = pathlib.Path(self.config.dbpath)

        # Symlink local database for upgrade determination.
        local_db_path = db_dir / 'local'
        local_db_symlink = cache_dir / 'local'
        # Check if the symlink or equivalent already exists.
        if local_db_path.resolve() != local_db_symlink.resolve():
            local_db_symlink.parent.mkdir(parents=True, exist_ok=True)
            if local_db_symlink.is_symlink():
                LOGGER.warning('unlinking existing symlink at %s', local_db_symlink)
                local_db_symlink.unlink()
            elif local_db_symlink.exists():
                LOGGER.error('refusing to unlink existing non-symlink path at %s', local_db_symlink)
                return
            local_db_symlink.symlink_to(local_db_path)

        # Copy current sync databases.
        sync_db_path = db_dir / 'sync'
        sync_db_copy = cache_dir / 'sync'
        shutil.copytree(sync_db_path, sync_db_copy, dirs_exist_ok=True)

    def download_sync_dbs(self):
        '''
        Download sync databases to cache directory.
        '''
        args = action_sync.parse_options((
            '--config', str(self.config_path),
            '-b', str(self.cache_dir),
            '-y'
        ))
        self.config = conf = config.init_with_config_and_options(args)
        with contextlib.redirect_stdout(io.StringIO()):
            try:
                trans = transaction.init_from_options(conf, args)
            except pyalpm.error as err:
                LOGGER.error('failed to sync databases to temporary directory: %s', err)
                if err.args[1] == 10:
                    LOGGER.error('lockpath exists: %s', self.cache_dir / 'db.lck')
                raise PaconkyError(f'failed to update databases: {err}') from err

            for sdb in conf.get_syncdbs():
                try:
                    sdb.update(False)
                except pyalpm.error as err:
                    LOGGER.error('failed to update repo %s: %s', sdb.name, err)
            trans.release()

    def get_upgradable_repo_pkgs(self, installed):
        '''
        Get a dictionary that maps repos to sets of upgradable packages and a
        set of foreign packages.

        Args:
            installed:
                A set of local database packages. This will be modified in
                place.

        Returns:
            A 2-tuple of the upgradable packages as a dictionary mapping repo
            names to sets of local and sync package pairs.
        '''
        conf = self.config
        installed = set(conf.get_localdb().pkgcache)
        upgradable = OrderedDict()
        for sdb in conf.get_syncdbs():
            for pkg in list(installed):
                syncpkg = sdb.get_pkg(pkg.name)
                if syncpkg:
                    if pyalpm.vercmp(syncpkg.version, pkg.version) > 0:
                        upgradable.setdefault(sdb.name, list()).append((pkg, syncpkg))
                    installed.remove(pkg)

        return upgradable

    @staticmethod
    def get_upgradable_aur_pkgs(installed):
        '''
        '''
        if AUR is None:
            LOGGER.warning('AUR.RPC is not installed: unable to check AUR for upgrades')
            return None

        # All remaining packages in the installed set are foreign.
        foreign = dict(((pkg.name, pkg) for pkg in installed))

        upgradable_aur = list()
        aur = AUR.AurRpc()
        try:
            aur_pkgs = aur.info(foreign.keys())
        except (AUR.AurError, urllib.error) as err:
            LOGGER.error('failed to retrieve AUR info: %s', err)
            return None

        for aur_pkg in aur_pkgs:
            try:
                installed_pkg = foreign[aur_pkg['Name']]
            except KeyError:
                upgradable_aur.append((None, aur_pkg))
                continue
            if pyalpm.vercmp(aur_pkg['Version'], installed_pkg.version) > 0:
                upgradable_aur.append((installed_pkg, aur_pkg))
            installed.remove(installed_pkg)

        return upgradable_aur

    def print_lists(self):
        '''
        Print the lists.
        '''
        conf = self.config
        installed = set(conf.get_localdb().pkgcache)
        upgradable_repo = self.get_upgradable_repo_pkgs(installed)
        upgradable_aur = self.get_upgradable_aur_pkgs(installed) if self.aur else None
        printer = ListPrinter(upgradable_repo, upgradable_aur)
        printer.print_lists()


class ListPrinter():
    '''
    Generate formatted lists of upgradable packages for Conky.
    '''
    HEADER_FMT = '${{color1}}${{hr}}\n${{color1}}[${{color2}}{name}${{color1}}] ${{alignr}}{status}'
    ERROR_MSG = '${color3}update check failed'
    COUNT_ZERO = '${color3}updated'
    COUNT_ONE = '${color2}1${color3} new package'
    COUNT_MANY_FMT = '${{color2}}{:d}${{color3}} new packages'
    # Text shown instead of repo when everything is up-to-date.
    TEXT_ZERO = 'local packages'
    FOOTER = '${color1}${hr}'
    ROW_FMT = '{name}{tag} ${{alignr}}{version}'
    TIME_FMT = '%Y-%m-%d %H:%M'

    def __init__(self, repo_pkgs, aur_pkgs, max_lines=80):
        self.repo_pkgs = repo_pkgs
        self.aur_pkgs = aur_pkgs
        self.max_lines = max_lines

    def print_list(self, name, count, row_data=None, max_rows=None):
        '''
        Print a list of upgradable package.

        Args:
            name:
                The name of this list (e.g. the repo name).

            count:
                The number of upgradable packages, which may be different from
                the number of lines when the list of packages is too long for
                the display.

            row_data:
                Dictionaries with keyword values for the row format string and
                the local database pkg.
        '''
        if count is None:
            status = self.ERROR_MSG
        elif count > 1:
            status = self.COUNT_MANY_FMT.format(count)
        elif count == 1:
            status = self.COUNT_ONE
        else:
            status = self.COUNT_ZERO

        print(self.HEADER_FMT.format(name=name, status=status))
        if row_data:
            row_data = sorted(row_data, key=lambda x: x['date'], reverse=True)
            row_fmt = self.ROW_FMT
            too_long = max_rows and count > max_rows
            if too_long:
                row_data = row_data[:max_rows - 1]
            for entry in row_data:
                is_orphan = not any(get_all_requiredby(entry['lpkg'], explicit=True))
                entry['tag'] = ' (orphan)' if is_orphan else ''
                print(row_fmt.format(**entry))
            if too_long:
                print('...')
        print(self.FOOTER)
        print('\n')

    def partition_available_lines(self):
        '''
        Partition available lines among the lists to display.
        '''
        pkgs_per_repo = OrderedDict()
        repo_pkgs = self.repo_pkgs
        aur_pkgs = self.aur_pkgs

        if repo_pkgs:
            for repo, pkgs in self.repo_pkgs.items():
                if not pkgs:
                    continue
                pkgs_per_repo[repo] = len(pkgs)
        if aur_pkgs:
            pkgs_per_repo[AUR_NAME] = len(aur_pkgs)

        n_repos = len(pkgs_per_repo)
        extra_lines_per_repo = len(self.HEADER_FMT.splitlines()) + len(self.FOOTER.splitlines())
        lines_between_repos = n_repos - 1

        available_lines = self.max_lines - (extra_lines_per_repo * n_repos + lines_between_repos)
        n_pkgs = sum(pkgs_per_repo.values())

        if n_pkgs <= available_lines:
            return pkgs_per_repo

        # Dumb but simple.
        repos = deque(pkgs_per_repo)
        lines_per_repo = dict((repo, 0) for repo in repos)
        while repos:
            repo = repos.popleft()
            lines_per_repo[repo] += 1
            available_lines -= 1
            if lines_per_repo[repo] < pkgs_per_repo[repo]:
                repos.append(repo)
        return lines_per_repo

    def format_date(self, pkg):
        '''
        Format the update date.
        '''
        timestamp = None
        if isinstance(pkg, pyalpm.Package):
            timestamp = pkg.builddate
        if isinstance(pkg, dict):
            try:
                timestamp = pkg['LastModified']
            except KeyError:
                pass

        if timestamp:
            return time.strftime(self.TIME_FMT, time.localtime(timestamp))
        return 'N/A'

    def print_lists(self):
        '''
        Print all lists.
        '''
        repo_pkgs = self.repo_pkgs
        aur_pkgs = self.aur_pkgs

        # Nothing to upgrade.
        if not (repo_pkgs or aur_pkgs):
            self.print_list(self.TEXT_ZERO, self.COUNT_ZERO)
            return

        lines_per_repo = self.partition_available_lines()
        if repo_pkgs:
            for repo, pkgs in repo_pkgs.items():
                row_data = tuple(
                    dict(
                        name=spkg.name,
                        version=spkg.version,
                        date=self.format_date(spkg),
                        lpkg=lpkg
                    ) for (lpkg, spkg) in pkgs
                )
                self.print_list(
                    repo,
                    len(pkgs),
                    row_data=row_data,
                    max_rows=lines_per_repo[repo]
                )

        if aur_pkgs:
            row_data = tuple(
                dict(
                    name=aurpkg['Name'],
                    version=aurpkg['Version'],
                    date=self.format_date(aurpkg),
                    lpkg=lpkg
                ) for (lpkg, aurpkg) in aur_pkgs
            )
            self.print_list(
                AUR_NAME,
                len(aur_pkgs),
                row_data=row_data,
                max_rows=lines_per_repo[AUR_NAME]
            )


def main(args=None):
    '''
    Main.
    '''
    parser = argparse.ArgumentParser(
        description='Print lists of upgradable Pacman packages for Conky.'
    )
    parser.add_argument(
        '--config', default=PACMAN_CONFIG,
        help='Pacman configuration file. Default: %(default)s'
    )
    parser.add_argument(
        '--cache',
        help='Cache directory for determining upgradable packages.'
    )
    parser.add_argument(
        '--aur', action='store_true',
        help='Check the AUR for available upgrades.'
    )
    pargs = parser.parse_args(args=args)

    paconky = Paconky(
        config_path=pargs.config,
        cache_dir=pargs.cache,
        aur=pargs.aur
    )
    #  paconky.download_sync_dbs()
    paconky.print_lists()


if __name__ == "__main__":
    logging.basicConfig(
        style='{',
        format='[{asctime}] {levelname} {message}',
        datefmt='%Y-%m-%d %H:%M:%S',
        level=logging.WARNING
    )
    try:
        main()
    except (KeyboardInterrupt, BrokenPipeError):
        pass
    except PaconkyError as err:
        sys.exit(str(err))