import gzip
import os
from argparse import ArgumentParser, FileType
from io import BytesIO
from pathlib import Path
from typing import List, AnyStr
import logging
import yaml
import createrepo_c as cr
from typing.io import BinaryIO


def grep_list_of_modules_yaml_gz(repo_path: AnyStr) -> List[BytesIO]:
    """
    Find all of valid *modules.yaml.gz in repos
    :param repo_path: path to a directory which contains repodirs
    :return: list of content from *modules.yaml.gz
    """

    result = []
    for path in Path(repo_path).rglob('repomd.xml'):
        repo_dir_path = Path(path.parent).parent
        repomd_obj = cr.Repomd(str(path))
        for record in repomd_obj.records:
            if record.type != 'modules':
                continue
            with open(os.path.join(
                    repo_dir_path,
                    record.location_href,
            ), 'rb') as fp:
                result.append(
                    BytesIO(fp.read())
                )
    return result


def collect_modules(modules_paths: List[BinaryIO], target_dir: str):
    """
    Read given modules.yaml.gz files and export modules
    and modulemd files from it.
    Returns:
        object:
    """
    modules_path = os.path.join(target_dir, 'modules')
    module_defaults_path = os.path.join(target_dir, 'module_defaults')
    os.makedirs(modules_path, exist_ok=True)
    os.makedirs(module_defaults_path, exist_ok=True)

    for module_file in modules_paths:
        data = gzip.decompress(module_file.read())
        documents = yaml.load_all(data, Loader=yaml.BaseLoader)
        for doc in documents:
            if doc['document'] == 'modulemd-defaults':
                name = doc['data']['module'] + '.yaml'
                path = os.path.join(module_defaults_path, name)
                logging.info('Found %s module defaults', name)
            else:
                # pungi.phases.pkgset.sources.source_koji.get_koji_modules
                stream = doc['data']['stream'].replace('-', '_')
                name = '%s-%s-%s.%s' % (
                    doc['data']['name'],
                    stream,
                    doc['data']['version'],
                    doc['data']['context']
                )
                path = os.path.join(modules_path, name)
                logging.info('Found module %s', name)

                if 'artifacts' not in doc['data']:
                    logging.warning(
                        'RPM %s does not have explicit list of artifacts',
                        name
                    )

            with open(path, 'w') as f:
                yaml.dump(doc, f, default_flow_style=False)


def cli_main():
    parser = ArgumentParser()
    path_group = parser.add_mutually_exclusive_group(required=True)
    path_group.add_argument(
        '-p', '--path',
        type=FileType('rb'), nargs='+',
        help='Path to modules.yaml.gz file. '
             'You may pass multiple files by passing -p path1 path2'
    )
    path_group.add_argument(
        '-rp', '--repo-path',
        required=False,
        type=str,
        default=None,
        help='Path to a directory which contains repodirs. E.g. /var/repos'
    )
    parser.add_argument('-t', '--target', required=True)

    namespace = parser.parse_args()
    if namespace.repo_path is None:
        modules = namespace.path
    else:
        modules = grep_list_of_modules_yaml_gz(namespace.repo_path)
    collect_modules(
        modules,
        namespace.target,
    )


if __name__ == '__main__':
    cli_main()