#!/usr/bin/env python3
# -*- coding: utf8 -*-

# Copyright (C) 2012-2021 Xyne
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# (version 2) as published by the Free Software Foundation.
#
#
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

# TODO
# Refactor handling of signature options.

import argparse
import collections
import logging
import os.path
import re
import sys
import xml.dom.minidom

import pyalpm

import XCGF
import XCPF
import XCPF.ArchPkg
import XCPF.PacmanConfig


# Optional AUR support.
try:
    import AUR.RPC
    import AUR.AurPkg
    AUR_SUPPORT = True
    DEFAULT_TTL = AUR.common.DEFAULT_TTL
except ImportError:
    AUR_SUPPORT = False
    DEFAULT_TTL = 900


# Optional Reflector support.
try:
    import Reflector
    REFLECTOR = Reflector
except ImportError:
    REFLECTOR = None


PACMAN_OPTIONS = set((
    '--arch',
    '-d', '--nodeps',
    '--debug',
    '--ignore',
    '--ignoregroup',
    '--needed',
    '-u', '--sysupgrade',
    '-v', '--verbose',
    '-y', '--refresh',
    '--noconfirm',
))

PREFERENCE_MAX = 100
PREFERENCE_MIN = 1


# --------------------------- Temporary Workaround --------------------------- #

def find_grp_pkgs(sync_dbs, name):
    '''
    Temporary function to work around segfault in pyalpm.find_grp_pkgs.
    '''
    print('looking for', name)
    for db in sync_dbs:
        for pkg in db.pkgcache:
            for group in pkg.groups:
                if group == name:
                    yield pkg


# -------------------------------- Functions --------------------------------- #

def determine_official_pkgs(pkgs):
    '''
    Determine which packages are in the official repositories and can therefore
    be downloaded more efficiently with Reflector.
    '''
    global REFLECTOR
    if not REFLECTOR:
        return XCPF.ArchPkg.PyalpmPkgSet(), pkgs
    else:
        official = XCPF.ArchPkg.PyalpmPkgSet()
        other = XCPF.ArchPkg.PyalpmPkgSet()
        for pkg in pkgs:
            if pkg.db.name in REFLECTOR.MirrorStatus.REPOSITORIES:
                official.add(pkg)
            else:
                other.add(pkg)
        return official, other


def iterate_over_all_deps(pkgs):
    '''
    Iterate over all dependencies of the given BuildablePkgs.
    '''
    for pkg in pkgs:
        for d in pkg.alldeps():
            yield d


def omit_cached_pkgs(conf, pkgs):
    for pkg in pkgs:
        for cache in conf.options['CacheDir']:
            fpath = os.path.join(cache, pkg.filename)
            for checksum in ('sha256', 'md5'):
                a = XCGF.get_checksum(fpath, checksum)
                b = getattr(pkg, checksum + 'sum')
                if a is None or a != b:
                    yield pkg
                    break
            else:
                continue
            break


# ---------------------------------- Pm2ml ----------------------------------- #

