#!/usr/bin/env python3

__license__ = 'MIT'

import argparse
import json
import hashlib
import os
import requirements
import subprocess
import sys
import tempfile
import urllib.request

from collections import OrderedDict


parser = argparse.ArgumentParser()
parser.add_argument('packages', nargs='*')
parser.add_argument('--python2', action='store_true',
                    help='Look for a Python 2 package')
parser.add_argument('--cleanup', choices=['scripts', 'all'],
                    help='Select what to clean up after build')
parser.add_argument('--requirements-file', '-r',
                    help='Specify requirements.txt file')
parser.add_argument('--build-only', action='store_const',
                    dest='cleanup', const='all',
                    help='Clean up all files after build')
parser.add_argument('--no-build-isolation',
                    help='https://pip.pypa.io/en/stable/reference/pip/#pep-517-and-518-support')
parser.add_argument('--output', '-o',
                    help='Specify output file name')
opts = parser.parse_args()


def get_pypi_url(name: str, filename: str) -> str:
    url = 'https://pypi.org/pypi/{}/json'.format(name)
    print('Extracting download url for', name)
    with urllib.request.urlopen(url) as response:
        body = json.loads(response.read().decode('utf-8'))
        for release in body['releases'].values():
            for source in release:
                if source['filename'] == filename:
                    return source['url']
        else:
            raise Exception('Failed to extract url from {}'.format(url))


def get_package_name(filename: str) -> str:
    if filename.endswith(('bz2', 'gz', 'xz', 'zip')):
        segments = filename.split('-')
        if len(segments) == 2:
            return segments[0]
        return '-'.join(segments[:len(segments) - 1])
    elif filename.endswith('whl'):
        segments = filename.split('-')
        if len(segments) == 5:
            return segments[0]
        return '-'.join(segments[:len(segments) - 4])
    else:
        raise Exception(
            'Downloaded filename: {} does not end with bz2, gz, xz, zip, or whl'.format(filename)
        )


def get_file_hash(filename: str) -> str:
    sha = hashlib.sha256()
    print('Generating hash for', filename)
    with open(filename, 'rb') as f:
        while True:
            data = f.read(1024 * 1024 * 32)
            if not data:
                break
            sha.update(data)
        return sha.hexdigest()


def parse_continuation_lines(fin):
    for line in fin:
        line = line.rstrip('\n')
        while line.endswith('\\'):
            try:
                line = line[:-1] + next(fin).rstrip('\n')
            except StopIteration:
                exit('Requirements have a wrong number of line continuation characters "\\"')
        yield line


packages = []
if opts.requirements_file and os.path.exists(opts.requirements_file):
    with open(opts.requirements_file, 'r') as req_file:
        reqs = parse_continuation_lines(req_file)
        reqs_as_str = '\n'.join([r.split('--hash')[0] for r in reqs])
        packages = list(requirements.parse(reqs_as_str))
        requirements_file = opts.requirements_file
elif opts.packages:
    packages = list(requirements.parse('\n'.join(opts.packages)))
    with tempfile.NamedTemporaryFile('w', delete=False) as req_file:
        req_file.write('\n'.join(opts.packages))
        requirements_file = req_file.name
else:
    exit('Please specifiy either packages or requirements file argument')

with open(requirements_file, 'r') as req_file:
    use_hash = '--hash=' in req_file.read()

python_version = '2' if opts.python2 else '3'
if opts.python2:
    pip_executable = 'pip2'
else:
    pip_executable = 'pip3'

if opts.output:
    output_package = opts.output
elif opts.requirements_file:
    output_package = 'python{}-{}'.format(
        python_version, opts.requirements_file.replace('.txt', ''),
    )
elif len(packages) == 1:
    output_package = 'python{}-{}'.format(
        python_version, packages[0].name,
    )
else:
    output_package = 'python{}-modules'.format(python_version)
output_filename = output_package + '.json'

modules = []
vcs_modules = []
sources = {}

