import { Construct } from "constructs";
import { readFileSync } from "fs";
import path from "path";
import { BlockPublicAccess, Bucket } from "aws-cdk-lib/aws-s3";
import { BucketDeployment, Source } from "aws-cdk-lib/aws-s3-deployment";
import { CustomResource, Duration, Fn, RemovalPolicy, Stack } from "aws-cdk-lib/core";
import {
  AllowedMethods,
  BehaviorOptions,
  CacheCookieBehavior,
  CacheHeaderBehavior,
  CachePolicy,
  CacheQueryStringBehavior,
  CachedMethods,
  Distribution,
  ICachePolicy,
  ViewerProtocolPolicy,
  FunctionEventType,
  OriginRequestPolicy,
  Function as CloudfrontFunction,
  FunctionCode,
} from "aws-cdk-lib/aws-cloudfront";
import { HttpOrigin, S3Origin } from "aws-cdk-lib/aws-cloudfront-origins";
import {
  Code,
  Function as CdkFunction,
  FunctionUrlAuthType,
  InvokeMode,
  Runtime,
} from "aws-cdk-lib/aws-lambda";
import { TableV2 as Table, AttributeType, Billing } from "aws-cdk-lib/aws-dynamodb";
import { Service, Source as AppRunnerSource, Memory, HealthCheck, Cpu } from "@aws-cdk/aws-apprunner-alpha";
import { DockerImageAsset } from "aws-cdk-lib/aws-ecr-assets";
import { Queue } from "aws-cdk-lib/aws-sqs";
import { SqsEventSource } from "aws-cdk-lib/aws-lambda-event-sources";
import { IGrantable } from "aws-cdk-lib/aws-iam";
import { Provider } from "aws-cdk-lib/custom-resources";
import { RetentionDays } from "aws-cdk-lib/aws-logs";
 
type BaseFunction = {
  handler: string;
  bundle: string;
};
 
type OpenNextFunctionOrigin = {
  type: "function";
  streaming?: boolean;
} & BaseFunction;
 
type OpenNextECSOrigin = {
  type: "ecs";
  bundle: string;
  dockerfile: string;
};
 
type OpenNextS3Origin = {
  type: "s3";
  originPath: string;
  copy: {
    from: string;
    to: string;
    cached: boolean;
    versionedSubDir?: string;
  }[];
};
 
type OpenNextOrigins = OpenNextFunctionOrigin | OpenNextECSOrigin | OpenNextS3Origin;
 
interface OpenNextOutput {
  edgeFunctions: {
  };
  origins: {
    s3: OpenNextS3Origin;
    default: OpenNextFunctionOrigin | OpenNextECSOrigin;
    imageOptimizer: OpenNextFunctionOrigin | OpenNextECSOrigin;
  };
  behaviors: {
    pattern: string;
    origin?: string;
    edgeFunction?: string;
  }[];
  additionalProps?: {
    disableIncrementalCache?: boolean;
    disableTagCache?: boolean;
    initializationFunction?: BaseFunction;
    warmer?: BaseFunction;
    revalidationFunction?: BaseFunction;
  };
}
 
interface OpenNextCdkReferenceImplementationProps {
  openNextPath: string;
}
 
export class OpenNextCdkReferenceImplementation extends Construct {
  private openNextOutput: OpenNextOutput;
  private bucket: Bucket;
  private table: Table;
  private queue: Queue;
 
  private staticCachePolicy: ICachePolicy;
  private serverCachePolicy: CachePolicy;
 
  public distribution: Distribution;
 
  constructor(scope: Construct, id: string, props: OpenNextCdkReferenceImplementationProps) {
    super(scope, id);
    this.openNextOutput = JSON.parse(
      readFileSync(path.join(props.openNextPath, "open-next.output.json"), "utf-8")
    ) as OpenNextOutput;
 
    this.bucket = new Bucket(this, "OpenNextBucket", {
      publicReadAccess: false,
      blockPublicAccess: BlockPublicAccess.BLOCK_ALL,
      autoDeleteObjects: true,
      removalPolicy: RemovalPolicy.DESTROY,
      enforceSSL: true,
    });
    this.table = this.createRevalidationTable();
    this.queue = this.createRevalidationQueue();
 
    const origins = this.createOrigins();
    this.serverCachePolicy = this.createServerCachePolicy();
    this.staticCachePolicy = this.createStaticCachePolicy();
    this.distribution = this.createDistribution(origins);
  }
 
