import binascii
import gzip
import lzma
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 _is_compressed_file(first_two_bytes: bytes, initial_bytes: bytes):
    return binascii.hexlify(first_two_bytes) == initial_bytes


def is_gzip_file(first_two_bytes):
    return _is_compressed_file(
        first_two_bytes=first_two_bytes,
        initial_bytes=b'1f8b',
    )


def is_xz_file(first_two_bytes):
    return _is_compressed_file(
        first_two_bytes=first_two_bytes,
        initial_bytes=b'fd37',
    )


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)
    # Defaults modules can be empty, but pungi detects
    # empty folder while copying and raises the exception in this case
    Path(os.path.join(module_defaults_path, '.empty')).touch()

    for module_file in modules_paths:
        data = module_file.read()
        if is_gzip_file(data[:2]):
            data = gzip.decompress(data)
        elif is_xz_file(data[:2]):
            data = lzma.decompress(data)
        documents = yaml.load_all(data, Loader=yaml.BaseLoader)
        for doc in documents:
            if doc['document'] == 'modulemd-defaults':
                name = f"{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('-', '_')
                doc_data = doc['data']
                name = f"{doc_data['name']}-{stream}-" \
                       f"{doc_data['version']}.{doc_data['context']}"
                arch_dir = os.path.join(
                    modules_path,
                    doc_data['arch']
                )
                os.makedirs(arch_dir, exist_ok=True)
                path = os.path.join(
                    arch_dir,
                    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()