configSchema.ts

/**
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 Licensed under the Apache License, Version 2.0 (the "License").
 You may not use this file except in compliance with the License.
 You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
import { z } from 'zod';

import { AmiHardwareType, ApplicationProtocol, ApplicationProtocolVersion, EcsSourceType, RemovalPolicy } from './cdk';
import { RagRepositoryConfigSchema, RdsInstanceConfig } from './ragSchema';

/**
 * Custom security groups for application.
 *
 * @property {T SecurityGroup} ecsModelAlbSg - ECS model application load balancer security group.
 * @property {T SecurityGroup} restApiAlbSg - REST API application load balancer security group.
 * @property {T SecurityGroup} lambdaSg - Lambda security group.
 * @property {T SecurityGroup} liteLlmSg - litellm security group.
 * @property {T SecurityGroup} openSearchSg - OpenSearch security group used by RAG.
 * @property {T SecurityGroup} pgVectorSg - PGVector security group used by RAG.
 */
export type SecurityGroups<T> = {
    ecsModelAlbSg: T;
    restApiAlbSg: T;
    lambdaSg: T;
    liteLlmSg?: T;
    openSearchSg?: T;
    pgVectorSg?: T;
};

/**
 * Enum for different types of models.
 */
export enum ModelType {
    TEXTGEN = 'textgen',
    EMBEDDING = 'embedding',
    VIDEOGEN = 'videogen',
}

/**
 * Configuration schema for Security Group imports.
 * These values are none/small/all, meaning a user can import any number of these or none of these.
 *
 * @property {string} modelSecurityGroupId - Security Group ID.
 * @property {string} restAlbSecurityGroupId - Security Group ID
 * @property {string} lambdaSecurityGroupId - Security Group ID
 * @property {string} liteLlmDbSecurityGroupId - Security Group ID.
 * @property {string} openSearchSecurityGroupId - Security Group ID.
 * @property {string} pgVectorSecurityGroupId - Security Group ID.
 */
export const SecurityGroupConfigSchema = z.object({
    modelSecurityGroupId: z.string().startsWith('sg-'),
    restAlbSecurityGroupId: z.string().startsWith('sg-'),
    lambdaSecurityGroupId: z.string().startsWith('sg-'),
    liteLlmDbSecurityGroupId: z.string().startsWith('sg-'),
    openSearchSecurityGroupId: z.string().startsWith('sg-').optional(),
    pgVectorSecurityGroupId: z.string().startsWith('sg-').optional(),
})
    .describe('Security Group Overrides used across stacks.');

// eslint-disable-next-line @typescript-eslint/no-unused-vars
const Ec2TypeSchema = z.object({
    memory: z.number().describe('Memory in megabytes (MB)'),
    gpuCount: z.number().min(0).describe('Number of GPUs'),
    nvmePath: z.string().default('').describe('Path to NVMe drive to mount'),
    maxThroughput: z.number().describe('Maximum network throughput in gigabits per second (Gbps)'),
    vCpus: z.number().describe('Number of virtual CPUs (vCPUs)'),
}).describe('Metadata for a specific EC2 instance type.');

type Ec2Type = z.infer<typeof Ec2TypeSchema>;

/**
 * Provides access to metadata for various EC2 instance types. This class is a utility that helps to retrieve the
 * metadata of specific EC2 instance types. The metadata includes properties like memory, GPU count, NVMe path,
 * maximum throughput, and virtual CPUs.
 */
export class Ec2Metadata {
    private static instances: Record<string, Ec2Type> = {
        'm5.large': {
            memory: 8 * 1000,
            gpuCount: 0,
            nvmePath: '',
            maxThroughput: 10,
            vCpus: 2,
        },
        'm5.xlarge': {
            memory: 16 * 1000,
            gpuCount: 0,
            nvmePath: '',
            maxThroughput: 10,
            vCpus: 4,
        },
        'm5.2xlarge': {
            memory: 32 * 1000,
            gpuCount: 0,
            nvmePath: '',
            maxThroughput: 10,
            vCpus: 8,
        },
        'm5d.xlarge': {
            memory: 16 * 1000,
            gpuCount: 0,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 4,
        },
        'm5d.2xlarge': {
            memory: 32 * 1000,
            gpuCount: 0,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 8,
        },
        'g4dn.xlarge': {
            memory: 16 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 4,
        },
        'g4dn.2xlarge': {
            memory: 32 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 8,
        },
        'g4dn.4xlarge': {
            memory: 64 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 16,
        },
        'g4dn.8xlarge': {
            memory: 128 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 50,
            vCpus: 32,
        },
        'g4dn.16xlarge': {
            memory: 256 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 50,
            vCpus: 64,
        },
        'g4dn.12xlarge': {
            memory: 192 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 50,
            vCpus: 48,
        },
        'g4dn.metal': {
            memory: 384 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 100,
            vCpus: 96,
        },
        'g5.xlarge': {
            memory: 16 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 4,
        },
        'g5.2xlarge': {
            memory: 32 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 8,
        },
        'g5.4xlarge': {
            memory: 64 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 16,
        },
        'g5.8xlarge': {
            memory: 128 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 32,
        },
        'g5.16xlarge': {
            memory: 256 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 64,
        },
        'g5.12xlarge': {
            memory: 192 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 40,
            vCpus: 48,
        },
        'g5.24xlarge': {
            memory: 384 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 50,
            vCpus: 96,
        },
        'g5.48xlarge': {
            memory: 768 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 100,
            vCpus: 192,
        },
        'g6.xlarge': {
            memory: 16 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 4,
        },
        'g6.2xlarge': {
            memory: 32 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 10,
            vCpus: 8,
        },
        'g6.4xlarge': {
            memory: 64 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 16,
        },
        'g6.8xlarge': {
            memory: 128 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 32,
        },
        'g6.16xlarge': {
            memory: 256 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 64,
        },
        'g6.12xlarge': {
            memory: 192 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 40,
            vCpus: 48,
        },
        'g6.24xlarge': {
            memory: 384 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 50,
            vCpus: 96,
        },
        'g6.48xlarge': {
            memory: 768 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 100,
            vCpus: 192,
        },
        'g6e.xlarge': {
            memory: 32 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 20,
            vCpus: 4,
        },
        'g6e.2xlarge': {
            memory: 64 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 20,
            vCpus: 8,
        },
        'g6e.4xlarge': {
            memory: 128 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 20,
            vCpus: 16,
        },
        'g6e.8xlarge': {
            memory: 256 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 25,
            vCpus: 32,
        },
        'g6e.16xlarge': {
            memory: 512 * 1000,
            gpuCount: 1,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 35,
            vCpus: 64,
        },
        'g6e.12xlarge': {
            memory: 384 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 100,
            vCpus: 48,
        },
        'g6e.24xlarge': {
            memory: 768 * 1000,
            gpuCount: 4,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 200,
            vCpus: 96,
        },
        'g6e.48xlarge': {
            memory: 1536 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 400,
            vCpus: 192,
        },
        'p4d.24xlarge': {
            memory: 1152 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 400,
            vCpus: 96,
        },
        'p4de.24xlarge': {
            memory: 1152 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 400,
            vCpus: 96,
        },
        'p5.48xlarge': {
            memory: 2000 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 3200,
            vCpus: 192,
        },
        'p5e.48xlarge': {
            memory: 2000 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 3200,
            vCpus: 192,
        },
        'p5en.48xlarge': {
            memory: 2000 * 1000,
            gpuCount: 8,
            nvmePath: '/dev/nvme1n1',
            maxThroughput: 3200,
            vCpus: 192,
        },
    };