  private createRevalidationTable() {
    const table = new Table(this, "RevalidationTable", {
      partitionKey: { name: "tag", type: AttributeType.STRING },
      sortKey: { name: "path", type: AttributeType.STRING },
      pointInTimeRecovery: true,
      billing: Billing.onDemand(),
      globalSecondaryIndexes: [
        {
          indexName: "revalidate",
          partitionKey: { name: "path", type: AttributeType.STRING },
          sortKey: { name: "revalidatedAt", type: AttributeType.NUMBER },
        },
      ],
      removalPolicy: RemovalPolicy.DESTROY,
    });
 
    const initFn = this.openNextOutput.additionalProps?.initializationFunction;
 
    const insertFn = new CdkFunction(this, "RevalidationInsertFunction", {
      description: "Next.js revalidation data insert",
      handler: initFn?.handler ?? "index.handler",
      // code: Code.fromAsset(initFn?.bundle ?? ""),
      code: Code.fromAsset(".open-next/dynamodb-provider"),
      runtime: Runtime.NODEJS_18_X,
      timeout: Duration.minutes(15),
      memorySize: 128,
      environment: {
        CACHE_DYNAMO_TABLE: table.tableName,
      },
    });
 
    const provider = new Provider(this, "RevalidationProvider", {
      onEventHandler: insertFn,
      logRetention: RetentionDays.ONE_DAY,
    });
 
    new CustomResource(this, "RevalidationResource", {
      serviceToken: provider.serviceToken,
      properties: {
        version: Date.now().toString(),
      },
    });
 
    return table;
  }
 
  private createOrigins() {
    const {
      s3: s3Origin,
      default: defaultOrigin,
      imageOptimizer: imageOrigin,
      ...restOrigins
    } = this.openNextOutput.origins;
    const s3 = new S3Origin(this.bucket, {
      originPath: s3Origin.originPath,
    });
    for (const copy of s3Origin.copy) {
      new BucketDeployment(this, `OpenNextBucketDeployment${copy.from}`, {
        sources: [Source.asset(copy.from)],
        destinationBucket: this.bucket,
        destinationKeyPrefix: copy.to,
        prune: false,
      });
    }
    const origins = {
      s3: new S3Origin(this.bucket, {
        originPath: s3Origin.originPath,
        originAccessIdentity: undefined,
      }),
      default:
        defaultOrigin.type === "function"
          ? this.createFunctionOrigin("default", defaultOrigin)
          : this.createAppRunnerOrigin("default", defaultOrigin),
      imageOptimizer:
        imageOrigin.type === "function"
          ? this.createFunctionOrigin("imageOptimizer", imageOrigin)
          : this.createAppRunnerOrigin("imageOptimizer", imageOrigin),
      ...Object.entries(restOrigins).reduce(
        (acc, [key, value]) => {
          if (value.type === "function") {
            acc[key] = this.createFunctionOrigin(key, value);
          } else if (value.type === "ecs") {
            acc[key] = this.createAppRunnerOrigin(key, value);
          }
          return acc;
        },
        {} as Record<string, HttpOrigin>
      ),
    };
    return origins;
  }
 
  private createRevalidationQueue() {
    const queue = new Queue(this, "RevalidationQueue", {
      fifo: true,
      receiveMessageWaitTime: Duration.seconds(20),
    });
    const consumer = new CdkFunction(this, "RevalidationFunction", {
      description: "Next.js revalidator",
      handler: "index.handler",
      code: Code.fromAsset(this.openNextOutput.additionalProps?.revalidationFunction?.bundle ?? ""),
      runtime: Runtime.NODEJS_18_X,
      timeout: Duration.seconds(30),
    });
    consumer.addEventSource(new SqsEventSource(queue, { batchSize: 5 }));
    return queue;
  }
 
  private getEnvironment() {
    return {
      CACHE_BUCKET_NAME: this.bucket.bucketName,
      CACHE_BUCKET_KEY_PREFIX: "_cache",
      CACHE_BUCKET_REGION: Stack.of(this).region,
      REVALIDATION_QUEUE_URL: this.queue.queueUrl,
      REVALIDATION_QUEUE_REGION: Stack.of(this).region,
      CACHE_DYNAMO_TABLE: this.table.tableName,
      // Those 2 are used only for image optimizer
      BUCKET_NAME: this.bucket.bucketName,
      BUCKET_KEY_PREFIX: "_assets",
    };
  }
 
  private grantPermissions(grantable: IGrantable) {
    this.bucket.grantReadWrite(grantable);
    this.table.grantReadWriteData(grantable);
    this.queue.grantSendMessages(grantable);
  }
 