class Pm2ml(object):

    def __init__(
        self,
        pargs,
        pacman_conf=None,
        aur=None,
        ttl=DEFAULT_TTL
    ):
        self.check_aur = False
        self.aur_only = False
        self.pargs = pargs

        self.load_pacman_conf(pargs.conf, pacman_conf=pacman_conf)
        self.configure_from_arguments(pargs)

        if aur is None:
            global AUR_SUPPORT
            if AUR_SUPPORT:
                aur = AUR.RPC.AurRpc()
        self.aur = aur

        if pargs.arch:
            self.arch = pargs.arch
        else:
            self.arch = self.pacman_conf.options['Architecture']
        self.create_factories(ttl=ttl)

    def load_pacman_conf(self, path, pacman_conf=None):
        '''
        Get a configuration object and pyalpm handle using the given arguments.
        '''
        if pacman_conf is None:
            pacman_conf = XCPF.PacmanConfig.PacmanConfig(conf=path)
        self.pacman_conf = pacman_conf
        self.initialize_alpm()

    def initialize_alpm(self):
        '''
        Initialize the ALPM database.
        '''
        self.handle = self.pacman_conf.initialize_alpm()

    def refresh_databases(self, opi=True, aur=True):
        '''
        Clear cached data from the database.
        '''
        if opi:
            self.obpf.opi.mdb.db_clean(wipe=True)
        if aur and self.aur:
            self.aur.mdb.db_clean(wipe=True)

    def create_factories(self, ttl=DEFAULT_TTL):
        '''
        Create internal BuildablePkgFactory instances.
        '''
        self.obpf = XCPF.ArchPkg.OfficialBuildablePkgFactory(
            self.arch,
            config=self.pacman_conf,
            handle=self.handle,
            ttl=ttl
        )
        if AUR_SUPPORT:
            self.abpf = AUR.AurPkg.AurBuildablePkgFactory(self.arch, ttl=ttl)
        else:
            self.abpf = None

    def search_aur(self, names):
        '''
        Search the AUR for pkgnames.
        '''
        if self.aur is not None:
            for pkg in AUR.RPC.insert_full_urls(self.aur.info(names)):
                yield pkg

    def configure_from_arguments(self, pargs=None):
        '''
        Configure settings from command-line arguments.
        '''
        if pargs is None:
            pargs = self.pargs
        if AUR_SUPPORT:
            if pargs.aur_only:
                pargs.aur = True
            self.check_aur = pargs.aur
            self.aur_only = pargs.aur_only

    def satisfy_deps(
        self,
        sync_pkgcache,
        target_pkgs,
        local_pkgcache=None,
        ignored=None,
        ignoredgroups=None
    ):
        '''
        Return sync packages required to satisfy the dependencies of the given
        packages. If the local package cache is given then dependencies satisfied by
        it will not be returned.
        '''
        queue = collections.deque(target_pkgs)
        seen = set(queue)
        found = XCPF.ArchPkg.PyalpmPkgSet()
        missing = set()

        if ignoredgroups:
            ignoredgroups = set(ignoredgroups)
        else:
            ignoredgroups = set()

        while queue:
            pkg = queue.popleft()
            for dep in pkg.depends:
                if local_pkgcache is None \
                        or pyalpm.find_satisfier(local_pkgcache, dep) is None:
                    prov = pyalpm.find_satisfier(sync_pkgcache, dep)
                    if prov is not None:
                        if (ignored and prov.name in ignored) \
                                or (set(prov.groups) & ignoredgroups):
                            continue
                        found.add(prov)
                        if prov.name not in seen:
                            seen.add(prov.name)
                            queue.append(prov)
                    else:
                        missing.add(dep)
        return found, missing

    def determine_upgradable(self, aur_names=None, downgrades=False):
        '''
        Determine upgradable packages in the repos and the AUR.

        aur_names: Additional AUR pkgnames to query when invoking the AUR RPC.
        '''

        sync_pkgs = XCPF.ArchPkg.PyalpmPkgSet()
        foreign_pkgs = XCPF.ArchPkg.PyalpmPkgSet()
        aur_pkgs = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()
        if not aur_names:
            aur_names = set()

        for pkg in self.handle.get_localdb().pkgcache:
            # Use low-level pyalpm function to offload checking for "replaces", etc.
            syncpkg = pyalpm.sync_newversion(pkg, self.handle.get_syncdbs())
            if not syncpkg and downgrades:
                syncpkg = XCPF.find_pkgname(self.handle.get_syncdbs(), pkg.name)
                if syncpkg and syncpkg.version == pkg.version:
                    syncpkg = None
            if syncpkg:
                if not self.aur_only:
                    sync_pkgs.add(syncpkg)
                continue
            else:
                # Even if no upgrade is available, we must determine if the package is
                # available in the repos to avoid attempting to update it from the AUR.
                for db in self.handle.get_syncdbs():
                    if db.get_pkg(pkg.name):
                        break
                else:
                    if self.check_aur:
                        foreign_pkgs.add(pkg)

        if self.aur is not None and self.check_aur and foreign_pkgs:
            found_aur_pkgs = self.search_aur(aur_names | set(foreign_pkgs.pkgs))

            for aur_pkg in found_aur_pkgs:
                # Append specifically requested AUR packages.
                if aur_pkg['Name'] in aur_names:
                    aur_pkgs.add(aur_pkg)
                    continue
                installed_pkg = foreign_pkgs.pkgs[aur_pkg['Name']]
                if pyalpm.vercmp(aur_pkg['Version'], installed_pkg.version) > 0 \
                        or (downgrades and aur_pkg['Version'] != installed_pkg.version):
                    # Append upgradable AUR packages.
                    aur_pkgs.add(aur_pkg)
                del foreign_pkgs.pkgs[aur_pkg['Name']]

        return sync_pkgs, aur_pkgs, foreign_pkgs

    def partition_pkgs(self, names, select=False, ignored=None):
        '''
        Partition a list of packages by their availability in the repos, the AUR or
        an unknown source.

        Groups and version requirements are handled.

        select: Prompt selection of packages.
        ignored: Packages to ignore.

        Returns the following:

          * The parsed version requirements as a map of names to a list of tuples.
            Each tuple contains the relationship and the version, as strings.
          * PyalpmPkgSet of sync packages satisfying the version requirements,
          * AurPkgSet of AUR packages satisfying the version requirements.
          * A set of names that could not be either found or satisfied.

        '''
        sync_dbs = self.handle.get_syncdbs()

        sync_pkgs = XCPF.ArchPkg.PyalpmPkgSet()
        # Use a basic PkgSet to avoid checks for None later.
        aur_pkgs = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()
        foreign = set()
        unsatisfied = set()

        ver_reqs = XCPF.collect_version_requirements(names)

        for qualified_name in set(ver_reqs):
            repo, name = XCPF.repo_and_package(qualified_name)
            if repo == 'AUR':
                foreign.add(qualified_name)
                continue

            reqs = list(XCPF.get_required_version_strings(ver_reqs, (qualified_name,)))
            pkg = XCPF.select_name_or_first(
                name,
                XCPF.find_satisfiers_in_dbs(sync_dbs, reqs)
            )
            if pkg:
                sync_pkgs.add(pkg)
            else:
                # Even if the version requirements are not satisfied, we must check for a
                # package with the given name. If one exists then it precludes a group or
                # an AUR package of the same name.
                for db in sync_dbs:
                    if repo is not None and db.name != repo:
                        continue
                    elif db.get_pkg(name):
                        unsatisfied.add(qualified_name)
                        break
                else:
                    # If there is a repo prefix or if there are version requirements then
                    # the name can be assumed to be a package and not a group.
                    if repo is None and not ver_reqs[qualified_name]:
                        # TODO: restore this once the error has been fixed
                        syncgrp = XCPF.ArchPkg.PyalpmPkgSet(pyalpm.find_grp_pkgs(sync_dbs, name))
                        #  syncgrp = XCPF.ArchPkg.PyalpmPkgSet(find_grp_pkgs(sync_dbs, name))
                        if syncgrp:
                            if select:
                                selected = XCPF.PacmanConfig.select_grp(name, syncgrp, ignored)
                                sync_pkgs |= selected
                            else:
                                sync_pkgs |= syncgrp
                        else:
                            # No such group so it must be a foreign package.
                            foreign.add(qualified_name)
                    else:
                        foreign.add(qualified_name)

        if self.aur is not None:
            # Perlish
            maybe_aur = dict(
                (name, repo) for repo, name in (XCPF.repo_and_package(f) for f in foreign)
                if repo is None or repo == 'AUR'
            )
            aur_pkgs.addmany(
                self.aur.info(XCPF.get_required_version_strings(ver_reqs, maybe_aur))
            )
            foreign -= set(
                '{}/{}'.format(repo, name) if repo else name for name, repo in (
                    (a, maybe_aur[a]) for a in aur_pkgs.pkgs
                )
            )
        return ver_reqs, sync_pkgs, aur_pkgs, (unsatisfied | foreign)

    def iterate_over_all_uninstalled_deps(self, pkgs):
        '''
        Iterate over all uninstalled dependencies for the given buildable packages.
        '''
        local_db = self.handle.get_localdb()
        for d in iterate_over_all_deps(pkgs):
            if not pyalpm.find_satisfier(local_db.pkgcache, d):
                yield d

    def resolve_aur_dependencies(
        self,
        pkgs,
        recursive=True,
        include_installed=False,
        ignored=None,
        ignoredgroups=None
    ):
        '''
        Resolve dependencies for the given AUR packages.

        Returns the same tuple as partition_pkgs without the version requirements.
        '''