    /**
     * Getter method to access EC2 metadata. Retrieves the metadata for a specific EC2 instance type.
     *
     * @param {string} key - .describe('The key representing the EC2 instance type (e.g., 'g4dn.xlarge').')
     * @throws {Error} Throws an error if no metadata is found for the specified EC2 instance type.
     * @returns {Ec2Type} The metadata for the specified EC2 instance type.
     */
    static get (key: string): Ec2Type {
        const instance = this.instances[key];
        if (!instance) {
            throw new Error(`No EC2 type found for key: ${key}`);
        }
        return instance;
    }

    /**
     * Get EC2 instances defined with metadata.
     *
     * @returns {string[]} Array of EC2 instances.
     */
    static getValidInstanceKeys (): string[] {
        return Object.keys(this.instances);
    }
}

export const VALID_INSTANCE_KEYS = Ec2Metadata.getValidInstanceKeys() as [string, ...string[]];

const ContainerHealthCheckConfigSchema = z.object({
    command: z.array(z.string()).default(['CMD-SHELL', 'exit 0']).describe('The command to run for health checks'),
    interval: z.number().default(10).describe('The time interval between health checks, in seconds.'),
    startPeriod: z.number().default(300).describe('The time to wait before starting the first health check, in seconds. Default 600s (10 min) to allow for large model loading.'),
    timeout: z.number().default(5).describe('The maximum time allowed for each health check to complete, in seconds'),
    retries: z.number().default(3).describe('The number of times to retry a failed health check before considering the container as unhealthy.'),
})
    .describe('Configuration for container health checks');

export { ContainerHealthCheckConfigSchema };
export type ContainerHealthCheckConfig = z.infer<typeof ContainerHealthCheckConfigSchema>;

export const ImageTarballAsset = z.object({
    path: z.string(),
    type: z.literal(EcsSourceType.TARBALL)
})
    .describe('Container image that will use tarball on disk');

export const ImageSourceAsset = z.object({
    baseImage: z.string(),
    path: z.string(),
    type: z.literal(EcsSourceType.ASSET),
})
    .describe('Container image that will be built based on Dockerfile and assets at the supplied path');

export const ImageECRAsset = z.object({
    repositoryArn: z.string(),
    tag: z.string().optional(),
    type: z.literal(EcsSourceType.ECR),
})
    .describe('Container image that will be pulled from the specified ECR repository');

export const ImageRegistryAsset = z.object({
    registry: z.string(),
    type: z.literal(EcsSourceType.REGISTRY),
})
    .describe('Container image that will be pulled from the specified public registry');

export const ImageExternalAsset = z.object({
    type: z.literal(EcsSourceType.EXTERNAL),
    code: z.any()
})
    .describe('Container image from external source. Use provided image without modification.');

export const ImageAssetSchema = z.union([ImageTarballAsset, ImageSourceAsset, ImageECRAsset, ImageRegistryAsset, ImageExternalAsset]);
export type ImageAsset = z.infer<typeof ImageAssetSchema>;

export const ContainerConfigSchema = z.object({
    image: ImageAssetSchema.describe('Base image for the container.'),
    environment: z
        .record(z.string(), z.any())
        .transform((obj) => {
            return Object.entries(obj).reduce(
                (acc, [key, value]) => {
                    acc[key] = String(value);
                    return acc;
                },
                {} as Record<string, string>,
            );
        })
        .default({})
        .describe('Environment variables for the container.'),
    sharedMemorySize: z.number().min(0).default(0).describe('The value for the size of the /dev/shm volume.'),
    healthCheckConfig: ContainerHealthCheckConfigSchema.default(ContainerHealthCheckConfigSchema.parse({})),
    privileged: z.boolean().optional(),
    memoryReservation: z.number().min(0).optional().describe('Memory reservation in MiB for the container.')
}).describe('Configuration for the container.');

