import sys
import os
import yaml
import traceback
import argparse
import configparser
import boto3
import logging
from time import sleep
from botocore.exceptions import ClientError
from typing import Dict
from zenyatta.aws.redis import ElasticCache
from zenyatta.aws import get_instance
from zenyatta.common.util import parse_config_file


def update_airflow_cfg_file(cfg_file_path: str, cfg_block: dict):
    if not os.path.isfile(cfg_file_path):
        raise Exception('could not find airflow config file {}'.format(cfg_file_path))
    try:
        config = configparser.SafeConfigParser()
        config.read(cfg_file_path)
        if cfg_block['engine'] in config.sections():
            if config.get(cfg_block['engine'], 'endpoint_addr') == cfg_block['endpoint_addr']:
                return
            else:
                config.set(cfg_block['engine'], 'endpoint_addr', cfg_block['endpoint_addr'])
        else:
            config.add_section(cfg_block['engine'])
            for item in cfg_block:
                if item != 'engine':
                    config.set(cfg_block['engine'], item, str(cfg_block[item]))
        with open(cfg_file_path, 'w') as f:
            config.write(f)
        logging.info('redis endpoint is appended.')

        # write to celery portion as well
        engine_string = "redis://{}:{}".format(cfg_block['endpoint_addr'], cfg_block['endpoint_port'])
        config.set("celery", "broker_url", engine_string)
        config.set("celery", "celery_result_backend", engine_string)
    except Exception as e:
        logging.info("exception {}".format(e))


def create_conn_block(conn_desc: Dict) -> Dict:
    block = {}
    block['engine'] = conn_desc['Engine']
    block['cluster_id'] = conn_desc['CacheClusterId']
    block['endpoint_addr'] = conn_desc['CacheNodes'][0]['Endpoint']['Address']
    block['endpoint_port'] = conn_desc['CacheNodes'][0]['Endpoint']['Port']
    return block


def get_or_create_elasticache_subnet(team: str, aws_region: str) -> str:
    subnet_name = "{team}-buddy".format(**locals())
    client = boto3.client('elasticache', region_name=aws_region)
    try:
        response = client.describe_cache_subnet_groups(CacheSubnetGroupName=subnet_name)
        return subnet_name
    except ClientError:
        instance = get_instance(aws_region)
        response = client.create_cache_subnet_group(
            CacheSubnetGroupName=subnet_name,
            CacheSubnetGroupDescription="cache subnet for {team} to the redis task queue in zenyatta"
            .format(**locals()),
            SubnetIds=[instance.subnet_id]
        )
        return response['CacheSubnetGroup']['CacheSubnetGroupName']


def launch_redis(cache_cluster_id, subnet_group, security_group_list):
    cfg_home = '/etc/zenyatta/'
    redis_cfg_name = 'hardware_cfg.yaml'
    airflow_cfg_name = 'airflow.cfg'
    cache_cluster_tag = cache_cluster_id

    redis_cfg = parse_config_file(cfg_home+redis_cfg_name)

    ec = ElasticCache('elasticache')
    airflow_cc_list = ec.fetch_cache_clusters(cache_cluster_id)
    if len(airflow_cc_list) == 1:
        logging.info("found one redis instance for airflow. updating endpoint change.")
        redis_endpoint_block = create_conn_block(airflow_cc_list[0])
        update_airflow_cfg_file(cfg_home+airflow_cfg_name, redis_endpoint_block)
    elif len(airflow_cc_list) == 0:
        logging.info("could not find a redis instance for airflow. buddy will create one.")
        rtn = ec.create_redis_cluster(redis_cfg, cache_cluster_id, cache_cluster_tag,
                                      subnet_group, security_group_list)
        if rtn is None:
            sys.exit(1)
        retry_attampt = 20
        while (len(airflow_cc_list) == 0 or airflow_cc_list[0]['CacheClusterStatus'] != 'available') \
                and retry_attampt > 0:

            sleep(30)
            airflow_cc_list = ec.fetch_cache_clusters(cache_cluster_id)
            retry_attampt -= 1
        if retry_attampt == 0:
            raise Exception("instance {} created but failed to fetch enpoint".format(cache_cluster_id))
            sys.exit(1)
        redis_endpoint_block = create_conn_block(airflow_cc_list[0])
        update_airflow_cfg_file(cfg_home+airflow_cfg_name, redis_endpoint_block)
    else:
        raise Exception("buddy found more than one cluster. please keep only one and delete the rest.")


if __name__ == '__main__':
    # buddy needs to provide team name, enviornment, subnet group, and a list of
    # security group in the form of calling
    # python setup_redis.py -n teamA -e dev or prod -s zenyatta-production
    #   -l security-group-1 security-group-2
    parser = argparse.ArgumentParser(description="set up redis instance for airflow")
    parser.add_argument('-n', '--name', required=True, dest='team', action='store',
                        help='buddy team name used as redis cluster identifier')
    parser.add_argument('-r', '--region', required=True, dest='region', action='store',
                        help='aws region to launch in')
    parser.add_argument('-s', '--subnet-group', required=False, dest='subnet_group', action='store',
                        help='subnet group name to set up redis cluster')
    parser.add_argument('-l', '--security-group-ids-list', required=False, dest='security_group_list',
                        type=str, nargs='*', action='store',
                        help='security group ids to set up redis cluster')
    args = parser.parse_args()
    if not args.subnet_group:
        # create subnet group
        args.subnet_group = get_or_create_elasticache_subnet(args.team, args.region)
    if not args.security_group_list:
        # coerce security group info from
        tmp = get_instance(args.region)
        args.security_group_list = [sg['GroupId'] for sg in tmp.security_groups]
    launch_redis('buddy-'+args.team, args.subnet_group, args.security_group_list)