#     sync_deps = XCPF.ArchPkg.BuildablePkgSet()
#     aur_deps = XCPF.ArchPkg.BuildablePkgSet()
        sync_deps = XCPF.ArchPkg.PyalpmPkgSet()
        aur_deps = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()
        unknown_deps = set()

        if include_installed:
            iterate = iterate_over_all_deps
        else:
            def iterate(ps):
                return self.iterate_over_all_uninstalled_deps(ps)

        deps = set(iterate(self.buildable_pkgs(pkgs)))
        seen = set()

        while deps:
            seen |= deps
            ver_reqs, sync_pkgs, aur_pkgs, unknown_pkgs = self.partition_pkgs(deps)

            sync_pkgs.ignore(ignored, ignoredgroups)
            aur_pkgs.ignore(ignored, ignoredgroups)

            sync_deps |= sync_pkgs
            aur_deps |= aur_pkgs
            unknown_deps |= unknown_pkgs

            if recursive:
                # Wrap in the BuildablePkg class to ensure proper dependency resolution.
                # TODO
                # Remove the wrapper when the AUR RPC interface is updated to include
                # the missing information.
                deps = set(d for d in iterate(self.buildable_pkgs(aur_pkgs)) if d not in seen)
            else:
                deps = None

        return sync_deps, aur_deps, unknown_deps

    def buildable_pkgs(self, pkgs):
        '''
        Wrap packages in BuildablePkgs. This only support alpm.Package and AUR.RPC
        dictionaries.
        '''
        for pkg in pkgs:
            if isinstance(pkg, pyalpm.Package):
                yield self.obpf.pkg(pkg)
            # No check is needed here because if a non-ALPM packages is passed then
            # AUR support is necessarily enabled (or someone is doing something wrong.
            else:
                yield self.abpf.pkg(pkg)

    def resolve_targets(
        self,
        targets,
        upgrade=0,
        deps=True,
        all_deps=False,
        select=False,
        ignored=None,
        ignoredgroups=None,
        ignore_ignored=False,
        only_needed=False,
    ):
        '''
        Determine the target sync and AUR packages along with their dependencies.
        '''
        if all_deps:
            deps = True

        if ignore_ignored:
            ignored = set()
            ignoredgroups = set()
        else:
            if ignored:
                ignored = set(ignored + self.pacman_conf.options['IgnorePkg'])
            else:
                ignored = set(self.pacman_conf.options['IgnorePkg'])
            if ignoredgroups:
                ignoredgroups = set(ignoredgroups + self.pacman_conf.options['IgnoreGroup'])
            else:
                ignoredgroups = set(self.pacman_conf.options['IgnoreGroup'])

        sync_pkgcache = XCPF.join_db_pkgcaches(self.handle.get_syncdbs())

        ver_reqs, sync_pkgs, aur_pkgs, not_found = self.partition_pkgs(
            targets,
            select=select,
            ignored=ignored
        )

        if upgrade:
            a, b, orphans = self.determine_upgradable(downgrades=(upgrade > 1))
            sync_pkgs |= a.ignore(ignored, ignoredgroups)
            aur_pkgs |= b.ignore(ignored, ignoredgroups)
        else:
            orphans = set()

        if deps:
            if self.aur is not None:
                sync_deps, aur_deps, unknown_deps = self.resolve_aur_dependencies(
                    aur_pkgs,
                    recursive=True,
                    include_installed=all_deps,
                    ignored=ignored,
                    ignoredgroups=ignoredgroups,
                )
            else:
                sync_deps = XCPF.ArchPkg.PyalpmPkgSet()
                aur_deps = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()
                unknown_deps = set()

            if all_deps:
                local_pkgcache = None
            else:
                local_pkgcache = self.handle.get_localdb().pkgcache

            more_sync_deps, missing_deps = self.satisfy_deps(
                sync_pkgcache,
                sync_pkgs | sync_deps,
                local_pkgcache=local_pkgcache,
                ignored=ignored,
                ignoredgroups=ignoredgroups,
            )

            sync_deps |= more_sync_deps
            unknown_deps |= missing_deps

        else:
            sync_deps = XCPF.ArchPkg.PyalpmPkgSet()
            aur_deps = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()
            unknown_deps = set()

    #   not_found |= set(ver_reqs) - (set(sync_pkgs.pkgs) | set(aur_pkgs.pkgs))

        if only_needed:
            for p in self.handle.get_localdb().pkgcache:
                for x in (sync_pkgs, sync_deps, aur_pkgs, aur_deps):
                    x.remove_version(p.name, p.version)

        return \
            sync_pkgs, sync_deps, \
            aur_pkgs, aur_deps, \
            not_found, unknown_deps, orphans

    def resolve_targets_from_arguments(self, pargs=None, ignore_ignored=False):
        '''
        Convenience wrapper around resolve_targets to fill in parameters from
        arguments.
        '''
        if pargs is None:
            pargs = self.pargs
        return self.resolve_targets(
            pargs.pkgs,
            upgrade=pargs.sysupgrade,
            deps=(not pargs.nodeps),
            all_deps=pargs.all_deps,
            select=(pargs.select and not pargs.noconfirm),
            ignored=pargs.ignore,
            ignoredgroups=pargs.ignoregroup,
            ignore_ignored=ignore_ignored,
            only_needed=pargs.needed,
        )

    def build_download_queue(self, pargs=None, sync_pkgs=None, aur_pkgs=None):
        '''
        Build a download queue that can be passed to download_queue_to_metalink.
        '''
        if pargs is None:
            pargs = self.pargs
        if sync_pkgs is None:
            sync_pkgs = XCPF.ArchPkg.PyalpmPkgSet()
        if aur_pkgs is None:
            aur_pkgs = AUR.AurPkg.AurPkgSet() if AUR_SUPPORT else XCPF.ArchPkg.PlaceholderPkgSet()

        download_queue = DownloadQueue()

        if pargs.databases:
            for db in self.handle.get_syncdbs():
                try:
                    siglevel = self.pacman_conf[db.name]['SigLevel'].split()[0]
                except KeyError:
                    siglevel = None
                download_sig = XCPF.PacmanConfig.needs_sig(siglevel, pargs.sigs, 'Database')
                download_queue.add_db(db, sig=download_sig, files=pargs.files)

        official, other = determine_official_pkgs(sync_pkgs)

        if official:
            global REFLECTOR
            mirrors = list()
            if pargs.reflector and REFLECTOR:
                reflector_options = REFLECTOR.parse_args(pargs.reflector)
                try:
                    mirrorstatus, mirrors = REFLECTOR.process_options(reflector_options)
                except REFLECTOR.MirrorStatusError as e:
                    logging.error(str(e))
                else:
                    mirrors = list(m['url'] for m in mirrors)
            for pkg in official:
                # This preserves the order of the user's mirrorlist files and prevents
                # duplicates from being added by Reflector.
                urls = list(os.path.join(url, pkg.filename) for url in pkg.db.servers)
                urls.extend(
                    set(
                        os.path.join(REFLECTOR.MirrorStatus.MIRROR_URL_FORMAT.format(
                            m, pkg.db.name, pkg.arch), pkg.filename
                        ) for m in mirrors
                    ) - set(urls)
                )

                try:
                    siglevel = self.pacman_conf[pkg.db.name]['SigLevel'].split()[0]
                except KeyError:
                    siglevel = None
                download_sig = XCPF.PacmanConfig.needs_sig(siglevel, pargs.sigs, 'Package')
                download_queue.add_sync_pkg(pkg, urls, download_sig)

        for pkg in other:
            try:
                siglevel = self.pacman_conf[pkg.db.name]['SigLevel'].split()[0]
            except KeyError:
                siglevel = None
            download_sig = XCPF.PacmanConfig.needs_sig(siglevel, pargs.sigs, 'Package')
            urls = set(os.path.join(url, pkg.filename) for url in pkg.db.servers)
            download_queue.add_sync_pkg(pkg, urls, download_sig)

        for pkg in aur_pkgs:
            download_queue.add_aur_pkg(pkg)

        return download_queue


