import argparse
import logging
import os
import re
import sys
import yaml


OS_LABEL_PREFIXES = {
    'os.template.cnv.io',
    'os.template.kubevirt.io'
}

RHEL_PATTERN = re.compile(r'rhel(\d+)\.(\d+)')
WINDOWS_PATTERN = re.compile(r'win(.*)')
FEDORA_PATTERN = re.compile(r'fedora(\d*)')
CENTOS_PATTERN = re.compile(r'centos(\d*)')

RHEL_MIN_VERSION = 6
RHEL_MIN_MINOR_VERSIONS = {6: 0, 7: 0, 8: 0}
FEDORA_MIN_VERSION = 31
CENTOS_MIN_VERSION = 7
WINDOWS_VERSIONS = {'2k12r2', '2k16', '2k19', '10'}


def os_names(labels):
    for label, value in labels.items():
        if value.lower() != 'true':
            continue

        try:
            prefix, os_name = label.split('/', 1)
        except ValueError:
            continue

        if prefix in OS_LABEL_PREFIXES:
            yield os_name


def is_rhel(os_name):
    match = RHEL_PATTERN.match(os_name)
    if not match:
        return False

    major, minor = map(int, match.groups())
    if major < RHEL_MIN_VERSION:
        return False

    return minor >= RHEL_MIN_MINOR_VERSIONS.get(major, 0)

def is_rhel6(os_name):
    return "6" in os_name

def is_fedora(os_name):
    match = FEDORA_PATTERN.match(os_name)
    if not match:
        return False

    version = int(match.group(1))
    return version >= FEDORA_MIN_VERSION


def is_centos(os_name):
    match = CENTOS_PATTERN.match(os_name)
    if not match:
        return False

    version = int(match.group(1))
    return version >= CENTOS_MIN_VERSION


def is_windows(doc):
    # Because all windows templates should exist DS and some of them had their os_name labels removed 
    # because they became deprecated, all templates that their names match the WINDOWS_PATTERN
    # should be included DS
    # This fix was explicitly made for the deprecated win2k12r2 template
    match = WINDOWS_PATTERN.match(doc['metadata']['name'])
    return match


def is_supported(doc):
    if is_windows(doc):
        annotations = doc['metadata']['annotations']

        annotations['template.kubevirt.io/provider'] = 'Red Hat'
        annotations['template.kubevirt.io/provider-support-level'] = 'Full'
        annotations['template.kubevirt.io/provider-url'] = 'https://www.redhat.com'

        return True

    try:
        labels = doc['metadata']['labels']
    except KeyError:
        return False

    for os_name in os_names(labels):
        # If a supported OS is found, apply support annotations to it
        annotations = doc['metadata']['annotations']

        if is_rhel(os_name):
            if not is_rhel6(os_name):
                annotations['template.kubevirt.io/provider'] = 'Red Hat'
                annotations['template.kubevirt.io/provider-support-level'] = 'Full'
                annotations['template.kubevirt.io/provider-url'] = 'https://www.redhat.com'
            return True

        elif is_fedora(os_name):
            annotations['template.kubevirt.io/provider'] = 'Red Hat'
            annotations['template.kubevirt.io/provider-support-level'] = 'Community'
            annotations['template.kubevirt.io/provider-url'] = 'https://www.fedoraproject.org'
            return True

        elif is_centos(os_name):
            annotations['template.kubevirt.io/provider'] = 'Red Hat'
            annotations['template.kubevirt.io/provider-support-level'] = 'Community'
            annotations['template.kubevirt.io/provider-url'] = 'https://www.centos.org'
            return True

    return False


def yaml_files(directory):
    for file in os.listdir(directory):
        path = os.path.join(directory, file)
        if os.path.isfile(path):
            yield path


def main():
    logging.basicConfig(level=logging.INFO)

    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("dir", help="directory containing manifests")
    args = arg_parser.parse_args()

    if not os.path.isdir(args.dir):
        logging.error('"%s" is not a directory', args.dir)
        sys.exit(2)

    for path in yaml_files(args.dir):
        logging.info("Filtering file: %s", path)
        with open(path, 'r') as f:
            yaml_str = f.read()

        try:
            filtered = list(filter(is_supported, yaml.safe_load_all(yaml_str)))
        except Exception as err:
            logging.warning("Error parsing file %s:\n%s", path, err)
            continue

        if len(filtered) == 0:
            logging.info("No templates passed filter, removing file: %s", path)
            os.remove(path)
            continue

        logging.info("%s templates passed filter", len(filtered))
        with open(path, 'w') as f:
            yaml.safe_dump_all(filtered, f, default_flow_style=False)


if __name__ == '__main__':
    main()