export type ContainerConfig = z.infer<typeof ContainerConfigSchema>;

export const LoadBalancerHealthCheckConfigSchema = z.object({
    path: z.string().describe('Path for the health check.'),
    interval: z.number().default(30).describe('Interval in seconds between health checks.'),
    timeout: z.number().default(10).describe('Timeout in seconds for each health check.'),
    healthyThresholdCount: z.number().default(2).describe('Number of consecutive successful health checks required to consider the target healthy.'),
    unhealthyThresholdCount: z.number().default(2).describe('Number of consecutive failed health checks required to consider the target unhealthy.'),
})
    .describe('Health check configuration for the load balancer.');

export type LoadBalancerHealthCheckConfig = z.infer<typeof LoadBalancerHealthCheckConfigSchema>;

export const LoadBalancerConfigSchema = z.object({
    sslCertIamArn: z.string().nullish().default(null).describe('SSL certificate IAM ARN for load balancer.'),
    healthCheckConfig: LoadBalancerHealthCheckConfigSchema,
    domainName: z.string().nullish().default(null).describe('Domain name to use instead of the load balancer\'s default DNS name.'),
})
    .describe('Configuration for load balancer settings.');

export const MetricConfigSchema = z.object({
    albMetricName: z.string().default('RequestCountPerTarget')
        .describe('ALB metric for scaling. Use "RequestCountPerTarget" for request volume scaling (good for embedding models) ' +
            'or "TargetResponseTime" for latency-based scaling (recommended for text generation LLMs). ' +
            'When using TargetResponseTime, set targetValue to the max acceptable p90 latency in seconds (e.g., 10).'),
    targetValue: z.number().default(30).describe('Target value for the metric. For RequestCountPerTarget this is requests per target. ' +
        'For TargetResponseTime this is the target p90 latency in seconds.'),
    duration: z.number().default(60).describe('Duration in seconds for metric evaluation.'),
    estimatedInstanceWarmup: z.number().min(0).max(3600).default(180).describe('Estimated warm-up time in seconds until a newly launched instance can send metrics to CloudWatch. Max 3600 (1 hour).'),
})
    .describe('Metric configuration for ECS auto scaling.');

export const AutoScalingConfigSchema = z.object({
    blockDeviceVolumeSize: z.number().min(30).default(50),
    minCapacity: z.number().min(1).default(1).describe('Minimum capacity for auto scaling. Must be at least 1.'),
    maxCapacity: z.number().min(1).default(2).describe('Maximum capacity for auto scaling. Must be at least 1.'),
    defaultInstanceWarmup: z.number().min(0).max(3600).default(180).describe('Default warm-up time in seconds until a newly launched instance can contribute to CloudWatch metrics. Max 3600 (1 hour). Larger models (e.g. gpt-oss-120b) may require values closer to the maximum.'),
    cooldown: z.number().min(1).default(420).describe('Cool down period in seconds between scaling activities.'),
    metricConfig: MetricConfigSchema,
})
    .describe('Configuration for auto scaling settings.');

const enumKeySchema = <E extends Record<string, string | number>>(e: E) => {
    const keys = Object.keys(e).filter((k) => Number.isNaN(Number(k))) as Array<Extract<keyof E, string>>;
    if (keys.length === 0) {
        throw new Error('Enum must have at least one valid key');
    }
    return z.enum(keys as unknown as [string, ...string[]]);
};

export const ApplicationTargetSchema = z.object({
    protocol: enumKeySchema(ApplicationProtocol).optional(),
    protocolVersion: enumKeySchema(ApplicationProtocolVersion).optional(),
    port: z.number().positive().optional(),
    priority: z.number().optional(),
    conditions: z.array(z.object({
        type: z.enum(['pathPatterns']),
        values: z.array(z.string())
    })).optional()
});

export type ApplicationTarget = z.infer<typeof ApplicationTargetSchema>;

export const TaskDefinitionSchema = z.object({
    containerConfig: ContainerConfigSchema,
    containerMemoryReservationMiB: z.number().default(1024 * 2)
        .describe('The amount (in MiB) of memory to present to the container.').optional(),
    memoryLimitMiB: z.number().positive().describe('The amount (in MiB) of memory to present to the container.').optional(),
    environment: z.record(z.string(), z.string()).describe('Environment variables set on the task container'),
    applicationTarget: ApplicationTargetSchema.optional().describe('How the load balancer should target the task.')
});

export type TaskDefinition = z.infer<typeof TaskDefinitionSchema>;


/**
 * Configuration schema for an ECS model.
 *
 * @property {AmiHardwareType} amiHardwareType - Name of the model.
 * @property {AutoScalingConfigSchema} autoScalingConfig - Configuration for auto scaling settings.
 * @property {Record<string,string>} buildArgs - Optional build args to be applied when creating the
 *                                              task container if containerConfig.image.type is ASSET
 * @property {ContainerConfigSchema} containerConfig - Configuration for the container.
 * @property {number} [containerMemoryBuffer=2048] - This is the amount of memory to buffer (or subtract off)
 *                                                from the total instance memory, if we don't include this,
 *                                                the container can have a hard time finding available RAM
 *                                                resources to start and the tasks will fail deployment
 * @property {Record<string,string>} environment - Environment variables set on the task container
 * @property {identifier} modelType - Unique identifier for the cluster which will be used when naming resources
 * @property {string} instanceType - EC2 instance type for running the model.
 * @property {boolean} [internetFacing=false] - Whether or not the cluster will be configured as internet facing
 * @property {LoadBalancerConfigSchema} loadBalancerConfig - Configuration for load balancer settings.
 */