# --------------------------------- Metalink --------------------------------- #

class Metalink(object):
    '''
    Metalink generation class.
    '''

    def __init__(self, impl, output_dir=None, set_preference=False):
        self.doc = impl.createDocument(None, "metalink", None)
        self.doc.documentElement.setAttribute('xmlns', "urn:ietf:params:xml:ns:metalink")
        self.files = self.doc.documentElement
        self.set_output_dir(output_dir)
        self.set_preference = set_preference

        # Map of pkgnames to tuples of (<pkg>, <url list>, <sig boolean>)
        self.sync_pkgs = dict()
        # Map of names to (<db>, <sig boolean>).
        self.dbs = dict()
        # Map of pkgbase to AUR packages.
        self.aur_pkgs = dict()

        self.stale = False

    def add_sync_pkg(self, pkg, urls, sig=False):
        self.sync_pkgs[pkg.name] = (pkg, urls, sig)
        self.stale = True

    def add_db(self, db, sig=False, files=False):
        self.dbs[db.name] = (db, sig, files)
        self.stale = True

    def add_aur_pkg(self, pkg):
        try:
            self.aur_pkgs[pkg['PackageBase']].append(pkg)
        except KeyError:
            self.aur_pkgs[pkg['PackageBase']] = [pkg]
        self.stale = True

    def clear(self):
        self.sync_pkgs.clear()
        self.dbs.clear()
        self.aur_pkgs.clear()
        self.files.childNodes.clear()
        self.stale = False

    def build_metalink(self, clear=False):
        if clear:
            self.files.childNodes.clear()
        for name in sorted(self.sync_pkgs):
            self.insert_sync_pkg(*self.sync_pkgs[name])
        for name in sorted(self.dbs):
            self.insert_db(*self.dbs[name])
        for pkgbase in sorted(self.aur_pkgs):
            self.insert_aur_pkg(self.aur_pkgs[pkgbase][0])
        self.stale = False

    def __str__(self):
        if self.stale:
            self.build_metalink(clear=True)
        return re.sub(
            r'(?<=>)\n\s*([^\s<].*?)\s*\n\s*',
            r'\1',
            self.doc.toprettyxml(indent=' ')
        )

    def set_output_dir(self, output_dir=None):
        if not output_dir:
            self.output_dir = None
        else:
            self.output_dir = os.path.relpath(output_dir)

    def insert_urls(self, element, urls):
        '''Add URL elements to the given element.'''
        if self.set_preference:
            # This is necessary to get the length if `urls` is a generator.
            if not isinstance(urls, list):
                urls = list(urls)
            p = float(PREFERENCE_MAX)
            n = len(urls)
            dp = max((p / n), 1.0)

        for url in urls:
            url_tag = self.doc.createElement('url')
            element.appendChild(url_tag)
            if self.set_preference:
                url_tag.setAttribute('preference', '{:0.0f}'.format(p))
                p = max((p - dp), PREFERENCE_MIN)
            url_val = self.doc.createTextNode(url)
            url_tag.appendChild(url_val)

    def insert_sync_pkg(self, pkg, urls, sig=False):
        '''Insert a sync db package.'''
        file_ = self.doc.createElement("file")
        if self.output_dir:
            file_.setAttribute("name", os.path.join(self.output_dir, pkg.filename))
        else:
            file_.setAttribute("name", pkg.filename)
        self.files.appendChild(file_)
        for tag, db_attr, attrs in (
            ('identity', 'name', ()),
            ('size', 'size', ()),
            ('version', 'version', ()),
            ('description', 'desc', ()),
            ('hash', 'sha256sum', (('type', 'sha256'),)),
            ('hash', 'md5sum', (('type', 'md5'),))
        ):
            tag = self.doc.createElement(tag)
            file_.appendChild(tag)
            val = self.doc.createTextNode(str(getattr(pkg, db_attr)))
            tag.appendChild(val)
            for key, val in attrs:
                tag.setAttribute(key, val)
        urls = list(urls)
        self.insert_urls(file_, urls)
        if sig:
            self.insert_file(pkg.filename + '.sig', (u + '.sig' for u in urls))

    def insert_file(self, name, urls):
        '''Insert a signature file.'''
        file_ = self.doc.createElement("file")
        if self.output_dir:
            file_.setAttribute("name", os.path.join(self.output_dir, name))
        else:
            file_.setAttribute("name", name)
        self.files.appendChild(file_)
        self.insert_urls(file_, urls)

    def insert_db(self, db, sig=False, files=False):
        '''Insert a sync db.'''
        file_ = self.doc.createElement("file")
        if files:
            db_ext = '.files'
        else:
            db_ext = '.db'
        name = db.name + db_ext
        if self.output_dir:
            file_.setAttribute("name", os.path.join(self.output_dir, name))
        else:
            file_.setAttribute("name", name)
        self.files.appendChild(file_)
        urls = list(os.path.join(url, db.name + db_ext) for url in db.servers)
        self.insert_urls(file_, urls)
        if sig:
            self.insert_file(name + '.sig', (u + '.sig' for u in urls))

    def insert_aur_pkg(self, pkg):
        '''Insert an AUR pkg.'''
        file_ = self.doc.createElement("file")
        fname = os.path.basename(pkg['URLPath'])
        if self.output_dir:
            file_.setAttribute("name", os.path.join(self.output_dir, fname))
        else:
            file_.setAttribute("name", fname)
        self.files.appendChild(file_)
        for tag, key in (
            ('identity', 'PackageBase'),
            ('version', 'Version'),
            ('description', 'Description'),
            ('url', 'URLPath')
        ):
            tag = self.doc.createElement(tag)
            file_.appendChild(tag)
            val = self.doc.createTextNode(pkg[key])
            tag.appendChild(val)


