/*
 * Copyright (c) 2023-2025 - Restate Software, Inc., Restate GmbH
 *
 * This file is part of the Restate CDK Construct Library,
 * which is released under the MIT license.
 *
 * You can find a copy of the license in file LICENSE in the root
 * directory of this repository or package, or at
 * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
 */

import path from "node:path";
import { Construct } from "constructs";
import * as cdk from "aws-cdk-lib";
import * as cx_api from "aws-cdk-lib/cx-api";
import * as iam from "aws-cdk-lib/aws-iam";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as lambda_node from "aws-cdk-lib/aws-lambda-nodejs";
import * as logs from "aws-cdk-lib/aws-logs";
import * as secrets from "aws-cdk-lib/aws-secretsmanager";
import { IRestateEnvironment } from "./restate-environment";
import type { RegistrationProperties } from "./register-service-handler/index.mts";

const DEFAULT_TIMEOUT = cdk.Duration.seconds(300);

export interface ServiceRegistrationProps {
  /**
   * Secrets Manager secret ARN for the authentication token to use when calling the admin API. Takes precedence
   * over the environment's token and JSON field configuration.
   */
  authToken?: secrets.ISecret;

  /**
   * Static headers to add to every admin API request made during registration (health check, deployment
   * registration, service visibility patch, and any pruning/deletion queries). Useful for tagging requests or
   * satisfying a proxy/gateway in front of the Restate admin endpoint.
   *
   * These values are stored in plaintext in the synthesized CloudFormation template; do not use them for credentials.
   * They take precedence over headers set by the handler, regardless of casing. Do not override `Content-Type` or
   * `Accept`, because registration and cleanup rely on JSON request and response bodies.
   */
  additionalHeaders?: Record<string, string>;

  /**
   * The external invoker role that Restate can assume to execute service handlers. If left unset, it's assumed that
   * the Restate deployment has sufficient permissions to invoke the handler directly. Takes precedence over the
   * environment's invokerRole.
   */
  invokerRole?: iam.IRole;

  /**
   * Whether to skip granting the invoker role permission to invoke the service handler. The deployer by default
   * will grant the invoker role permission to invoke the handler, but you can set this to `true` to handle this
   * manually.
   */
  skipInvokeFunctionGrant?: boolean;

  /**
   * Private services are only available to other Restate services in the same environment, and are not accessible for
   * ingress-based invocation. If multiple services are exposed by the same handler, all of them will be updated.
   * Default: `false`, i.e. services will be made public and reachable via ingress by default.
   *
   * @see https://docs.restate.dev/operate/registration#private-services
   */
  private?: boolean;

  /**
   * A dummy parameter to force CloudFormation to update the deployment when the configuration changes. Useful if
   * you want to target the "latest version" of a service handler and need to force a deployment in order to trigger
   * discovery. Set this to a new value every time you want to force a service registration to happen, e.g. a timestamp.
   */
  configurationVersion?: string;

  /**
   * Accept self-signed certificates.
   */
  insecure?: boolean;

  /**
   * Specify a custom admin endpoint URL, overriding the one exposed by the target environment. You may need this if
   * the `Environment` construct is reporting a different URL from the one that the deployer can reach, e.g. if your
   * Restate service is behind a load balancer.
   */
  adminUrl?: string;

  /**
   * What to do when the handler is removed from the stack.
   * - RETAIN: Leave deployment registered in Restate (default). Use this if you want to transition to managing
   *   deployments manually, or if you want to remove the ServiceDeployer without affecting existing registrations.
   * - DESTROY: Force-remove the deployment from Restate. Use this if you are decommissioning the service.
   *
   * Default: RETAIN
   */
  removalPolicy?: cdk.RemovalPolicy;

  /**
   * Prune fully drained deployments of the same handler after each successful registration. Only removes deployments
   * that have no associated services and no pinned invocations. This helps clean up old deployment versions that
   * accumulate over time as new versions are registered.
   *
   * Default: false
   */
  pruneDrainedDeployments?: boolean;