  private createFunctionOrigin(key: string, origin: OpenNextFunctionOrigin) {
    const environment = this.getEnvironment();
    const fn = new CdkFunction(this, `${key}Function`, {
      runtime: Runtime.NODEJS_18_X,
      handler: origin.handler,
      code: Code.fromAsset(origin.bundle),
      environment,
      memorySize: 1024,
    });
    const fnUrl = fn.addFunctionUrl({
      authType: FunctionUrlAuthType.NONE,
      invokeMode: origin.streaming ? InvokeMode.RESPONSE_STREAM : InvokeMode.BUFFERED,
    });
    this.grantPermissions(fn);
    return new HttpOrigin(Fn.parseDomainName(fnUrl.url));
  }
 
  // 我们使用 AppRunner 是因为它是演示新特性最简单的方式
  // 您可以使用任何其他容器服务，如 ECS、EKS、Fargate 等
  private createAppRunnerOrigin(key: string, origin: OpenNextECSOrigin): HttpOrigin {
    const imageAsset = new DockerImageAsset(this, `${key}ImageAsset`, {
      directory: origin.bundle,
      // file: origin.dockerfile,
    });
    const service = new Service(this, `${key}Service`, {
      source: AppRunnerSource.fromAsset({
        asset: imageAsset,
 
        imageConfiguration: {
          port: 3000,
          environmentVariables: this.getEnvironment(),
        },
      }),
      serviceName: key,
      autoDeploymentsEnabled: false,
      cpu: Cpu.HALF_VCPU,
      memory: Memory.ONE_GB,
      healthCheck: HealthCheck.http({
        path: "/__health",
      }),
    });
    this.grantPermissions(service);
    return new HttpOrigin(service.serviceUrl);
  }
 
  private createDistribution(origins: Record<string, HttpOrigin | S3Origin>) {
    const cloudfrontFunction = new CloudfrontFunction(this, "OpenNextCfFunction", {
      code: FunctionCode.fromInline(`
			function handler(event) {
				var request = event.request;
				request.headers["x-forwarded-host"] = request.headers.host;
				return request;
			}
			`),
    });
    const fnAssociations = [
      {
        function: cloudfrontFunction,
        eventType: FunctionEventType.VIEWER_REQUEST,
      },
    ];
 
    const distribution = new Distribution(this, "OpenNextDistribution", {
      defaultBehavior: {
        origin: origins.default,
        viewerProtocolPolicy: ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
        allowedMethods: AllowedMethods.ALLOW_GET_HEAD_OPTIONS,
        cachedMethods: CachedMethods.CACHE_GET_HEAD_OPTIONS,
        cachePolicy: this.serverCachePolicy,
        originRequestPolicy: OriginRequestPolicy.ALL_VIEWER_EXCEPT_HOST_HEADER,
        functionAssociations: fnAssociations,
      },
      additionalBehaviors: this.openNextOutput.behaviors
        .filter((b) => b.pattern !== "*")
        .reduce(
          (acc, behavior) => {
            return {
              ...acc,
                origin: behavior.origin ? origins[behavior.origin] : origins.default,
                viewerProtocolPolicy: ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
                allowedMethods: AllowedMethods.ALLOW_GET_HEAD_OPTIONS,
                cachedMethods: CachedMethods.CACHE_GET_HEAD_OPTIONS,
                cachePolicy: behavior.origin === "s3" ? this.staticCachePolicy : this.serverCachePolicy,
                originRequestPolicy:
                  behavior.origin === "s3" ? undefined : OriginRequestPolicy.ALL_VIEWER_EXCEPT_HOST_HEADER,
                functionAssociations: fnAssociations,
              },
            };
          },
          {} as Record<string, BehaviorOptions>
        ),
    });
    return distribution;
  }
 
  private createServerCachePolicy() {
    return new CachePolicy(this, "OpenNextServerCachePolicy", {
      queryStringBehavior: CacheQueryStringBehavior.all(),
      headerBehavior: CacheHeaderBehavior.allowList(
        "accept",
        "accept-encoding",
        "rsc",
        "next-router-prefetch",
        "next-router-state-tree",
        "next-url",
        "x-prerender-revalidate"
      ),
      cookieBehavior: CacheCookieBehavior.none(),
      defaultTtl: Duration.days(0),
      maxTtl: Duration.days(365),
      minTtl: Duration.days(0),
    });
  }
 
  private createStaticCachePolicy() {
    return CachePolicy.CACHING_OPTIMIZED;
  }
}