from botocore.exceptions import ClientError
from botocore.client import Config
from typing import Dict, Any
import boto3
import uuid
from airflow.settings import conf
from airflow.exceptions import AirflowConfigException
import zenyatta
from subprocess import check_output
import logging


def get_creds(arn: str, region: str=None) -> Dict:
    sts = boto3.client('sts')
    assumed_role = sts.assume_role(RoleArn=arn, RoleSessionName="zenyatta-"+str(uuid.uuid4()))
    if not region:
        try:
            region = conf.get('aws', 'region')
        except AirflowConfigException:
            region = "us-west-2"
    assumed_role.update({
        'access_key': assumed_role['Credentials']['AccessKeyId'],
        'secret_key': assumed_role['Credentials']['SecretAccessKey'],
        'token': assumed_role['Credentials']['SessionToken'],
        'region': region
    })
    return assumed_role


def boto_session(arn: str, region: str=None) -> (boto3.session.Session, Dict):
    aws_creds = get_creds(arn, region)
    return (boto3.Session(aws_access_key_id=aws_creds['access_key'],
                          aws_secret_access_key=aws_creds['secret_key'],
                          aws_session_token=aws_creds['token'],
                          region_name=aws_creds['region']),
            aws_creds)


def boto_client(resource: str, arn: str=None, region: str='us-west-2', signature_version: str=None) -> (Any, Any):
    config = None

    if signature_version:
        config = Config(signature_version=signature_version)

    if arn is None:
        return boto3.client(resource, region_name=region, config=config), None
    else:
        sesh, creds = boto_session(arn, region)
        return sesh.client(resource,
                           region_name=region,
                           aws_access_key_id=creds['Credentials']['AccessKeyId'],
                           aws_secret_access_key=creds['Credentials']['SecretAccessKey'],
                           aws_session_token=creds['Credentials']['SessionToken'],
                           config=config,
                           ), creds


def boto_resource(resource: str, arn: str=None, region: str='us-west-2') -> (Any, Any):
    if arn is None:
        return boto3.resource(resource, region_name=region), None
    else:
        sesh, creds = boto_session(arn, region)
        return sesh.resource(resource,
                             region_name=region,
                             aws_access_key_id=creds['Credentials']['AccessKeyId'],
                             aws_secret_access_key=creds['Credentials']['SecretAccessKey'],
                             aws_session_token=creds['Credentials']['SessionToken'],
                             ), creds


def get_role_from_conn_id(conn_id: str) -> str:
    """
    order of precedence is to check the connections.yaml, then the config
    :param conn_id:
    :return:
    """
    meta = get_aws_metadata(conn_id)
    try:
        return meta['role_arn']
    except KeyError:
        return meta['ec2_arn']
    except TypeError:  # nothing in connections.yaml, fallback to conf
        from airflow.configuration import conf
        return conf.get('aws', 'role_arn')


def get_aws_metadata(name: str) -> Dict[str, str]:
    """
    return { 'AWS_KEY': "from sandstorm"
             'AWS_SECRET': "from sandstorm"
             'S3_BUCKET': "bucket_name_string"
    """
    import zenyatta.common as common
    try:
        conns = common.get_connections()['aws'].get(name)  # type: Dict
    except Exception as e:
        raise common.ZenyattaError("couldn't load conns: {0}".format(e))
    return conns


def get_instance(aws_region: str) -> boto3.resources:
    """
    return the underlying ec2 instance so that the caller can get ec2 settings like region, security group,
    and etc to create or configure new aws objects running in the same account
    :param aws_region
    :return:
    """
    ec2 = boto3.resource('ec2', region_name=aws_region)
    instance_id = check_output(["ec2metadata", "--instance-id"]).decode().strip()
    instance = ec2.Instance(instance_id)
    return instance
