import json
import pulumi
import pulumi_aws as aws
import pulumi_aws_apigateway as apigateway
import api_gw

config = pulumi.Config()
aws_config = pulumi.Config("aws")
aws_region = aws_config.require("region")
account_id = aws.get_caller_identity().account_id


def create_lambda_role(lambda_name, iam_config=None):
    """Create IAM role for Lambda with configurable policies"""
    
    # Base managed policies
    managed_policies = [aws.iam.ManagedPolicy.AWS_LAMBDA_BASIC_EXECUTION_ROLE, 
                       aws.iam.ManagedPolicy.AWS_LAMBDA_VPC_ACCESS_EXECUTION_ROLE]
    
    if iam_config and "managed_policies" in iam_config:
        managed_policies.extend(iam_config["managed_policies"])
    
    role = aws.iam.Role(f"role-{lambda_name}", 
        assume_role_policy=json.dumps({
            "Version": "2012-10-17",
            "Statement": [{
                "Action": "sts:AssumeRole",
                "Effect": "Allow",
                "Principal": {"Service": "lambda.amazonaws.com"},
            }],
        }),
        managed_policy_arns=managed_policies
    )
    
    # Create custom inline policies from YAML config
    if iam_config and "custom_policies" in iam_config:
        for policy in iam_config["custom_policies"]:
            aws.iam.RolePolicy(f"{lambda_name}-{policy['name']}",
                role=role.id,
                policy=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [{
                        "Effect": policy["effect"],
                        "Action": policy["actions"],
                        "Resource": policy["resources"]
                    }]
                })
            )
    
    return role

lambda_api_configs = config.require_object("lambda-api")
for l_api in lambda_api_configs:
    ecr_repo = aws.ecr.get_repository(name=l_api["ecr"]["repo_name"])
    ecr_image = aws.ecr.get_image(repository_name=l_api["ecr"]["repo_name"], image_tag=l_api["ecr"]["tag"])
    lambda_name = l_api["name"]
    if not l_api["network_config"]["is_private"]:
        raise "Not implemented yet: Public lambda function"
    else:
        lambda_sg = aws.ec2.SecurityGroup(f"lambda-sg-{lambda_name}",
            description=f"SG for Lambda {lambda_name}",
            egress=[{
                "cidr_blocks": ["0.0.0.0/0"],
                "from_port": 0,
                "protocol": "-1",
                "to_port": 0,
            }],
            name=lambda_name,
            vpc_id=l_api["network_config"]["vpc_id"],
        )

        # print(pulumi.Output.all(ecr_repo.repository_url, ecr_image.image_digest).apply(lambda args: f'{args[0]}@{args[1]}'))
        # Create role for this specific Lambda
        role = create_lambda_role(lambda_name, l_api.get("iam"))
        if "env_vars" in l_api:
            variables={k:v for k,v in l_api["env_vars"].items()}
        else:
            variables={}
        # Define the Lambda function, replacing <IMAGE_URI> with your actual image URI
        fn = aws.lambda_.Function(f"{lambda_name}",
            package_type="Image",
            # image_uri=ecr_repo.repository_url.apply(lambda url: f"{url}:latest"),  # Assuming 'latest' tag
            image_uri=pulumi.Output.all(ecr_repo.repository_url, ecr_image.image_digest).apply(lambda args: f'{args[0]}@{args[1]}'),
            role=role.arn,
            timeout=l_api["timeout"],
            memory_size=l_api["memory"],
            environment={
                "variables": variables
            },
            vpc_config=dict(
                ipv6_allowed_for_dual_stack = False,
                subnet_ids = l_api["network_config"]["private_subnet_ids"],
                security_group_ids=[lambda_sg.id]
            ),
            publish=l_api["provisioned_concurrency"]>0, #necessary for provisioned concurrency
        )

        if l_api["provisioned_concurrency"]>0: 
            lambda_concurrency = aws.lambda_.ProvisionedConcurrencyConfig(l_api["name"],
                provisioned_concurrent_executions=l_api["provisioned_concurrency"],
                function_name=fn.name,
                qualifier=fn.version
            )

    api_config = l_api["api_gateway"]
    if api_config["use_api_gw"]:
        if api_config["communication_type"]=="HTTP":
            api_gw.create_api_gatewayv2(api_config, l_api, fn, aws_region, config.require('project_name'))
        else:
            api_gw.create_api_gateway(api_config, l_api, fn, account_id, aws_region, config.require('project_name'))