from os import environ
from .client import aws_session, aws_client
from base64 import b64decode
import logging
from cachetools import cached, TTLCache


ROLE_ARN = environ.get("PARAMETER_STORE_ACCESS_ROLE_ARN", "arn:aws:iam::803813156102:role/ParameterStoreAccessRole")
REGION = environ.get("PARAMETER_STORE_REGION", 'us-east-1')
PROFILE = environ.get("PARAMETER_STORE_PROFILE")
CONFIG_ROOT = environ.get("PARAMETER_CONFIG_ROOT", '/config')

log_level = logging.getLevelName(environ['LOG_LEVEL']) if 'LOG_LEVEL' in environ else logging.INFO
logging.basicConfig(format='%(asctime)s,%(msecs)d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s',
                    datefmt='%d-%m-%Y:%H:%M:%S',
                    level=log_level)
logger = logging.getLogger('py-aws-util-config')

cache = TTLCache(maxsize=100, ttl=300)


@cached(cache)
def get_environment():
    environment = environ.get('BMC_ENVIRONMENT') or environ.get('bmc_environment') \
                  or environ.get('BMC_ENV') or environ.get('bmc_env') \
                  or environ.get('ENVIRONMENT') or environ.get('environment') \
                  or environ.get('ENV') or environ.get('env')
    if not environment:
        environment = 'development'
    return environment


@cached(cache)
def get_config_parameter(name, with_decryption=False, environment: str=None):
    _env = environment or get_environment()
    _name = f"{CONFIG_ROOT}/{_env}/{name}".replace("//", "/")
    ssm_client = aws_client('ssm', region=REGION)
    # Trying to improve performance by avoiding the assume role call unless necessary
    try:
        logger.debug("Requesting config parameter: " + _name)
        response = ssm_client.get_parameter(Name=_name, WithDecryption=with_decryption)
        return response['Parameter']['Value']
    except:
        ssm_client = aws_client('ssm', role_arn=ROLE_ARN, session_name='config_session', region=REGION)
        logger.debug("Requesting config parameter: " + _name)
        response = ssm_client.get_parameter(Name=_name, WithDecryption=with_decryption)
        return response['Parameter']['Value']


@cached(cache)
def get_config(path: str, recursive: bool = False, with_decryption=False, environment: str=None, asDict: bool=False):
    _env = environment or get_environment()
    _path = f"{CONFIG_ROOT}/{_env}/{path}".replace("//", "/")
    try:
        ssm_client = aws_client('ssm', region=REGION)
        logger.debug("Requesting config path: " + _path)
        response = ssm_client.get_parameters_by_path(Path=_path,
                                                     Recursive=recursive,
                                                     WithDecryption=with_decryption)
        params = response['Parameters']
        if asDict:
            return {x['Name'][len(_path) + 1:]: x['Value'] for x in params}
        else:
            return params
    except:
        ssm_client = aws_client('ssm', role_arn=ROLE_ARN, session_name='config_session', region=REGION)
        logger.debug("Requesting config path, with role: " + _path)
        response = ssm_client.get_parameters_by_path(Path=_path,
                                                     Recursive=recursive,
                                                     WithDecryption=with_decryption)
        params = response['Parameters']
        if asDict:
            return {x['Name'][len(_path) + 1:]: x['Value'] for x in params}
        else:
            return params


@cached(cache)
def decrypt(secret):
    config_session = aws_session(role_arn=ROLE_ARN, session_name='config_session', region=REGION)
    kms_client = aws_client('kms', session=config_session, region=REGION)
    return kms_client.decrypt(CiphertextBlob=b64decode(secret))['Plaintext'].decode("utf-8").strip('\n')