tempdir_prefix = 'pip-generator-{}'.format(output_filename.replace('.json', ''))
with tempfile.TemporaryDirectory(prefix=tempdir_prefix) as tempdir:
    pip_download = [
        pip_executable,
        'download',
        '--exists-action=i',
        '--dest',
        tempdir,
        '-r',
        requirements_file
    ]

    if use_hash:
        pip_download.append('--require-hashes')

    try:
        # May download the package twice, the first time it allows pip to
        # select the preferred package, if the package downloaded is not
        # platform generic, then it forces pip to download the sdist (using
        # the --no-binary option).

        subprocess.run(pip_download, check=True)
        for filename in os.listdir(tempdir):
            if not filename.endswith(('gz', 'any.whl', 'zip')):
                os.remove(os.path.join(tempdir, filename))

        subprocess.run(pip_download + [
            '--no-binary', ':all:'
            ], check=True)

    except subprocess.CalledProcessError:
        print('Failed to download')
        print('Please fix the module manually in the generated file')

    files = {get_package_name(f): [] for f in os.listdir(tempdir)}

    for filename in os.listdir(tempdir):
        name = get_package_name(filename)
        files[name].append(filename)

    # Delete redundant sources, where .zip (vcs sources) take precedence
    # over .whl, and .whl over tar.gz sources.
    for name in files:
        if len(files[name]) > 1:
            whl_source = False
            zip_source = False
            for f in files[name]:
                if f.endswith(('any.whl')):
                    whl_source = True
                if f.endswith(('.zip')):
                    zip_source = True
            if zip_source:
                for f in files[name]:
                    if not f.endswith(('.zip')):
                        os.remove(os.path.join(tempdir, f))
            elif whl_source:
                for f in files[name]:
                    if not f.endswith(('any.whl')):
                        os.remove(os.path.join(tempdir, f))

    vcs_packages = {
        x.name: {'vcs': x.vcs, 'revision': x.revision, 'uri': x.uri}
        for x in packages
        if x.vcs
    }

    for filename in os.listdir(tempdir):
        name = get_package_name(filename)
        sha256 = get_file_hash(os.path.join(tempdir, filename))

        if name in vcs_packages:
            uri = vcs_packages[name]['uri']
            revision = vcs_packages[name]['revision']
            vcs = vcs_packages[name]['vcs']
            url = 'https://' + uri.split('://', 1)[1]
            s = 'commit'
            if vcs == 'svn':
                s = 'revision'
            source = OrderedDict([
                ('type', vcs),
                ('url', url),
                (s, revision),
            ])
            is_vcs = True
        else:
            url = get_pypi_url(name, filename)
            source = OrderedDict([
                ('type', 'file'),
                ('url', url),
                ('sha256', sha256)])
            is_vcs = False
        sources[name] = {'source': source, 'vcs': is_vcs}

# Python3 packages that come as part of org.freedesktop.Sdk.
system_packages = ['Mako', 'Markdown', 'meson', 'pip', 'setuptools', 'six']
for package in packages:

    if package.name is None:
        print('Warning: skipping invalid requirement specification {} because it is missing a name'.format(package.line), file=sys.stderr)
        print('Append #egg=<pkgname> to the end of the requirement line to fix', file=sys.stderr)
        continue
    elif package.name in system_packages:
        continue

    if len(package.extras) > 0:
        extras = '[' + ','.join(extra for extra in package.extras) + ']'
    else:
        extras = ''

    version_list = [x[0] + x[1] for x in package.specs]
    version = ','.join(version_list)

    if package.vcs:
        revision = ''
        if package.revision:
            revision = '@' + package.revision
        pkg = package.uri + revision + '#egg=' + package.name
    else:
        pkg = package.name + extras + version

    dependencies = []
    # Downloads the package again to list dependencies

    tempdir_prefix = 'pip-generator-{}'.format(package.name)
    with tempfile.TemporaryDirectory(prefix=tempdir_prefix) as tempdir:
        pip_download = [
            pip_executable,
            'download',
            '--exists-action=i',
            '--dest',
            tempdir,
        ]
        try:
            print('Generating dependencies for {}'.format(package.name))
            subprocess.run(pip_download + [pkg], check=True, stdout=subprocess.DEVNULL)
            for filename in os.listdir(tempdir):
                dependencies.append(get_package_name(filename))

        except subprocess.CalledProcessError:
            print('Failed to download {}'.format(package.name))

    is_vcs = True if package.vcs else False
    package_sources = [sources[x]['source']
                       for x in dependencies
                       if x in sources
                       # Add vcs source only if package is vcs
                       and (not sources[x]['vcs'] or is_vcs)]

    if package.vcs:
        name_for_pip = '.'
    else:
        name_for_pip = pkg

    module_name = 'python{}-{}'.format(python_version, package.name)

    pip_command = [
        pip_executable,
        'install',
        '--exists-action=i',
        '--no-index',
        '--find-links="file://${PWD}"',
        '--prefix=${FLATPAK_DEST}',
        '"{}"'.format(name_for_pip)
    ]
    if opts.no_build_isolation:
        pip_command.append('--no-build-isolation')

    module = OrderedDict([
        ('name', module_name),
        ('buildsystem', 'simple'),
        ('build-commands', [' '.join(pip_command)]),
        ('sources', package_sources),
    ])
    if opts.cleanup == 'all':
        module['cleanup'] = ['*']
    elif opts.cleanup == 'scripts':
        module['cleanup'] = ['/bin', '/share/man/man1']

    if package.vcs:
        vcs_modules.append(module)
    else:
        modules.append(module)

if not opts.requirements_file:
    try:
        os.remove(requirements_file)
    except FileNotFoundError:
        pass

modules = vcs_modules + modules
if len(modules) == 1:
    pypi_module = modules[0]
else:
    pypi_module = {
        'name': output_package,
        'buildsystem': 'simple',
        'build-commands': [],
        'modules': modules,
    }

with open(output_filename, 'w') as output:
    output.write(json.dumps(pypi_module, indent=4))
    print('Output saved to {}'.format(output_filename))
