#!/usr/bin/env python3

"""
  This script depends on pm2ml, rsync and aria2. It accepts an Arch Linux rsync
  mirror as its first argument optionally followed by a local Pacserve server
  URL and then pm2ml arguments to control the download the download.

  The easiest way to use this script is to create a wrapper script or alias with
  a fixed rsync server (along with an optional pacserve server).

  See the Arch Linux mirror status page

    https://www.archlinux.org/mirrors/status

  or use `reflector -p rsync` to get the URL of an rsync mirror.

"""

from os import chdir, makedirs
from os.path import join
from pm2ml import build_download_queue, DownloadQueue, download_queue_to_metalink
from pycman import config
from subprocess import Popen, PIPE
from sys import argv, stderr, exit
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import urlopen

RSYNC_REPOS = (
  'community',
  'community-staging',
  'community-testing',
  'core',
  'extra',
  'gnome-unstable',
  'kde-unstable',
  'multilib',
  'multilib-testing'
  'staging',
  'testing'
)

def download_queue_to_rsync_cmd(conf, url_str, queue, output_dir=None):
  url = urlparse(url_str)
  host = url.netloc
  # remove initial slash
  path = url.path[1:].replace('$arch', conf.options['Architecture'])

  cmd = ['rsync', '-aL', '--progress', '--no-motd']
  host_added = False
  for db, sigs in queue.dbs:
    db_path = '::' + join(path.replace('$repo', db.name), '%s.db' % db.name)
    if not host_added:
      db_path = host + db_path
    cmd.append(db_path)
    if sigs:
      cmd.append(db_path + '.sig')

  for pkg, urls, sigs in queue.sync_pkgs:
    pkg_path = '::' + join(path.replace('$repo', pkg.db.name), pkg.filename)
    if not host_added:
      pkg_path = host + pkg_path
    cmd.append(pkg_path)
    if sigs:
      cmd.append(pkg_path + '.sig')

  if not output_dir:
    output_dir = '.'
  cmd.append(output_dir)
  return cmd





def main(rsync_url, args=None):
  if args and args[0][:4] == 'http':
    pacserver = args[0]
    args = args[1:]
  else:
    pacserver = None

  pargs, conf, download_queue, not_found, missing_deps = build_download_queue(args)

  rsync_queue = DownloadQueue()
  metalink_queue = DownloadQueue()

  for db, sigs in download_queue.dbs:
    if rsync_url and db.name in RSYNC_REPOS:
      rsync_queue.add_db(db, sigs)
    else:
      metalink_queue.add_db(db, sigs)

  if pacserver:
    try_pacserver = True
  else:
    try_pacserver = False

  for pkg, urls, sigs in download_queue.sync_pkgs:
    use_pacserver = False
    if try_pacserver:
      pacserve_url = join(pacserver, 'search', pkg.db.name, pkg.arch, pkg.filename)
      try:
        with urlopen(pacserve_url) as f:
          found_url = f.geturl().replace('/search/', '/request/')
          if found_url.startswith(pacserver):
            continue
          urls = [found_url]
          use_pacserver = True
      except HTTPError as e:
        if e.code != 418:
          raise e
      except URLError as e:
        stderr.write('warning: failed to query Pacserve: %s\n' % e.reason)
        try_pacserver = False
    if not use_pacserver and rsync_url and pkg.db.name in RSYNC_REPOS:
      rsync_queue.add_sync_pkg(pkg, urls, sigs)
    else:
      metalink_queue.add_sync_pkg(pkg, urls, sigs)

  metalink_queue.aur_pkgs = download_queue.aur_pkgs

  rsync_cmd = download_queue_to_rsync_cmd(conf, rsync_url, rsync_queue)

  metalink = str(download_queue_to_metalink(metalink_queue)).encode()
  aria2_cmd = ('aria2c', '--metalink-file=-', '--conf-path=/etc/ppl.conf')


  if pargs.output_dir:
    makedirs(pargs.output_dir, exist_ok=True)
    chdir(pargs.output_dir)

  if rsync_queue:
    rsync_p = Popen(rsync_cmd)

  if metalink_queue:
    aria2c_p = Popen(aria2_cmd, stdin=PIPE)
    aria2c_p.communicate(input=metalink)

  if rsync_queue:
    e = rsync_p.wait()
    if e != 0:
      stderr.write('error: rsync exited with %d\n' % e)

  if metalink_queue:
    e = aria2c_p.wait()
    if e != 0:
      stderr.write('error: aria2c exited with %d\n' % e)
      stderr.write('%s\n' % metalink.decode())




if __name__ == '__main__':
  try:
    url, opts = argv[1], argv[2:]
  except IndexError:
    stderr.write('usage: %s <rsync_url> [pacserve url] [pm2ml options]\n' % argv[0])
    exit(1)
  main(url, opts)