# Copyright (c) OpenMMLab. All rights reserved.

import logging

import os

from argparse import ArgumentParser

import requests

import yaml as yml

from mmseg.utils import get_root_logger

def check_url(url):

    """Check url response status.

    Args:

        url (str): url needed to check.

    Returns:

        int, bool: status code and check flag.

    """

    flag = True

    r = requests.head(url)

    status_code = r.status_code

    if status_code == 403 or status_code == 404:

        flag = False

    return status_code, flag

def parse_args():

    parser = ArgumentParser('url valid check.')

    parser.add_argument(

        '-m',

        '--model-name',

        type=str,

        help='Select the model needed to check')

    args = parser.parse_args()

    return args

def main():

    args = parse_args()

    model_name = args.model_name

    # yml path generate.

    # If model_name is not set, script will check all of the models.

    if model_name is not None:

        yml_list = [(model_name, f'configs/{model_name}/{model_name}.yml')]

    else:

        # check all

        yml_list = [(x, f'configs/{x}/{x}.yml') for x in os.listdir('configs/')

                    if x != '_base_']

    logger = get_root_logger(log_file='url_check.log', log_level=logging.ERROR)

    for model_name, yml_path in yml_list:

        # Default yaml loader unsafe.

        model_infos = yml.load(

            open(yml_path, 'r'), Loader=yml.CLoader)['Models']

        for model_info in model_infos:

            config_name = model_info['Name']

            checkpoint_url = model_info['Weights']

            # checkpoint url check

            status_code, flag = check_url(checkpoint_url)

            if flag:

                logger.info(f'checkpoint | {config_name} | {checkpoint_url} | '

                            f'{status_code} valid')

            else:

                logger.error(

                    f'checkpoint | {config_name} | {checkpoint_url} | '

                    f'{status_code} | error')

            # log_json check

            checkpoint_name = checkpoint_url.split('/')[-1]

            model_time = '-'.join(checkpoint_name.split('-')[:-1]).replace(

                f'{config_name}_', '')

            # two style of log_json name

            # use '_' to link model_time (will be deprecated)

            log_json_url_1 = f'https://download.openmmlab.com/mmsegmentation/v0.5/{model_name}/{config_name}/{config_name}_{model_time}.log.json'  # noqa

            status_code_1, flag_1 = check_url(log_json_url_1)

            # use '-' to link model_time

            log_json_url_2 = f'https://download.openmmlab.com/mmsegmentation/v0.5/{model_name}/{config_name}/{config_name}-{model_time}.log.json'  # noqa

            status_code_2, flag_2 = check_url(log_json_url_2)

            if flag_1 or flag_2:

                if flag_1:

                    logger.info(

                        f'log.json | {config_name} | {log_json_url_1} | '

                        f'{status_code_1} | valid')

                else:

                    logger.info(

                        f'log.json | {config_name} | {log_json_url_2} | '

                        f'{status_code_2} | valid')

            else:

                logger.error(

                    f'log.json | {config_name} | {log_json_url_1} & '

                    f'{log_json_url_2} | {status_code_1} & {status_code_2} | '

                    'error')

if __name__ == '__main__':

    main()