import argparse
import json
import os
from pathlib import Path

import repomd

def get_repos(repo_dir):
    # Get a list of all repo dirs
    repos = []

    for path in Path(repo_dir).rglob('repomd.xml'):
        repos.append(path.parent.parent.absolute())

    return repos

def get_distribution(repo):
    distribution_dir = repo.name

    parts = distribution_dir.split('-')
    distribution = '-'.join(parts[:2])

    return distribution

def get_package_info(repo, base_path):
    packages = []

    repo_data = repomd.load('file://' + str(repo))
    for package_data in repo_data:
        package = {'Package': package_data.name,
                   'Version': package_data.version,
                   'Build': package_data.release,
                   'Architecture': package_data.arch,
                   'Filename': os.path.relpath(str(repo) + '/' +
                                               package_data.location,
                                               base_path),
                   'Description': package_data.description,
                   'Distribution': get_distribution(repo),
                   'Url': package_data.url,
                   'Maintainer': package_data.vendor,
                   'Repo': os.path.relpath(str(repo), base_path).split('/')[0],
                   'Licence': package_data.license
                   }

        packages.append(package)

    return packages

def main():

    # Command line arguments
    parser = argparse.ArgumentParser(description='Scan a set of APT repos and '
                                                 'generate a JSON catalog of '
                                                 'the contents.')
    parser.add_argument("repo", help="the repo directory, or directory "
                                     "containing multiple repos")

    args = parser.parse_args()

    package_info = []
    repos = get_repos(args.repo)

    for repo in repos:
        package_info.extend(get_package_info(repo, args.repo))

    with open('yum.json', 'w') as output_file:
        json.dump(package_info, output_file, indent=2, sort_keys=True)

if __name__ == "__main__":
    main()