import {
	DynamoDBClient,
	UpdateTableCommand,
	DescribeTableCommand,
} from '@aws-sdk/client-dynamodb';
import { GlobalSecondaryIndexProps } from 'aws-cdk-lib/aws-dynamodb';
import { CloudFormationCustomResourceEvent } from 'aws-lambda';

const ddb = new DynamoDBClient({});

function sleep(ms: number) {
	return new Promise((res) => setTimeout(res, ms));
}

/**
 * Polls DDB until status matches the desired status.
 * 
 * @param tableName 
 * @param indexName 
 * @param desiredStatus 
 */
async function waitForGsiStatus(
	tableName: string,
	indexName: string,
	desiredStatus: string
) {
	while (true) {
		const table = await ddb.send(new DescribeTableCommand({ TableName: tableName }));
		const gsis = table.Table?.GlobalSecondaryIndexes ?? [];
		const index = gsis.find((gsi) => gsi.IndexName === indexName);
		if (!index && desiredStatus === 'DELETED') break;
		if (index?.IndexStatus === desiredStatus) break;
		await sleep(5000);
	}
}

export const handler = async (event: CloudFormationCustomResourceEvent) => {
	console.log('Event:', JSON.stringify(event));

	const tableName = event.ResourceProperties.TableName;
	const desiredGsis: GlobalSecondaryIndexProps[] =
		event.ResourceProperties.GlobalSecondaryIndexes || [];
	const desiredNames = new Set(desiredGsis.map((gsi) => gsi.indexName));

	const tableDesc = await ddb.send(new DescribeTableCommand({ TableName: tableName }));
	const existingGsis = tableDesc.Table?.GlobalSecondaryIndexes ?? [];
	const existingNames = new Set(existingGsis.map((gsi) => gsi.IndexName));

	if (event.RequestType === 'Delete') {
		// table-level delete will take care of deleting everything.
		return { PhysicalResourceId: tableName };
	}

	// Delete unwanted GSIs first.
	for (const gsi of existingGsis) {
		if (!gsi.IndexName) {
			console.warn('"Existing" GSI without IndexName found (1)', gsi);
			continue;
		}
		if (!desiredNames.has(gsi.IndexName)) {
			console.log(`Deleting GSI: ${gsi.IndexName}`);
			await ddb.send(new UpdateTableCommand({
				TableName: tableName,
				GlobalSecondaryIndexUpdates: [
					{ Delete: { IndexName: gsi.IndexName } },
				],
			}));
			await waitForGsiStatus(tableName, gsi.IndexName, 'DELETED');
		}
	}

	// Add new GSIs
	for (const gsi of desiredGsis) {
		if (!existingNames.has(gsi.indexName)) {
			console.log(`Creating GSI: ${gsi.indexName}`);
			await ddb.send(new UpdateTableCommand({
				TableName: tableName,
				AttributeDefinitions: [
					{
						AttributeName: gsi.partitionKey!.name,
						AttributeType: gsi.partitionKey!.type,
					},
					...(gsi.sortKey ? [{
						AttributeName: gsi.sortKey.name,
						AttributeType: gsi.sortKey.type,
					}] : []),
				],
				GlobalSecondaryIndexUpdates: [{
					Create: {
						IndexName: gsi.indexName,
						KeySchema: [
							{
								AttributeName: gsi.partitionKey!.name,
								KeyType: 'HASH'
							},
							...(gsi.sortKey ? [{
								AttributeName: gsi.sortKey.name,
								KeyType: 'RANGE' as const
							}] : []),
						],
						Projection: {
							ProjectionType: gsi.projectionType ?? 'ALL',
						}
					}
				}],
			}));
			await waitForGsiStatus(tableName, gsi.indexName, 'ACTIVE');
		}
	}

	return {
		PhysicalResourceId: tableName,
		Data: { Message: 'GSIs synced successfully' },
	};
};