export const EcsBaseConfigSchema = z.object({
    amiHardwareType: z.enum(AmiHardwareType).describe('Name of the model.'),
    amiId: z.string().optional()
        .describe('Optional AMI ID for a custom ECS machine image (e.g. ami-0123456789abcdef0). If not provided, the default ECS-optimized AMI will be used (AL2 for ADC/iso regions, AL2023 otherwise).'),
    autoScalingConfig: AutoScalingConfigSchema.describe('Configuration for auto scaling settings.'),
    buildArgs: z.record(z.string(), z.string()).optional()
        .describe('Optional build args to be applied when creating the task container if containerConfig.image.type is ASSET'),
    containerMemoryBuffer: z.number().default(1024 * 2)
        .describe('This is the amount of memory to buffer (or subtract off) from the total instance memory, ' +
            'if we don\'t include this, the container can have a hard time finding available RAM resources to start and the tasks will fail deployment'),
    tasks: z.record(z.string(), TaskDefinitionSchema),
    instanceType: z.enum(VALID_INSTANCE_KEYS).describe('EC2 instance type for running the model.'),
    internetFacing: z.boolean().default(false).describe('Whether or not the cluster will be configured as internet facing'),
    loadBalancerConfig: LoadBalancerConfigSchema
});


/**
 * Details and configurations of a registered model.
 *
 * @property {string} provider - Model provider, of the form <engine>.<type>.
 * @property {string} modelName - The unique name that identifies the model.
 * @property {string} modelId - The unique user-provided name for the model.
 * @property {ModelType} modelType - Specifies the type of model (e.g., 'textgen', 'embedding').
 * @property {string} endpointUrl - The URL endpoint where the model can be accessed or invoked.
 * @property {boolean} streaming - Indicates whether the model supports streaming capabilities.
 */
export type RegisteredModel = {
    provider: string;
    modelId: string;
    modelName: string;
    modelType: ModelType;
    endpointUrl: string;
    streaming?: boolean;
};

/**
 * Type representing configuration for an ECS model.
 */
export type ECSConfig = z.infer<typeof EcsBaseConfigSchema>;

/**
 * Configuration schema for an ECS model.
 *
 * @property {string} modelName - Name of the model.
 * @property {string} modelId - An optional short id to use when creating cdk ids for model related
 *                              resources
 * @property {boolean} [deploy=true] - Whether to deploy model.
 * @property {boolean} [streaming=null] - Whether the model supports streaming.
 * @property {string} modelType - Type of model.
 * @property {string} instanceType - EC2 instance type for running the model.
 * @property {string} inferenceContainer - Prebuilt inference container for serving model.
 * @property {ContainerConfigSchema} containerConfig - Configuration for the container.
 * @property {AutoScalingConfigSchema} autoScalingConfig - Configuration for auto scaling settings.
 * @property {LoadBalancerConfigSchema} loadBalancerConfig - Configuration for load balancer settings.
 * @property {string} [localModelCode='/opt/model-code'] - Path in container for local model code.
 * @property {string} [modelHosting='ecs'] - Model hosting.
 */
export const EcsClusterConfigSchema = z
    .object({
        modelName: z.string(),
        modelId: z.string().optional(),
        deploy: z.boolean().default(true),
        streaming: z.boolean().nullable().default(null),
        modelType: z.enum(ModelType),
        instanceType: z.enum(VALID_INSTANCE_KEYS),
        inferenceContainer: z
            .union([z.literal('tgi'), z.literal('tei'), z.literal('instructor'), z.literal('vllm')])
            .refine((data) => {
                return !data.includes('.'); // string cannot contain a period
            }),
        containerConfig: ContainerConfigSchema,
        autoScalingConfig: AutoScalingConfigSchema,
        loadBalancerConfig: LoadBalancerConfigSchema,
        localModelCode: z.string().default('/opt/model-code'),
        containerMemoryBuffer: z.number().default(1024 * 2)
            .describe('Memory in MiB to reserve for the host OS/ECS agent. Container gets (instance memory - buffer). Default: 2048 MiB'),
        modelHosting: z
            .string()
            .default('ecs')
            .refine((data) => {
                return !data.includes('.'); // string cannot contain a period
            }),
    })
    .refine(
        (data) => {
            // 'textgen' type must have boolean streaming, 'embedding' and 'videogen' types must have null streaming
            const isValidForTextgen = data.modelType === 'textgen' && typeof data.streaming === 'boolean';
            const isValidForEmbedding = data.modelType === 'embedding' && data.streaming === null;
            const isValidForVideogen = data.modelType === 'videogen' && data.streaming === null;

            return isValidForTextgen || isValidForEmbedding || isValidForVideogen;
        },
        {
            message: `For 'textgen' models, 'streaming' must be true or false.
            For 'embedding' and 'videogen' models, 'streaming' must not be set.`,
            path: ['streaming'],
        },
    );

/**
 * Type representing configuration for an ECS model.
 */
export type EcsClusterConfig = z.infer<typeof EcsClusterConfigSchema>;

const InferenceContainer = {
    tgi: 'tgi',
    tei: 'tei',
    instructor: 'instructor',
    vllm: 'vllm',
} as const;