class DownloadQueue(object):
    def __init__(self):
        self.dbs = list()
        self.sync_pkgs = list()
        self.aur_pkgs = list()

    def __bool__(self):
        return bool(self.dbs or self.sync_pkgs or self.aur_pkgs)

    def __nonzero__(self):
        return self.dbs or self.sync_pkgs or self.aur_pkgs

    def add_db(self, db, sig=False, files=False):
        self.dbs.append((db, sig, files))

    def add_sync_pkg(self, pkg, urls, sig=False):
        self.sync_pkgs.append((pkg, urls, sig))

    def add_aur_pkg(self, pkg):
        self.aur_pkgs.append(pkg)


# ----------------------------- Argument Parsing ----------------------------- #

def parse_args(args):
    parser = argparse.ArgumentParser(
        description='Generate metalinks for parallel and segmented package downloads.',
        prog='pm2ml',
    )
    parser.add_argument(
        'pkgs', nargs='*', default=[], metavar='<pkgname>',
        help='Packages or groups to download.'
    )
    parser.add_argument(
        '--all-deps', action='store_true',
        help='Include all dependencies even if they are already installed.'
    )
    if AUR_SUPPORT:
        parser.add_argument(
            '-a', '--aur', action='store_true',
            help='Enable AUR support.'
        )
    parser.add_argument(
        '--select', action='store_true',
        help='Present a package selection dialogue for package groups. This may be overridden by --noconfirm.'
    )
    parser.add_argument(
        '--arch', metavar='<architecture>',
        help='Override the architecture in the Pacman configuration file.'
    )
    if AUR_SUPPORT:
        parser.add_argument(
            '--aur-only', action='store_true',
            help='Only download AUR archives. Use with "-u" to only download tarballs for upgradable AUR packages.'
        )
    parser.add_argument(
        '-c', '--conf', metavar='<path>', default='/etc/pacman.conf',
        help='Set the pacman.conf file. Default: %(default)s'
    )
    parser.add_argument(
        '--noconfirm', action='store_true',
        help='Suppress user prompts.'
    )
    parser.add_argument(
        '-d', '--nodeps', action='store_true',
        help='Skip dependencies.'
    )
    parser.add_argument(
        '--debug', action='store_true',
        help='Display debugging information.'
    )
    parser.add_argument(
        '--files', action='store_true',
        help='Download Pacman "files" databases.'
    )
    parser.add_argument(
        '--ignore', action='append', default=[], metavar='<pkgname>',
        help='Ignore designated package.'
    )
    parser.add_argument(
        '--ignoregroup', action='append', default=[], metavar='<grpname>',
        help='Ignore packages belonging to designated groups.'
    )
    parser.add_argument(
        '--needed', action='store_true',
        help='Skip packages if they already exist in the cache.'
    )
    parser.add_argument(
        '-o', '--output-dir', metavar='<path>',
        help='Set the output directory.'
    )
    parser.add_argument(
        '-p', '--preference', action='store_true',
        help='Use preference attributes in URL elements in the metalink.'
    )
    parser.add_argument(
        '-r', '--reflector', nargs=argparse.REMAINDER,
        help='Enable Reflector support and treat remaining arguments as Reflector arguments. E.g. "-r --latest 50'
    )
    parser.add_argument(
        '-s', '--sigs', action='count', default=0,
        help='Include signature files for repos with optional and required SigLevels. Pass this flag twice to attempt to download signature for all databases and packages.'
    )
    parser.add_argument(
        '-u', '--sysupgrade', action='count', default=0,
        help='Download upgradable packages.'
    )
    parser.add_argument(
        '-v', '--verbose', action='store_true',
        help='Verbose output.'
    )
    parser.add_argument(
        '-y', '--databases', '--refresh', action='store_true',
        help='Download databases.'
    )
    pargs = parser.parse_args(args)
    if AUR_SUPPORT and pargs.aur_only:
        pargs.aur = True
    return pargs