  /**
   * Number of old drained deployment revisions to retain. Only applies if `pruneDrainedDeployments` is enabled.
   * Drained deployments beyond this limit will be removed, oldest first.
   *
   * Default: 0
   */
  revisionHistoryLimit?: number;

  /**
   * Maximum number of drained deployments to prune per registration. Limits the cleanup work done in each
   * deployment to avoid long-running operations. Only applies if `pruneDrainedDeployments` is enabled.
   *
   * Default: 10
   */
  maxPrunedPerRun?: number;

  /**
   * Force deployment registration, overwriting any existing deployment at the same endpoint. This allows
   * breaking changes such as removing service handlers.
   *
   * When enabled, both breaking schema changes and deployment overwrites are allowed. This is the most
   * permissive option but may cause issues with in-flight invocations that are pinned to the existing
   * deployment.
   *
   * | Setting               | Breaking changes | Overwrites |
   * |-----------------------|------------------|------------|
   * | `force: true`         | Allowed          | Allowed    |
   * | `breaking: true`      | Allowed          | Forbidden  |
   * | Neither               | Forbidden        | Forbidden  |
   *
   * Note: If both `force` and `breaking` are set, `force` takes precedence.
   *
   * @see breaking for a safer alternative that allows breaking changes without overwriting
   * @default false
   */
  force?: boolean;

  /**
   * Allow breaking schema changes (e.g., removing handlers, changing service types) without overwriting
   * the existing deployment. This is safer than `force` because existing deployments are preserved,
   * allowing in-flight invocations to complete on the previous version.
   *
   * Use this when evolving services that may have breaking API changes but you want to preserve existing
   * deployment versions for in-progress work.
   *
   * | Setting               | Breaking changes | Overwrites |
   * |-----------------------|------------------|------------|
   * | `force: true`         | Allowed          | Allowed    |
   * | `breaking: true`      | Allowed          | Forbidden  |
   * | Neither               | Forbidden        | Forbidden  |
   *
   * Note: If both `force` and `breaking` are set, `force` takes precedence. To use the safer `breaking`
   * behavior, set `force: false` and `breaking: true`.
   *
   * @see force for allowing both breaking changes and overwrites
   * @default false
   */
  breaking?: boolean;

  /**
   * Maximum number of admin health check attempts before the deployer gives up and reports failure to
   * CloudFormation. Defaults to a value that, combined with `healthCheckMaxBackoff`, keeps the worst-case
   * loop comfortably below the deployer Lambda's 5-minute default timeout. Increase only if you also
   * raise the deployer's Lambda `timeout` to match.
   *
   * @default 10
   */
  healthCheckRetryAttempts?: number;

  /**
   * Cap on the per-iteration backoff sleep used during admin health check retries. Without a cap, the
   * exponential backoff grows fast enough that the deployer Lambda can be killed by the runtime mid-loop,
   * leaving CloudFormation to wait for its 60-minute step timeout.
   *
   * @default Duration.seconds(20)
   */
  healthCheckMaxBackoff?: cdk.Duration;
}

/**
 * Register Lambda-backed restate services with Restate environments.
 *
 * You can reuse the same deployer to register the services exposed by multiple handlers. You can also reuse the
 * deployer to target multiple Restate environments, provided the configuration options are compatible (e.g. the Restate
 * environments it needs to communicate with are all accessible from the same VPC and Security Groups, accept the same
 * authentication token, and so on). Conversely, you can create multiple deployers in cases when you need to deploy to
 * multiple environments that require distinct configuration.
 *
 * Deployment logs are retained for 30 days by default.
 *
 * @see {register}
 */
export class ServiceDeployer extends Construct {
  /** The custom resource provider for handling "deployment" resources. */
  readonly eventHandler: lambda_node.NodejsFunction;

  private invocationPolicy?: iam.Policy;