const InferenceContainerSchema = z.object({
    tgi: z.literal(InferenceContainer.tgi).describe('Text Generation Inference'),
    tei: z.literal(InferenceContainer.tei).describe('Text Embeddings Inference'),
    instructor: z.literal(InferenceContainer.instructor).describe('Instructor Library'),
    vllm: z.literal(InferenceContainer.vllm).describe('vLLM (Virtual Large Language Model) Library'),
});

const EcsModelConfigSchema = z
    .object({
        modelName: z.string().describe('Name of the model.'),
        baseImage: z.string().describe('Base image for the container.'),
        inferenceContainer: z
            .union([InferenceContainerSchema.shape.tgi, InferenceContainerSchema.shape.tei, InferenceContainerSchema.shape.instructor, InferenceContainerSchema.shape.vllm])
            .refine((data) => {
                return !data.includes('.'); // string cannot contain a period
            })
            .describe('Prebuilt inference container for serving model.'),
    })
    .describe('Configuration schema for an ECS model.');

/**
 * Type representing configuration for an ECS model.
 */
type EcsModelConfig = z.infer<typeof EcsModelConfigSchema>;

/**
 * Union model type representing various model configurations.
 */
export type ModelConfig = EcsModelConfig;

const AuthConfigSchema = z.object({
    authority: z.string().transform((value) => {
        if (value.endsWith('/')) {
            return value.slice(0, -1);
        } else {
            return value;
        }
    })
        .describe('URL of OIDC authority.'),
    clientId: z.string().describe('Client ID for OIDC IDP .'),
    adminGroup: z.string().default('').describe('Name of the admin group.'),
    userGroup: z.string().default('').describe('Name of the user group.'),
    apiGroup: z.string().default('').describe('Name of the API group for API token access.'),
    ragAdminGroup: z.string().default('').describe('Name of the RAG admin group for RAG management access.'),
    jwtGroupsProperty: z.string().default('').describe('Name of the JWT groups property.'),
    additionalScopes: z.array(z.string()).default([]).describe('Additional JWT scopes to request.'),
}).describe('Configuration schema for authorization.');

const FastApiContainerConfigSchema = z.object({
    internetFacing: z.boolean().default(true).describe('Whether the REST API ALB will be configured as internet facing.'),
    domainName: z.string().nullish().default(null),
    sslCertIamArn: z.string().nullish().default(null).describe('ARN of the self-signed cert to be used throughout the system'),
    imageConfig: ImageAssetSchema.optional().describe('Override image configuration for ECS FastAPI Containers'),
    buildConfig: z.object({
        PRISMA_CACHE_DIR: z.string().optional().describe('Override with a path relative to the build directory for a pre-cached prisma directory. Defaults to PRISMA_CACHE.')
    }).default({}),
    rdsConfig: RdsInstanceConfig
        .default({
            dbName: 'postgres',
            username: 'postgres',
            dbPort: 5432,
        })
        .refine(
            (config) => {
                return !config.dbHost;
            },
            {
                message:
                    'We do not allow using an existing DB for LiteLLM because of its requirement in internal model management ' +
                    'APIs. Please do not define the dbHost field for the FastAPI container DB config.',
            },
        ),
    rateLimitRpm: z.number().int().positive().default(60).describe(
        'Per-user sustained request rate limit in requests per minute. Each ECS task enforces this independently.'
    ),
    rateLimitBurst: z.number().int().nonnegative().default(10).describe(
        'Per-user burst allowance above the sustained rate limit. Allows short spikes without throttling.'
    ),
    rateLimitEnabled: z.boolean().default(true).describe(
        'Enable or disable per-user rate limiting on the serve API.'
    ),
    rateLimitOverrides: z.record(
        z.string(),
        z.object({
            rpm: z.number().int().positive().optional().describe('Override RPM for this user/token.'),
            burst: z.number().int().nonnegative().optional().describe('Override burst for this user/token.'),
        })
    ).default({}).describe(
        'Per-user rate limit overrides. Keys use the format "token:<tokenUUID>" for API tokens, ' +
        '"oidc:<sub>" for OIDC users, or "user:<username>". Values can override "rpm" and/or "burst".'
    ),
}).describe('Configuration schema for REST API.');

/** Custom domain / TLS for the MCP Workbench ALB only (separate from Serve’s `restApiConfig`). */
const McpWorkbenchRestApiConfigSchema = z
    .object({
        domainName: z
            .string()
            .nullish()
            .default(null)
            .describe(
                'Hostname for the MCP Workbench ALB (HTTPS listener and SSM …/mcpWorkbench/endpoint). Configure here for the same YAML shape as `restApiConfig.domainName` for LISA Serve.',
            ),
        sslCertIamArn: z
            .string()
            .nullish()
            .default(null)
            .describe(
                'ACM certificate ARN for the MCP Workbench ALB. Same role as `restApiConfig.sslCertIamArn` for Serve; if omitted, falls back to `mcpWorkbenchEcsConfig.sslCertIamArn` then `restApiConfig.sslCertIamArn`.',
            ),
    })
    .describe('Optional load balancer domain and TLS for MCP Workbench (parallel to `restApiConfig` for LISA Serve).');

const RagFileProcessingConfigSchema = z.object({
    chunkSize: z.number().min(100).max(10000),
    chunkOverlap: z.number().min(0),
})
    .describe('Configuration schema for RAG file processing. Determines the chunk size and chunk overlap when processing documents.');

const PypiConfigSchema = z.object({
    indexUrl: z.string().default('').describe('URL for the pypi index.'),
    trustedHost: z.string().default('').describe('Trusted host for pypi.'),
})
    .describe('Configuration schema for pypi');

/**
 * Enum for different types of stack synthesizers
 */