# The intermediate DownloadQueue object is used to allow scripts to hook into
# the code.
def download_queue_to_metalink(download_queue, output_dir=None, set_preference=False):
    impl = xml.dom.minidom.getDOMImplementation()
    metalink = Metalink(impl, output_dir=output_dir, set_preference=set_preference)

    for args in download_queue.dbs:
        metalink.add_db(*args)

    for args in download_queue.sync_pkgs:
        metalink.add_sync_pkg(*args)

    for pkg in download_queue.aur_pkgs:
        metalink.add_aur_pkg(pkg)

    return metalink


def main(args=None):
    pargs = parse_args(args)
    pm2ml = Pm2ml(pargs)

    sync_pkgs, sync_deps, \
        aur_pkgs, aur_deps, \
        not_found, unknown_deps, orphans = pm2ml.resolve_targets_from_arguments()

    download_queue = pm2ml.build_download_queue(
        pargs, (sync_pkgs | sync_deps), (aur_pkgs | aur_deps)
    )

    metalink = download_queue_to_metalink(
        download_queue,
        output_dir=pargs.output_dir,
        set_preference=pargs.preference
    )
    print(metalink)

    orphans = set(p.name for p in orphans)
    for label, xs, f in (
        ('Not found', not_found, logging.error),
        ('Unresolved dependencies', unknown_deps, logging.error),
        ('Installed orphans', orphans, logging.warning)
    ):
        if xs:
            f('{}: {}'.format(label, ' '.join(sorted(xs))))


def run_main(args=None):
    try:
        main(args)
    except (KeyboardInterrupt, BrokenPipeError):
        pass
    except pyalpm.error as e:
        sys.exit(e)


if __name__ == '__main__':
    run_main()