  constructor(
    scope: Construct,
    id: string,
    /**
     * Allows the custom resource event handler properties to be overridden. The main use case for this is specifying
     * VPC and security group settings for Restate environments that require it. The event handler must be able
     * to reach the S3 API, so if the subnet has no egress, it will need an S3 VPC endpoint.
     */
    props?: Partial<
      Pick<
        lambda.FunctionProps,
        | "allowPublicSubnet"
        | "architecture"
        | "runtime"
        | "code"
        | "handler"
        | "functionName"
        | "logGroup"
        | "role"
        | "securityGroups"
        | "timeout"
        | "vpc"
        | "vpcSubnets"
      > &
        Pick<logs.LogGroupProps, "removalPolicy">
    >,
  ) {
    super(scope, id);

    this.eventHandler = new lambda.Function(this, "EventHandler", {
      functionName: props?.functionName,
      logGroup: props?.logGroup,
      description: "Restate custom registration handler",
      code: props?.code ?? cdk.aws_lambda.Code.fromAsset(path.join(__dirname, "register-service-handler")),
      handler: props?.handler ?? "entrypoint.handler",
      architecture: props?.architecture ?? lambda.Architecture.ARM_64,
      runtime: props?.runtime ?? lambda.Runtime.NODEJS_22_X,
      role: props?.role,
      memorySize: 128,
      timeout: props?.timeout ?? DEFAULT_TIMEOUT,
      ...(props?.vpc
        ? ({
            vpc: props?.vpc,
            vpcSubnets: props?.vpcSubnets,
            securityGroups: props?.securityGroups,
          } satisfies Pick<lambda.FunctionOptions, "vpc" | "vpcSubnets" | "securityGroups">)
        : {}),
      allowPublicSubnet: props?.allowPublicSubnet,
    });

    // Skip creating an explicit LogGroup when CDK manages it automatically via the
    // useCdkManagedLogGroup feature flag - both would resolve to the same name and conflict.
    const cdkManagedLogGroup = cdk.FeatureFlags.of(this).isEnabled(cx_api.USE_CDK_MANAGED_LAMBDA_LOGGROUP);
    if (!props?.logGroup && !cdkManagedLogGroup) {
      // By default, Lambda Functions have a log group with never-expiring retention policy.
      new logs.LogGroup(this, "DeploymentLogs", {
        logGroupName: `/aws/lambda/${this.eventHandler.functionName}`,
        retention: logs.RetentionDays.ONE_MONTH,
        removalPolicy: props?.removalPolicy ?? cdk.RemovalPolicy.RETAIN_ON_UPDATE_OR_DELETE,
      });
    }
  }

  /**
   * Deploy a Lambda-backed Restate handler to a given environment.
   *
   * Note that a change in the handler properties is necessary to trigger re-discovery due to how CloudFormation updates
   * work. If you deploy a fixed Lambda alias such as `$LATEST` which isn't changing on every handler code or
   * configuration update, you will want to set the `configurationVersion` property in `options` to a new value (e.g. a
   * timestamp) to ensure an update to the Restate environment is triggered on stack deployment.
   *
   * @param handler service handler - must be a specific function version, use "latest" if you don't care about explicit versioning
   * @param environment target Restate environment
   * @param options additional options; see field documentation for details
   * @see {ServiceRegistrationProps}
   */
  register(handler: lambda.IVersion, environment: IRestateEnvironment, options?: ServiceRegistrationProps) {
    this.registerServiceInternal(undefined, handler, environment, options);
  }

  /**
   * Deploy a Lambda-backed Restate handler to a given environment, ensuring that a particular service name exists.
   *
   * Note that a change in the handler properties is necessary to trigger re-discovery due to how CloudFormation updates
   * work. If you deploy a fixed Lambda alias such as `$LATEST` which isn't changing on every handler code or
   * configuration update, you will want to set the `configurationVersion` property in `options` to a new value (e.g. a
   * timestamp) to ensure an update to the Restate environment is triggered on stack deployment.
   *
   * @param serviceName the service name within Restate - as a safety mechanism, this must match the service's
   *        self-reported name during discovery; if there are multiple services, one of them must match or the
   *        deployment fails
   * @param handler service handler - must be a specific function version, use "latest" if you don't care about explicit versioning
   * @param environment target Restate environment
   * @param options additional options; see field documentation for details
   */
  deployService(
    serviceName: string,
    handler: lambda.IVersion,
    environment: IRestateEnvironment,
    options?: ServiceRegistrationProps,
  ) {
    this.registerServiceInternal(serviceName, handler, environment, options);
  }