export enum stackSynthesizerType {
    CliCredentialsStackSynthesizer = 'CliCredentialsStackSynthesizer',
    DefaultStackSynthesizer = 'DefaultStackSynthesizer',
    LegacyStackSynthesizer = 'LegacyStackSynthesizer',
}

const ApiGatewayConfigSchema = z
    .object({
        domainName: z.string().nullish().default(null).describe('Custom domain name for API Gateway Endpoint'),
    })
    .optional()
    .describe('Configuration schema for API Gateway Endpoint');

const AuditLoggingConfigSchema = z
    .object({
        enabled: z.boolean().default(false).describe('Whether to enable audit logging for opted-in API Gateway paths.'),
        auditAll: z.boolean().default(false).describe('If true, enable audit logging for all API Gateway paths.'),
        enabledPaths: z
            .array(z.string().min(1))
            .default([])
            .describe('API Gateway path prefixes (e.g. "/session") to include in audit logging. Prefix match is used.'),
        maxRequestBodyBytes: z
            .number()
            .int()
            .min(1)
            .default(20000)
            .describe('Maximum request body bytes to include in audit logs (oversized bodies are replaced with a placeholder).'),
        includeJsonBody: z
            .boolean()
            .default(false)
            .describe(
                'When true, emit AUDIT_API_GATEWAY_REQUEST_BODY for opted-in paths. When false (default), request bodies are never included in audit logs.'
            ),
    })
    .describe('Audit logging configuration for API Gateway request auditing.');

const LiteLLMConfig = z.object({
    db_key: z.string().refine(
        (key) => key.startsWith('sk-'), // key needed for model management actions
        'Key string must be defined for model management operations, and it must start with "sk-".' +
        'This can be any string, and a random UUID is recommended. Example: sk-f132c7cc-059c-481b-b5ca-a42e191672aa',
    ),
    general_settings: z.any().optional(),
    litellm_settings: z.any().optional(),
    router_settings: z.any().optional(),
    environment_variables: z.any().optional(),
    // LiteLLM callback-specific settings (e.g., OpenTelemetry message logging toggles).
    // This must be allowed here so Zod doesn't strip it at deploy time.
    callback_settings: z.any().optional(),
})
    .describe('Core LiteLLM configuration - see https://litellm.vercel.app/docs/proxy/configs#all-settings for more details about each field.');

const RoleConfig = z.object({
    DockerImageBuilderDeploymentRole: z.string().max(64),
    DockerImageBuilderEC2Role: z.string().max(64),
    DockerImageBuilderRole: z.string().max(64),
    DocsRole: z.string().max(64).optional(),
    DocsDeployerRole: z.string().max(64).optional(),
    ECSModelDeployerRole: z.string().max(64),
    ECSModelTaskRole: z.string().max(64),
    ECSRestApiRole: z.string().max(64),
    ECSRestApiExRole: z.string().max(64),
    LambdaExecutionRole: z.string().max(64),
    ECSMcpWorkbenchApiRole: z.string().max(64),
    ECSMcpWorkbenchApiExRole: z.string().max(64),
    LambdaConfigurationApiExecutionRole: z.string().max(64),
    ModelApiRole: z.string().max(64),
    ModelsSfnLambdaRole: z.string().max(64),
    ModelSfnRole: z.string().max(64),
    RagLambdaExecutionRole: z.string().max(64).optional(),
    RestApiAuthorizerRole: z.string().max(64),
    S3ReaderRole: z.string().max(64).optional(),
    UIDeploymentRole: z.string().max(64).optional(),
    VectorStoreCreatorRole: z.string().max(64).optional(),
    McpServerDeployerRole: z.string().max(64),
    ApiTokensApiRole: z.string().max(64),
})
    .describe('Role overrides used across stacks.');


