import fs from "fs";
import url from "url";
import path from "path";
import crypto from "crypto";
import spawn from "cross-spawn";
import { Construct } from "constructs";
import { Effect, Role, Policy, PolicyStatement, CompositePrincipal, ServicePrincipal, ManagedPolicy, } from "aws-cdk-lib/aws-iam";
import { Version, Code, Runtime, Function as CdkFunction, } from "aws-cdk-lib/aws-lambda";
import { Asset } from "aws-cdk-lib/aws-s3-assets";
import { Lazy, Duration as CdkDuration, CustomResource, } from "aws-cdk-lib/core";
import { useProject } from "../project.js";
import { useRuntimeHandlers } from "../runtime/handlers.js";
import { Stack } from "./Stack.js";
import { useFunctions } from "./Function.js";
import { bindEnvironment, bindPermissions, getReferencedSecrets, } from "./util/functionBinding.js";
import { toCdkSize } from "./util/size.js";
import { toCdkDuration } from "./util/duration.js";
import { attachPermissionsToRole } from "./util/permission.js";
const __dirname = path.dirname(url.fileURLToPath(import.meta.url));
/////////////////////
// Construct
/////////////////////
export class EdgeFunction extends Construct {
    role;
    functionArn;
    function;
    functionCR;
    assetReplacer;
    assetReplacerPolicy;
    scope;
    versionId;
    bindingEnvs;
    props;
    constructor(scope, id, props) {
        super(scope, id);
        // Override scope
        // note: this is intended to be used internally by SST to make constructs
        //       backwards compatible when the hirechical structure of the constructs
        //       changes. When the hirerchical structure changes, the child AWS
        //       resources' logical ID will change. And CloudFormation will recreate
        //       them.
        this.scope = props.scopeOverride || this;
        this.props = {
            ...props,
            environment: props.environment || {},
            permissions: props.permissions || [],
        };
        // Bind first b/e function's environment variables cannot be added after
        this.bindingEnvs = {};
        this.bind(props.bind || []);
        // Create function with placeholder code
        const assetBucket = "placeholder";
        const assetKey = "placeholder";
        const handlerFilename = "placeholder";
        const { assetReplacer, assetReplacerPolicy } = this.createCodeReplacer(assetBucket, assetKey, handlerFilename);
        this.role = this.createRole();
        const lambdaBucket = this.createSingletonBucketInUsEast1();
        const { fn, fnArn } = this.createFunctionInUsEast1(assetBucket, assetKey, lambdaBucket);
        const { versionId } = this.createVersionInUsEast1(fn, fnArn);
        fn.node.addDependency(assetReplacer);
        this.function = CdkFunction.fromFunctionAttributes(this.scope, "ICdkFunction", {
            functionArn: fnArn,
            role: this.role,
        });
        this.functionCR = fn;
        this.functionArn = fnArn;
        this.versionId = versionId;
        this.assetReplacer = assetReplacer;
        this.assetReplacerPolicy = assetReplacerPolicy;
    }
    get currentVersion() {
        return Version.fromVersionArn(this, `${this.node.id}FunctionVersion`, `${this.functionArn}:${this.versionId}`);
    }
    async build() {
        const { bundle, handler } = this.props;
        const { asset, handlerFilename } = bundle
            ? await this.buildAssetFromBundle(bundle, handler)
            : await this.buildAssetFromHandler();
        this.updateCodeReplacer(asset.s3BucketName, asset.s3ObjectKey, handlerFilename);
        this.updateFunctionInUsEast1(asset.s3BucketName, asset.s3ObjectKey);
    }
    attachPermissions(permissions) {
        attachPermissionsToRole(this.role, permissions);
    }
    addEnvironment(key, value) {
        // Note: addEnvironment currently only updates AssetReplacer's
        //       "_SST_FUNCTION_ENVIRONMENT_" replacements
        this.props.environment[key] = value;
        const cfnReplacer = this.assetReplacer.node
            .defaultChild;
        cfnReplacer.addPropertyOverride("replacements.0.replace", JSON.stringify({
            ...this.props.environment,
            ...this.bindingEnvs,
        }));
    }
    async buildAssetFromHandler() {
        const { nodejs } = this.props;
        useFunctions().add(this.node.addr, {
            ...this.props,
            nodejs: {
                ...nodejs,
                banner: [
                    `process.env = { ...process.env, ..."{{ _SST_FUNCTION_ENVIRONMENT_ }}" };`,
                    nodejs?.banner || "",
                ].join("\n"),
            },
        });
        // Build function
        const bundle = await useRuntimeHandlers().build(this.node.addr, "deploy");
        // create wrapper that calls the handler
        if (bundle.type === "error")
            throw new Error([
                `There was a problem bundling the SSR function for the "${this.scope.node.id}" Site.`,
                ...bundle.errors,
            ].join("\n"));
        const asset = new Asset(this.scope, `FunctionAsset`, {
            path: bundle.out,
        });
        // Get handler filename
        const isESM = (nodejs?.format || "esm") === "esm";
        const parsed = path.parse(bundle.handler);
        const handlerFilename = `${parsed.dir}/${parsed.name}${isESM ? ".mjs" : ".cjs"}`;
        return { asset, handlerFilename };
    }
    async buildAssetFromBundle(bundle, handler) {
        // We expose an environment variable token which is used by the code
        // replacer to inject the environment variables assigned to the
        // EdgeFunction construct.
        //
        // "{{ _SST_FUNCTION_ENVIRONMENT_ }}" will get replaced during
        // deployment with an object of environment key-value pairs, ie.
        // const environment = {"API_URL": "https://api.example.com"};
        //
        // This inlining strategy is required as Lambda@Edge doesn't natively
        // support runtime environment variables. A downside of this approach
        // is that environment variables cannot be toggled after deployment,
        // each change to one requires a redeployment.
        const { dir: inputPath, name: inputFilename, ext: inputHandlerFunction, } = path.parse(handler);
        const inputFileExt = this.getHandlerExtension(path.join(bundle, inputPath, inputFilename));
        const handlerFilename = handler.replace(inputHandlerFunction, inputFileExt);
        const filePath = path.join(bundle, handlerFilename);
        const fileData = fs.readFileSync(filePath, "utf8");
        fs.writeFileSync(filePath, `process.env = { ...process.env, ..."{{ _SST_FUNCTION_ENVIRONMENT_ }}" };\n${fileData}`);
        // Note: cannot point the bundle to the `.open-next/server-function`
        //       b/c the folder contains node_modules. And pnpm node_modules
        //       contains symlinks. CDK cannot zip symlinks correctly.
        //       https://github.com/aws/aws-cdk/issues/9251
        //       We will zip the folder ourselves.
        const outputPath = path.resolve(useProject().paths.artifacts, `SsrFunction-${this.node.id}-${this.node.addr}`);
        const script = path.resolve(__dirname, "../support/ssr-site-function-archiver.mjs");
        const result = spawn.sync("node", [script, path.join(bundle), path.join(outputPath, "server-function.zip")], { stdio: "inherit" });
        if (result.status !== 0) {
            throw new Error(`There was a problem generating the assets package.`);
        }
        // Create asset
        const asset = new Asset(this.scope, `FunctionAsset`, {
            path: path.join(outputPath, "server-function.zip"),
        });
        return { handlerFilename, asset };
    }
    bind(constructs) {
        const app = this.node.root;
        this.bindingEnvs = {
            SST_APP: app.name,
            SST_STAGE: app.stage,
            SST_REGION: app.region,
            SST_SSM_PREFIX: useProject().config.ssmPrefix,
        };
        // Get referenced secrets
        const referencedSecrets = [];
        constructs.forEach((c) => referencedSecrets.push(...getReferencedSecrets(c)));
        [...constructs, ...referencedSecrets].forEach((c) => {
            // Bind environment
            this.bindingEnvs = {
                ...this.bindingEnvs,
                ...bindEnvironment(c),
            };
            // Bind permissions
            if (this.props.permissions !== "*") {
                this.props.permissions.push(...Object.entries(bindPermissions(c)).map(([action, resources]) => new PolicyStatement({
                    actions: [action],
                    effect: Effect.ALLOW,
                    resources,
                })));
            }
        });
    }
    createCodeReplacer(assetBucket, assetKey, handlerFilename) {
        const { environment } = this.props;
        const stack = Stack.of(this);
        // Note: Source code for the Lambda functions have "{{ ENV_KEY }}" in them.
        //       They need to be replaced with real values before the Lambda
        //       functions get deployed.
        const replacements = [
            {
                files: handlerFilename,
                search: '"{{ _SST_FUNCTION_ENVIRONMENT_ }}"',
                replace: JSON.stringify({
                    ...environment,
                    ...this.bindingEnvs,
                }),
            },
            ...Object.entries(environment).map(([key, value]) => ({
                files: "**/*.@(*js|json|html)",
                search: `{{ ${key} }}`,
                replace: value,
            })),
        ];
        const policy = new Policy(this, "AssetReplacerPolicy", {
            statements: [
                new PolicyStatement({
                    effect: Effect.ALLOW,
                    actions: ["s3:GetObject", "s3:PutObject"],
                    resources: [`arn:${stack.partition}:s3:::${assetBucket}/*`],
                }),
            ],
        });
        stack.customResourceHandler.role?.attachInlinePolicy(policy);
        const resource = new CustomResource(this.scope, "AssetReplacer", {
            serviceToken: stack.customResourceHandler.functionArn,
            resourceType: "Custom::AssetReplacer",
            properties: {
                bucket: assetBucket,
                key: assetKey,
                replacements,
            },
        });
        resource.node.addDependency(policy);
        return { assetReplacer: resource, assetReplacerPolicy: policy };
    }
    updateCodeReplacer(assetBucket, assetKey, handlerFilename) {
        const stack = Stack.of(this);
        const cfnReplacer = this.assetReplacer.node
            .defaultChild;
        cfnReplacer.addPropertyOverride("bucket", assetBucket);
        cfnReplacer.addPropertyOverride("key", assetKey);
        cfnReplacer.addPropertyOverride("replacements.0.files", handlerFilename);
        const cfnPolicy = this.assetReplacerPolicy.node.defaultChild;
        cfnPolicy.addPropertyOverride("PolicyDocument.Statement.0.Resource", `arn:${stack.partition}:s3:::${assetBucket}/*`);
    }
    createRole() {
        const { permissions } = this.props;
        // Create function role
        const role = new Role(this.scope, `ServerLambdaRole`, {
            assumedBy: new CompositePrincipal(new ServicePrincipal("lambda.amazonaws.com"), new ServicePrincipal("edgelambda.amazonaws.com")),
            managedPolicies: [
                ManagedPolicy.fromManagedPolicyArn(this, "EdgeLambdaPolicy", `arn:${Stack.of(this).partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole`),
            ],
        });
        // Attach permission
        if (permissions) {
            attachPermissionsToRole(role, permissions);
        }
        return role;
    }
    createSingletonBucketInUsEast1() {
        // Create a S3 bucket in us-east-1 to store Lambda code. Create
        // 1 bucket for all Edge functions.
        // Do not recreate if exist
        const providerId = "EdgeLambdaBucketProvider";
        const resId = "EdgeLambdaBucket";
        const stack = Stack.of(this);
        const existingResource = stack.node.tryFindChild(resId);
        if (existingResource) {
            return existingResource;
        }
        // Create provider
        const provider = new CdkFunction(stack, providerId, {
            code: Code.fromAsset(path.join(__dirname, "../support/edge-function")),
            handler: "s3-bucket.handler",
            runtime: Runtime.NODEJS_16_X,
            timeout: CdkDuration.minutes(15),
            memorySize: 1024,
            initialPolicy: [
                new PolicyStatement({
                    effect: Effect.ALLOW,
                    actions: ["s3:*"],
                    resources: ["*"],
                }),
            ],
        });
        // Create custom resource
        const resource = new CustomResource(stack, resId, {
            serviceToken: provider.functionArn,
            resourceType: "Custom::SSTEdgeLambdaBucket",
            properties: {
                BucketNamePrefix: `${stack.stackName}-${resId}`,
            },
        });
        return resource;
    }
    createFunctionInUsEast1(assetBucket, assetKey, lambdaBucket) {
        const { handler, runtime, timeout, memorySize } = this.props;
        // Do not recreate if exist
        const providerId = "EdgeLambdaProvider";
        const resId = `${this.node.id}EdgeLambda`;
        const stack = Stack.of(this);
        let provider = stack.node.tryFindChild(providerId);
        // Create provider if not already created
        if (!provider) {
            provider = new CdkFunction(stack, providerId, {
                code: Code.fromAsset(path.join(__dirname, "../support/edge-function")),
                handler: "edge-lambda.handler",
                runtime: Runtime.NODEJS_16_X,
                timeout: CdkDuration.minutes(15),
                memorySize: 1024,
                initialPolicy: [
                    new PolicyStatement({
                        effect: Effect.ALLOW,
                        actions: ["lambda:*", "s3:*"],
                        resources: ["*"],
                    }),
                ],
            });
            if (provider.role) {
                this.role.grantPassRole(provider.role);
            }
        }
        // Create custom resource
        const fn = new CustomResource(this.scope, resId, {
            serviceToken: provider.functionArn,
            resourceType: "Custom::SSTEdgeLambda",
            properties: {
                FunctionNamePrefix: `${Stack.of(this).stackName}-${resId}`,
                FunctionBucket: lambdaBucket.getAttString("BucketName"),
                FunctionParams: {
                    Description: `${this.node.id} handler`,
                    Handler: handler,
                    Code: {
                        S3Bucket: assetBucket,
                        S3Key: assetKey,
                    },
                    Runtime: runtime === "nodejs14.x"
                        ? Runtime.NODEJS_14_X.name
                        : runtime === "nodejs16.x"
                            ? Runtime.NODEJS_16_X.name
                            : Runtime.NODEJS_18_X.name,
                    MemorySize: typeof memorySize === "string"
                        ? toCdkSize(memorySize).toMebibytes()
                        : memorySize,
                    Timeout: typeof timeout === "string"
                        ? toCdkDuration(timeout).toSeconds()
                        : timeout,
                    Role: this.role.roleArn,
                },
            },
        });
        return { fn, fnArn: fn.getAttString("FunctionArn") };
    }
    updateFunctionInUsEast1(assetBucket, assetKey) {
        const cfnLambda = this.functionCR.node.defaultChild;
        cfnLambda.addPropertyOverride("FunctionParams.Code", {
            S3Bucket: assetBucket,
            S3Key: assetKey,
        });
    }
    createVersionInUsEast1(fn, fnArn) {
        // Do not recreate if exist
        const providerId = "EdgeLambdaVersionProvider";
        const resId = `${this.node.id}EdgeLambdaVersion`;
        const stack = Stack.of(this);
        let provider = stack.node.tryFindChild(providerId);
        // Create provider if not already created
        if (!provider) {
            provider = new CdkFunction(stack, providerId, {
                code: Code.fromAsset(path.join(__dirname, "../support/edge-function")),
                handler: "edge-lambda-version.handler",
                runtime: Runtime.NODEJS_16_X,
                timeout: CdkDuration.minutes(15),
                memorySize: 1024,
                initialPolicy: [
                    new PolicyStatement({
                        effect: Effect.ALLOW,
                        actions: ["lambda:*"],
                        resources: ["*"],
                    }),
                ],
            });
        }
        // Create custom resource
        const version = new CustomResource(this.scope, resId, {
            serviceToken: provider.functionArn,
            resourceType: "Custom::SSTEdgeLambdaVersion",
            properties: {
                FunctionArn: fnArn,
            },
        });
        // Override the version's logical ID with a lazy string which includes the
        // hash of the function itself, so a new version resource is created when
        // the function configuration changes.
        const cfn = version.node.defaultChild;
        const originalLogicalId = Stack.of(version).resolve(cfn.logicalId);
        cfn.overrideLogicalId(Lazy.uncachedString({
            produce: () => {
                const hash = this.calculateHash(fn);
                const logicalId = this.trimFromStart(originalLogicalId, 255 - 32);
                return `${logicalId}${hash}`;
            },
        }));
        return { version, versionId: version.getAttString("Version") };
    }
    getHandlerExtension(pathWithoutExtension) {
        const ext = [
            ".ts",
            ".tsx",
            ".mts",
            ".cts",
            ".js",
            ".jsx",
            ".mjs",
            ".cjs",
        ].find((ext) => fs.existsSync(pathWithoutExtension + ext));
        if (!ext) {
            throw new Error(`Cannot find the SSR function handler file for the "${this.scope.node.id}" Site.`);
        }
        return ext;
    }
    trimFromStart(s, maxLength) {
        const desiredLength = Math.min(maxLength, s.length);
        const newStart = s.length - desiredLength;
        return s.substring(newStart);
    }
    calculateHash(resource) {
        // render the cloudformation resource from this function
        // config is of the shape:
        // {
        //  Resources: {
        //    LogicalId: {
        //      Type: 'Function',
        //      Properties: { ... }
        // }}}
        const cfnResource = resource.node.defaultChild;
        const config = Stack.of(resource).resolve(cfnResource._toCloudFormation());
        const resources = config.Resources;
        const resourceKeys = Object.keys(resources);
        if (resourceKeys.length !== 1) {
            throw new Error(`Expected one rendered CloudFormation resource but found ${resourceKeys.length}`);
        }
        const logicalId = resourceKeys[0];
        const properties = resources[logicalId].Properties.FunctionParams;
        const hash = crypto.createHash("md5");
        hash.update(JSON.stringify(properties));
        return hash.digest("hex");
    }
}