  private registerServiceInternal(
    serviceName: string | undefined,
    handler: lambda.IVersion,
    environment: IRestateEnvironment,
    options?: ServiceRegistrationProps,
  ) {
    if (environment.authTokenJsonField !== undefined && !environment.authToken) {
      throw new Error("authTokenJsonField requires an authToken on the target environment.");
    }

    const authToken = options?.authToken ?? environment.authToken;
    // The JSON field describes the environment's secret, so do not apply it to a registration-level token override.
    const authTokenJsonField = options?.authToken === undefined ? environment.authTokenJsonField : undefined;
    authToken?.grantRead(this.eventHandler);

    const invokerRole = options?.invokerRole ?? environment.invokerRole;

    const deployment = new cdk.CustomResource(handler, "RestateServiceDeployment", {
      serviceToken: this.eventHandler.functionArn,
      resourceType: "Custom::RestateServiceDeployment",
      properties: {
        servicePath: serviceName,
        adminUrl: options?.adminUrl ?? environment.adminUrl,
        authTokenSecretArn: authToken?.secretArn,
        // Forward JSON-field extraction and extra headers only when set, to avoid CFN property diffs for existing users.
        ...(authTokenJsonField !== undefined ? { authTokenJsonField } : {}),
        ...(options?.additionalHeaders !== undefined ? { additionalHeaders: options.additionalHeaders } : {}),
        serviceLambdaArn: handler.functionArn,
        invokeRoleArn: invokerRole?.roleArn,
        removalPolicy: options?.removalPolicy === cdk.RemovalPolicy.DESTROY ? "destroy" : ("retain" as const),
        private: (options?.private ?? false).toString() as "true" | "false",
        configurationVersion:
          options?.configurationVersion ??
          (handler.functionArn.endsWith(":$LATEST") ? new Date().toISOString() : undefined),
        insecure: (options?.insecure ?? false).toString() as "true" | "false",
        pruneDrainedDeployments: (options?.pruneDrainedDeployments ?? false).toString() as "true" | "false",
        revisionHistoryLimit: options?.revisionHistoryLimit ?? 0,
        maxPrunedPerRun: options?.maxPrunedPerRun ?? 10,
        force: (options?.force ?? false).toString() as "true" | "false",
        breaking: (options?.breaking ?? false).toString() as "true" | "false",
        // Forward retry knobs only when the caller sets them, to avoid CFN property diffs for existing users.
        ...(options?.healthCheckRetryAttempts !== undefined
          ? { healthCheckRetryAttempts: options.healthCheckRetryAttempts }
          : {}),
        ...(options?.healthCheckMaxBackoff !== undefined
          ? { healthCheckMaxBackoffSeconds: options.healthCheckMaxBackoff.toSeconds() }
          : {}),
      } satisfies RegistrationProperties,
    });

    if (invokerRole && !options?.skipInvokeFunctionGrant) {
      // We create a separate policy which we'll attach to the provided invoker role. This breaks a circular cross-stack
      // dependency that would otherwise be created between the service deployer and the invoker role.
      if (!this.invocationPolicy) {
        this.invocationPolicy = new iam.Policy(this, "InvocationPolicy");
        // Despite the ARN reference above, CloudFormation sometimes tries to invoke the custom resource handler before
        // all permissions are applied. Adding an explicit dependency includes a dependency on any pending policy updates
        // defined in the same stack as the service deployer, which seems to help. Some propagation delay might still mean
        // we lean on retries in the deployer event handler in any event but this reduces the probability of failure.
        deployment.node.addDependency(this.invocationPolicy);
        this.invocationPolicy.attachToRole(invokerRole);
      }
      this.invocationPolicy.addStatements(
        new iam.PolicyStatement({
          actions: ["lambda:InvokeFunction"],
          resources: handler.lambda.resourceArnsForGrantInvoke,
        }),
      );
    }
  }
}