export const RawConfigObject = z.object({
    appName: z.string().default('lisa').describe('Name of the application.'),
    profile: z
        .string()
        .nullish()
        .transform((value) => value ?? '')
        .describe('AWS CLI profile for deployment.'),
    deploymentName: z.string().default('prod').describe('Name of the deployment.'),
    accountNumber: z
        .number()
        .or(z.string())
        .transform((value) => value.toString())
        .refine((value) => value.length === 12, {
            error: 'AWS account number should be 12 digits. If your account ID starts with 0, then please surround the ID with quotation marks.'
        })
        .describe('AWS account number for deployment. Must be 12 digits.'),
    region: z.string().describe('AWS region for deployment.'),
    partition: z.string().default('aws').describe('AWS partition for deployment.'),
    domain: z.string().default('amazonaws.com').describe('AWS domain for deployment'),
    restApiConfig: FastApiContainerConfigSchema.describe('Image override for Rest API'),
    mcpWorkbenchRestApiConfig: McpWorkbenchRestApiConfigSchema.optional().describe(
        'Custom domain and certificate for the MCP Workbench ALB. Same usage as `restApiConfig.domainName` / `sslCertIamArn` for LISA Serve.',
    ),
    mcpWorkbenchConfig: ImageAssetSchema.optional().describe('Image override for MCP Workbench'),
    mcpWorkbenchBuildConfig: z.object({
        S6_OVERLAY_NOARCH_SOURCE: z.string().optional().describe('Override the URL with a path relative to the build directory for the architecture independent S6 overlay tar.xz.'),
        S6_OVERLAY_ARCH_SOURCE: z.string().optional().describe('Override the URL with a path relative to the build directory for the architecture specific S6 overlay tar.xz.'),
        RCLONE_SOURCE: z.string().optional().describe('Override the URL with a path relative to the build directory for an rclone .zip')
    }).default({}),
    batchIngestionConfig: ImageAssetSchema.optional().describe('Image override for Batch Ingestion'),
    vpcId: z.string().optional().describe('VPC ID for the application. (e.g. vpc-0123456789abcdef)'),
    subnets: z.array(z.object({
        subnetId: z.string().startsWith('subnet-'),
        ipv4CidrBlock: z.string(),
        availabilityZone: z.string().describe('Specify the availability zone for the subnet in the format {region}{letter} (e.g., us-east-1a).'),
    })).optional().describe('Array of subnet objects for the application. These contain a subnetId(e.g. [subnet-fedcba9876543210] and ipv4CidrBlock'),
    securityGroupConfig: SecurityGroupConfigSchema.optional(),
    deploymentStage: z.string().default('prod').describe('Deployment stage for the application.'),
    removalPolicy: z.enum([RemovalPolicy.DESTROY, RemovalPolicy.RETAIN])
        .default(RemovalPolicy.DESTROY)
        .describe('Removal policy for resources (destroy or retain).'),
    maxAzs: z.int().optional().default(2).describe('Number of AZs that LISA will deploy to. This must be set before deploying the network stack.'),
    runCdkNag: z.boolean().default(false).describe('Whether to run CDK Nag checks.'),
    privateEndpoints: z.boolean().default(false).describe('Whether to use privateEndpoints for REST API.'),
    s3BucketModels: z.string().describe('S3 bucket for models.'),
    mountS3DebUrl: z.string().describe('URL for S3-mounted Debian package.'),
    imageBuilderVolumeSize: z.number().default(50).describe('EC2 volume size for image builder. Needs to be large enough for system plus inference container.'),
    accountNumbersEcr: z
        .array(z.union([z.number(), z.string()]))
        .transform((arr) => arr.map(String))
        .refine((value) => value.every((num) => num.length === 12), {
            error: 'AWS account number should be 12 digits. If your account ID starts with 0, then please surround the ID with quotation marks.'
        })
        .optional()
        .describe('List of AWS account numbers for ECR repositories.'),
    deployRag: z.boolean().default(true).describe('Whether to deploy RAG stacks.'),
    deployChat: z.boolean().default(true).describe('Whether to deploy chat stacks.'),
    deployDocs: z.boolean().default(true).describe('Whether to deploy docs stacks.'),
    deployUi: z.boolean().default(true).describe('Whether to deploy UI stacks.'),
    useCustomBranding: z.boolean().optional().describe('Whether to use custom branding assets in the UI.'),
    customDisplayName: z.string().optional().describe('Custom display name to replace "LISA" branding in titles and descriptions. Requires "useCustomBranding" to be enabled.'),
    deployMetrics: z.boolean().default(true).describe('Whether to deploy Metrics stack.'),
    deployHealthDashboard: z.boolean().default(true).describe('Whether to deploy the ECS Model Health CloudWatch dashboard for monitoring model container health, errors, latency, and resource utilization.'),
    deployMcp: z.boolean().default(true).describe('Whether to deploy LISA MCP stack.'),
    deployServe: z.boolean().default(true).describe('Whether to deploy LISA Serve stack.'),
    deployMcpWorkbench: z.boolean().default(true).describe('Whether to deploy MCP Workbench stack.'),
    mcpWorkbenchEcsConfig: z
        .object({
            instanceType: z.enum(VALID_INSTANCE_KEYS).optional().describe('EC2 instance type for the MCP Workbench ECS cluster.'),
            blockDeviceVolumeSize: z.number().min(30).optional().describe('Root volume size (GiB) for cluster instances.'),
            minCapacity: z.number().min(1).optional().describe('Minimum ASG capacity for the MCP Workbench cluster.'),
            maxCapacity: z.number().min(1).optional().describe('Maximum ASG capacity for the MCP Workbench cluster.'),
            cooldown: z.number().min(1).optional().describe('Cooldown (seconds) between scaling activities.'),
            domainName: z
                .string()
                .nullish()
                .describe(
                    'Optional hostname for the MCP Workbench ALB (same effect as `mcpWorkbenchRestApiConfig.domainName`; use that block for parity with `restApiConfig`). ' +
                    'If omitted and restApiConfig.domainName is set, a default is derived (e.g. first label `lisa-serve` → `lisa-mcp-workbench`, or `serve` → `mcp-workbench`) so the workbench does not reuse the Serve API hostname. ' +
                    'Otherwise the published endpoint uses this ALB’s DNS name. You must create DNS for the chosen or derived name pointing at the MCP Workbench ALB.',
                ),
            sslCertIamArn: z
                .string()
                .nullish()
                .describe(
                    'Optional ACM certificate ARN for the MCP Workbench ALB HTTPS listener (same effect as `mcpWorkbenchRestApiConfig.sslCertIamArn`). If omitted, inherits restApiConfig.sslCertIamArn when set; ' +
                    'otherwise the workbench ALB uses HTTP on port 80 (browser MCP from an https UI will fail). Set explicitly when using a dedicated workbench hostname.',
                ),
        })
        .optional()
        .describe(
            'Optional sizing and load-balancer settings for the MCP Workbench ECS cluster. The workbench HTTP server always runs on its own ECS cluster and ALB (separate from the Serve REST API).',
        ),
    mcpWorkbenchCorsOrigins: z
        .string()
        .default('*')
        .describe(
            'Comma-separated CORS allowed origins for the MCP Workbench HTTP server container (CORS_ORIGINS). Use * to allow any browser origin (typical when the UI is served from varying hosts or ports). More restrictive deployments can list explicit origins.',
        ),
    logLevel: z.union([z.literal('DEBUG'), z.literal('INFO'), z.literal('WARNING'), z.literal('ERROR')])
        .default('DEBUG')
        .describe('Log level for application.'),
    authConfig: AuthConfigSchema.optional().describe('Authorization configuration.'),
    roles: RoleConfig.optional(),
    pypiConfig: PypiConfigSchema.default({
        indexUrl: '',
        trustedHost: '',
    }).describe('Pypi configuration.'),
    baseImage: z.string().default('public.ecr.aws/docker/library/python:3.13-slim').describe('Base image used for LISA serve components'),
    nodejsImage: z.string().default('public.ecr.aws/lambda/nodejs:24').describe('Base image used for LISA NodeJS lambda deployments'),
    condaUrl: z.string().default('').describe('Conda URL configuration'),
    certificateAuthorityBundle: z.string().default('').describe('Certificate Authority Bundle file'),
    ragRepositories: z.array(RagRepositoryConfigSchema).describe('Rag Repository configuration.'),
    ragFileProcessingConfig: RagFileProcessingConfigSchema.optional().describe('Rag file processing configuration.'),
    ecsModels: z.array(EcsModelConfigSchema).optional().describe('Array of ECS model configurations.'),
    apiGatewayConfig: ApiGatewayConfigSchema,
    nvmeHostMountPath: z.string().default('/nvme').describe('Host path for NVMe drives.'),
    nvmeContainerMountPath: z.string().default('/nvme').describe('Container path for NVMe drives.'),
    tags: z
        .array(
            z.object({
                Key: z.string(),
                Value: z.string(),
            }),
        )
        .optional()
        .describe('Array of key-value pairs for tagging.'),
    tagContainers: z.boolean().default(false).describe('Add repeatable container tags to each cdk deployed container'),
    deploymentPrefix: z.string().optional().describe('Prefix for deployment resources.'),
    webAppAssetsPath: z.string().optional().describe('Optional path to precompiled webapp assets. If not specified the web application will be built at deploy time.'),
    ecsModelDeployerPath: z.string().optional().describe('Optional path to precompiled ecs model deployer. If not specified the ecs model deployer will be built at deploy time.'),
    mcpServerDeployerPath: z.string().optional().describe('Optional path to precompiled mcp server deployer. If not specified the mcp server deployer will be built at deploy time.'),
    vectorStoreDeployerPath: z.string().optional().describe('Optional path to precompiled vector store deployer. If not specified the vector store deployer will be built at deploy time.'),
    documentsPath: z.string().optional().describe('Optional path to precompiled LISA documents. If not specified the LISA docs will be built at deploy time.'),
    lambdaPath: z.any().optional().describe('Optional path to precompiled LISA lambda. If not specified the LISA lambda will be built at deploy time.'),
    lambdaLayerAssets: z
        .object({
            authorizerLayerPath: z.string().optional().describe('Lambda Authorizer code path'),
            commonLayerPath: z.string().optional().describe('Lambda common layer code path'),
            fastapiLayerPath: z.string().optional().describe('Lambda API code path'),
            ragLayerPath: z.string().optional().describe('Lambda RAG layer code path'),
            cdkLayerPath: z.string().optional().describe('Lambda CDK layer code path'),
        })
        .optional()
        .describe('Configuration for local Lambda layer code'),
    auditLoggingConfig: AuditLoggingConfigSchema.optional(),
    permissionsBoundaryAspect: z
        .object({
            permissionsBoundaryPolicyName: z.string(),
            rolePrefix: z.string().max(20).optional(),
            policyPrefix: z.string().max(20).optional(),
            instanceProfilePrefix: z.string().optional(),
        })
        .optional()
        .describe('Aspect CDK injector for permissions. Ref: https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_iam.PermissionsBoundary.html'),
    stackSynthesizer: z.enum(stackSynthesizerType).optional().describe('Set the stack synthesize type. Ref: https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.StackSynthesizer.html'),
    bootstrapQualifier: z.string().optional().describe('CDK bootstrap qualifier to use for stack synthesis. Defaults to CDK default if not specified.'),
    bootstrapRolePrefix: z.string().optional().describe('Prefix for CDK bootstrap role names. Useful when roles have custom prefixes like My_User_Roles_. Leave empty for standard role names.'),
    litellmConfig: LiteLLMConfig,
    convertInlinePoliciesToManaged: z.boolean().optional().default(false).describe('Convert inline policies to managed policies'),
    iamRdsAuth: z.boolean().optional().default(false)
        .describe('Enable IAM authentication for RDS. When true (default), IAM authentication is used and the bootstrap password is deleted after setup. When false, password-based authentication is used. WARNING: Switching from true to false after deployment is not supported - the master password is permanently deleted when IAM auth is enabled. This is a one-way migration.'),
});

export const RawConfigSchema = RawConfigObject
    .refine((config) => (config.pypiConfig.indexUrl && config.region.includes('iso')) || !config.region.includes('iso'), {
        message: 'Must set PypiConfig if in an iso region',
    })
    .refine(
        (config) => {
            return !config.deployUi || config.deployChat;
        },
        {
            message: 'Chat stack is needed for UI stack. You must set deployChat to true if deployUi is true.',
        },
    )
    .refine(
        (config) => {
            return (
                !(config.deployChat || config.deployRag || config.deployUi) ||
                config.authConfig
            );
        },
        {
            message:
                'An auth config must be provided when deploying the chat, RAG, or UI stacks or when deploying an internet ' +
                'facing ALB. Check that `deployChat`, `deployRag`, `deployUi`, and `restApiConfig.internetFacing` are all ' +
                'false or that an `authConfig` is provided.',
        },
    )
    .describe('Raw application configuration schema.');