import pulumi
import pulumi_aws as aws
import pulumi_docker as docker
import conf as config
import json


def setup_kms():
    # KMS Key Setup
    app_key = aws.kms.Key(f"{config.project_name}-key",
        description="Key for encrypting secrets",
        enable_key_rotation=True,
        policy=json.dumps({
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Sid": "",
                    "Principal": {
                        "AWS": f"arn:aws:iam::{config.account_id}:root",
                    },
                    "Action": [
                        "kms:Create*", "kms:Describe*", "kms:Enable*", "kms:List*", "kms:Put*", "kms:Update*",
                        "kms:Revoke*", "kms:Disable*", "kms:Get*", "kms:Delete*", "kms:ScheduleKeyDeletion",
                        "kms:CancelKeyDeletion", "kms:Tag*", "kms:UntagResource",
                    ],
                    "Resource": "*",
                },
                {
                    "Effect": "Allow",
                    "Principal": {
                        "AWS": f"arn:aws:iam::{config.account_id}:root",
                    },
                    "Action": [
                        "kms:Encrypt", "kms:Decrypt", "kms:ReEncrypt*", "kms:GenerateDataKey*", "kms:DescribeKey",
                    ],
                    "Resource": "*",
                },
                {
                    "Sid": 'Allow access to EFS for all principals in the account that are authorized to use EFS',
                    "Effect": 'Allow',
                    "Principal": {"AWS": "*"},
                    "Action": [
                        "kms:Encrypt", "kms:Decrypt", "kms:ReEncrypt*", "kms:GenerateDataKey*",
                        "kms:CreateGrant", "kms:DescribeKey",
                    ],
                    "Resource": "*",
                    "Condition": {
                        "StringEquals": {
                            "kms:ViaService": f"elasticfilesystem.{config.aws_region}.amazonaws.com",
                            "kms:CallerAccount": config.account_id,
                        },
                    },
                },
            ],
        }),
        tags={
            "pulumi-application": config.project_name,
            "pulumi-environment": config.stack_name,
        },
    )

    # SSM Parameter Setup
    app_ssm_parameter = aws.ssm.Parameter(f"{config.project_name}-ssm-parameter",
        type="SecureString",
        value=config.config.require_secret("bedrock_api_key"),
        key_id=app_key.key_id,
        name=f"/{config.project_name}/{config.stack_name}/BEDROCK_API_KEY",
        tags={
            "pulumi-application": config.project_name,
            "pulumi-environment": config.stack_name,
        },
    )
    return app_ssm_parameter, app_key