from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

import logging
import os
import sys

from botocore.exceptions import ClientError

from stacker.session_cache import get_session
from stacker.hooks import utils
from stacker.ui import get_raw_input


logger = logging.getLogger(__name__)

KEYPAIR_LOG_MESSAGE = "keypair: %s (%s) %s"


def get_existing_key_pair(ec2, keypair_name):
    resp = ec2.describe_key_pairs()
    keypair = next((kp for kp in resp["KeyPairs"]
                    if kp["KeyName"] == keypair_name), None)

    if keypair:
        logger.info(KEYPAIR_LOG_MESSAGE,
                    keypair["KeyName"],
                    keypair["KeyFingerprint"],
                    "exists")
        return {
            "status": "exists",
            "key_name": keypair["KeyName"],
            "fingerprint": keypair["KeyFingerprint"],
        }

    logger.info("keypair: \"%s\" not found", keypair_name)
    return None


def import_key_pair(ec2, keypair_name, public_key_data):
    keypair = ec2.import_key_pair(
        KeyName=keypair_name,
        PublicKeyMaterial=public_key_data.strip(),
        DryRun=False)
    logger.info(KEYPAIR_LOG_MESSAGE,
                keypair["KeyName"],
                keypair["KeyFingerprint"],
                "imported")
    return keypair


def read_public_key_file(path):
    try:
        with open(utils.full_path(path), 'rb') as f:
            data = f.read()

        if not data.startswith(b"ssh-rsa"):
            raise ValueError(
                "Bad public key data, must be an RSA key in SSH authorized "
                "keys format (beginning with `ssh-rsa`)")

        return data.strip()
    except (ValueError, IOError, OSError) as e:
        logger.error("Failed to read public key file {}: {}".format(
            path, e))
        return None


def create_key_pair_from_public_key_file(ec2, keypair_name, public_key_path):
    public_key_data = read_public_key_file(public_key_path)
    if not public_key_data:
        return None

    keypair = import_key_pair(ec2, keypair_name, public_key_data)
    return {
        "status": "imported",
        "key_name": keypair["KeyName"],
        "fingerprint": keypair["KeyFingerprint"],
    }


def create_key_pair_in_ssm(ec2, ssm, keypair_name, parameter_name,
                           kms_key_id=None):
    keypair = create_key_pair(ec2, keypair_name)
    try:
        kms_key_label = 'default'
        kms_args = {}
        if kms_key_id:
            kms_key_label = kms_key_id
            kms_args = {"KeyId": kms_key_id}

        logger.info("Storing generated key in SSM parameter \"%s\" "
                    "using KMS key \"%s\"", parameter_name, kms_key_label)

        ssm.put_parameter(
            Name=parameter_name,
            Description="SSH private key for KeyPair \"{}\" "
                        "(generated by Stacker)".format(keypair_name),
            Value=keypair["KeyMaterial"],
            Type="SecureString",
            Overwrite=False,
            **kms_args)
    except ClientError:
        # Erase the key pair if we failed to store it in SSM, since the
        # private key will be lost anyway

        logger.exception("Failed to store generated key in SSM, deleting "
                         "created key pair as private key will be lost")
        ec2.delete_key_pair(KeyName=keypair_name, DryRun=False)
        return None

    return {
        "status": "created",
        "key_name": keypair["KeyName"],
        "fingerprint": keypair["KeyFingerprint"],
    }


def create_key_pair(ec2, keypair_name):
    keypair = ec2.create_key_pair(KeyName=keypair_name, DryRun=False)
    logger.info(KEYPAIR_LOG_MESSAGE,
                keypair["KeyName"],
                keypair["KeyFingerprint"],
                "created")
    return keypair


def create_key_pair_local(ec2, keypair_name, dest_dir):
    dest_dir = utils.full_path(dest_dir)
    if not os.path.isdir(dest_dir):
        logger.error("\"%s\" is not a valid directory", dest_dir)
        return None

    file_name = "{0}.pem".format(keypair_name)
    key_path = os.path.join(dest_dir, file_name)
    if os.path.isfile(key_path):
        # This mimics the old boto2 keypair.save error
        logger.error("\"%s\" already exists in \"%s\" directory",
                     file_name, dest_dir)
        return None

    # Open the file before creating the key pair to catch errors early
    with open(key_path, "wb") as f:
        keypair = create_key_pair(ec2, keypair_name)
        f.write(keypair["KeyMaterial"].encode("ascii"))

    return {
        "status": "created",
        "key_name": keypair["KeyName"],
        "fingerprint": keypair["KeyFingerprint"],
        "file_path": key_path
    }


def interactive_prompt(keypair_name, ):
    if not sys.stdin.isatty():
        return None, None

    try:
        while True:
            action = get_raw_input(
                "import or create keypair \"%s\"? (import/create/cancel) " % (
                    keypair_name,
                )
            )

            if action.lower() == "cancel":
                break

            if action.lower() in ("i", "import"):
                path = get_raw_input("path to keypair file: ")
                return "import", path.strip()

            if action.lower() == "create":
                path = get_raw_input("directory to save keyfile: ")
                return "create", path.strip()
    except (EOFError, KeyboardInterrupt):
        return None, None

    return None, None


def ensure_keypair_exists(provider, context, **kwargs):
    """Ensure a specific keypair exists within AWS.

    If the key doesn't exist, upload it.

    Args:
        provider (:class:`stacker.providers.base.BaseProvider`): provider
            instance
        context (:class:`stacker.context.Context`): context instance
        keypair (str): name of the key pair to create
        ssm_parameter_name (str, optional): path to an SSM store parameter to
            receive the generated private key, instead of importing it or
            storing it locally.
        ssm_key_id (str, optional): ID of a KMS key to encrypt the SSM
            parameter with. If omitted, the default key will be used.
        public_key_path (str, optional): path to a public key file to be
            imported instead of generating a new key. Incompatible with the SSM
            options, as the private key will not be available for storing.

    Returns:
        In case of failure ``False``, otherwise a dict containing:
            status (str): one of "exists", "imported" or "created"
            key_name (str): name of the key pair
            fingerprint (str): fingerprint of the key pair
            file_path (str, optional): if a new key was created, the path to
                the file where the private key was stored

    """

    keypair_name = kwargs["keypair"]
    ssm_parameter_name = kwargs.get("ssm_parameter_name")
    ssm_key_id = kwargs.get("ssm_key_id")
    public_key_path = kwargs.get("public_key_path")

    if public_key_path and ssm_parameter_name:
        logger.error("public_key_path and ssm_parameter_name cannot be "
                     "specified at the same time")
        return False

    session = get_session(region=provider.region,
                          profile=kwargs.get("profile"))
    ec2 = session.client("ec2")

    keypair = get_existing_key_pair(ec2, keypair_name)
    if keypair:
        return keypair

    if public_key_path:
        keypair = create_key_pair_from_public_key_file(
            ec2, keypair_name, public_key_path)

    elif ssm_parameter_name:
        ssm = session.client('ssm')
        keypair = create_key_pair_in_ssm(
            ec2, ssm, keypair_name, ssm_parameter_name, ssm_key_id)
    else:
        action, path = interactive_prompt(keypair_name)
        if action == "import":
            keypair = create_key_pair_from_public_key_file(
                ec2, keypair_name, path)
        elif action == "create":
            keypair = create_key_pair_local(ec2, keypair_name, path)
        else:
            logger.warning("no action to find keypair, failing")

    if not keypair:
        return False

    return keypair
