import os
from typing import Callable, Mapping, Optional, Sequence
import cdk_nag
import aws_cdk.aws_iam as aws_iam
from aws_cdk import BundlingOptions, aws_lambda, Stack
from constructs import Construct


class AppProjectFunction(Construct):
    def __init__(
        self,
        scope: Construct,
        id: str,
        entry: str,
        root: str,
        runtime: aws_lambda.Runtime,
        layers: Sequence[aws_lambda.ILayerVersion],
        function_name: str,
        handler: str = "handler",
        handler_filename: str = "handler.py",
        environment: Optional[Mapping[str, str]] = None,
        permissions: Sequence[Callable[[aws_iam.IGrantable], aws_iam.Grant]] = [],
    ) -> None:
        super().__init__(scope, id)

        stack = Stack.of(self)

        current_dir = "."
        code = aws_lambda.Code.from_asset(
            path=current_dir,
            bundling=BundlingOptions(
                image=runtime.bundling_image,
                command=[
                    "bash",
                    "-c",
                    f"rsync -r {root} /asset-output",
                ],
            ),
        )

        lambda_policy = aws_iam.ManagedPolicy(self, 'lambda-managed-policy',
            managed_policy_name=f"{function_name}-managed-policy",
            statements=[aws_iam.PolicyStatement(
                actions=[
                    "logs:CreateLogGroup",
                    "logs:CreateLogStream",
                    "logs:PutLogEvents",
                ],
                resources=[
                    f"arn:aws:logs:{stack.region}:{stack.account}:log-group:{function_name}",
                    f"arn:aws:logs:{stack.region}:{stack.account}:log-group:{function_name}:log-stream:*"
                ],
            ), aws_iam.PolicyStatement(
                actions=[
                    "xray:PutTraceSegments",
                    "xray:PutTelemetryRecords",
                ],
                resources=['*'],
            )],
        )

        cdk_nag.NagSuppressions.add_resource_suppressions(
            construct=lambda_policy,
            suppressions=[
                cdk_nag.NagPackSuppression(
                    id="AwsSolutions-IAM5",
                    reason="Log stream IDs are autogenerated.",
                    applies_to=["Resource::arn:aws:logs:<AWS::Region>:<AWS::AccountId>:log-group:simple-crud-api:log-stream:*"]
                ),
                cdk_nag.NagPackSuppression(
                    id="AwsSolutions-IAM5",
                    reason="X-Ray actions have no resources for the policy.",
                    applies_to=["Resource::*"]
                ),
            ],
        )

        lambda_role = aws_iam.Role(self, 'lambda-role',
            assumed_by=aws_iam.ServicePrincipal("lambda.amazonaws.com"),
            role_name=f"{function_name}-role",
            managed_policies=[lambda_policy],
        )

        self._func = aws_lambda.Function(
            self,
            "AppProjectFunction",
            code=code,
            runtime=runtime,
            handler=f"{entry}/{os.path.splitext(handler_filename)[0]}.{handler}",
            layers=layers,
            function_name=function_name,
            environment=environment,
            tracing=aws_lambda.Tracing.ACTIVE,
            role=lambda_role
        )

        default_policy = [p for p in lambda_role.node.children if isinstance(p, aws_iam.Policy)][0]
        cdk_nag.NagSuppressions.add_resource_suppressions(
            construct=default_policy,
            suppressions=[
                cdk_nag.NagPackSuppression(
                    id="AwsSolutions-IAM5",
                    reason="X-Ray actions have no resources for the policy.",
                    applies_to=["Resource::*"]
                ),
            ],
        )

        for p in permissions:
            p(self._func)

    @property
    def function(self) -> aws_lambda.Function:
        return self._func