"""
Collect secrets from AWS SSM Parameter Store and output as environment variable exports.
"""
import argparse
from dataclasses import dataclass
import sys
from typing import Any, Dict, Iterable, List, Optional

import boto3


@dataclass
class EnvironmentVariable:
    name: str
    value: str


def get_parameters(path: str) -> List[Dict[str, Any]]:
    """
    Retrieve parameters from AWS SSM Parameter Store. Decrypts any encrypted parameters.

    Relies on the appropriate environment variables to authenticate against AWS:
    https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html
    """
    ssm = boto3.client("ssm")
    next_token: Optional[bool] = True
    parameters: List[Dict[str, Any]] = []
    while next_token is not None:
        kwargs = {"Path": path, "Recursive": False, "WithDecryption": True}
        if next_token is not True:
            kwargs["NextToken"] = next_token
        response = ssm.get_parameters_by_path(**kwargs)
        new_parameters = response.get("Parameters", [])
        parameters.extend(new_parameters)
        next_token = response.get("NextToken")

    return parameters


def parameter_to_env(parameter_object: Dict[str, Any]) -> EnvironmentVariable:
    """
    Transforms parameters returned by the AWS SSM API into EnvironmentVariables.
    """
    parameter_path = parameter_object.get("Name")
    if parameter_path is None:
        raise ValueError('Did not find "Name" in parameter object')
    name = parameter_path.split("/")[-1].upper()

    value = parameter_object.get("Value")
    if value is None:
        raise ValueError('Did not find "Value" in parameter object')

    return EnvironmentVariable(name, value)


def env_string(env_vars: Iterable[EnvironmentVariable], with_export: bool) -> str:
    """
    Produces a string which, when executed in a shell, exports the desired environment variables as
    specified by env_vars.
    """
    prefix = "export " if with_export else ""
    return "\n".join([f'{prefix}{var.name}="{var.value}"' for var in env_vars])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Materialize environment variables from AWS SSM Parameter Store"
    )
    parser.add_argument(
        "path",
        help="SSM path from which to pull environment variables (pull is NOT recursive)",
    )
    parser.add_argument(
        "-o", "--outfile", type=argparse.FileType("w"), default=sys.stdout
    )
    parser.add_argument(
        "--export",
        action="store_true",
        help="Set to output environment strings with export statements",
    )

    args = parser.parse_args()

    result = env_string(map(parameter_to_env, get_parameters(args.path)), args.export)
    with args.outfile as ofp:
        print(result, file=ofp)