"use strict";
/*
 *  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *  SPDX-License-Identifier: Apache-2.0
 */
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const workbench_core_base_1 = require("@aws/workbench-core-base");
const workbench_core_environments_1 = require("@aws/workbench-core-environments");
const lodash_1 = __importDefault(require("lodash"));
const uuid_1 = require("uuid");
class SagemakerNotebookEnvironmentLifecycleService {
    constructor() {
        this._envType = 'sagemakerNotebook';
        this.helper = new workbench_core_environments_1.EnvironmentLifecycleHelper();
        this.aws = new workbench_core_base_1.AwsService({ region: process.env.AWS_REGION, ddbTableName: process.env.STACK_NAME });
        this.envService = new workbench_core_environments_1.EnvironmentService(this.aws.helpers.ddb);
    }
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    async launch(envMetadata) {
        var _a;
        const cidr = lodash_1.default.find(envMetadata.ETC.params, { key: 'CIDR' }).value;
        const instanceSize = lodash_1.default.find(envMetadata.ETC.params, { key: 'InstanceType' }).value;
        const autoStopIdleTimeInMinutes = lodash_1.default.find(envMetadata.ETC.params, { key: 'AutoStopIdleTimeInMinutes' })
            .value;
        const { datasetsBucketArn, mainAccountRegion, mainAccountId, mainAcctS3ArtifactEncryptionArn, mainAcctS3DatasetsEncryptionArn } = await this.helper.getCfnOutputs();
        const datasetIds = ((_a = envMetadata.DATASETS) === null || _a === void 0 ? void 0 : _a.map((dataSet) => {
            return dataSet.id;
        })) || [];
        const mainAcctEncryptionArnList = [mainAcctS3ArtifactEncryptionArn, mainAcctS3DatasetsEncryptionArn];
        const { s3Mounts, iamPolicyDocument } = await this.helper.getDatasetsToMount(datasetIds, envMetadata, mainAcctEncryptionArnList);
        const ssmParameters = {
            InstanceName: [`basicnotebookinstance-${Date.now()}`],
            VPC: [envMetadata.PROJ.vpcId],
            Subnet: [envMetadata.PROJ.subnetId],
            ProvisioningArtifactId: [envMetadata.ETC.provisioningArtifactId],
            ProductId: [envMetadata.ETC.productId],
            Namespace: [`${this._envType}-${Date.now()}`],
            EncryptionKeyArn: [envMetadata.PROJ.encryptionKeyArn],
            CIDR: [cidr],
            InstanceType: [instanceSize],
            DatasetsBucketArn: [datasetsBucketArn],
            EnvId: [envMetadata.id],
            EnvironmentInstanceFiles: [envMetadata.PROJ.environmentInstanceFiles],
            AutoStopIdleTimeInMinutes: [autoStopIdleTimeInMinutes],
            IamPolicyDocument: [iamPolicyDocument],
            S3Mounts: [s3Mounts],
            MainAccountS3ArtifactKeyArn: [mainAcctS3ArtifactEncryptionArn],
            MainAccountS3DatasetsKeyArn: [mainAcctS3DatasetsEncryptionArn],
            MainAccountRegion: [mainAccountRegion],
            MainAccountId: [mainAccountId]
        };
        await this.helper.launch({
            ssmParameters,
            operation: 'Launch',
            envType: this._envType,
            envMetadata
        });
        return Object.assign(Object.assign({}, envMetadata), { status: 'PENDING' });
    }
    async terminate(envId) {
        // Get value from env in DDB
        const envDetails = await this.envService.getEnvironment(envId, true);
        const provisionedProductId = envDetails.provisionedProductId; // This is updated by status handler
        const ssmParameters = {
            ProvisionedProductId: [provisionedProductId],
            TerminateToken: [(0, uuid_1.v4)()],
            EnvId: [envId]
        };
        // Execute termination doc
        await this.helper.executeSSMDocument({
            ssmParameters,
            operation: 'Terminate',
            envType: this._envType,
            envMgmtRoleArn: envDetails.PROJ.envMgmtRoleArn,
            externalId: envDetails.PROJ.externalId
        });
        // Delete access point(s) for this workspace
        await this.helper.removeAccessPoints(envDetails);
        // Store env row in DDB
        await this.envService.updateEnvironment(envId, { status: 'TERMINATING' });
        return { envId, status: 'TERMINATING' };
    }
    async start(envId) {
        // Get value from env in DDB
        const envDetails = await this.envService.getEnvironment(envId, true);
        const instanceName = envDetails.instanceId;
        // Assume hosting account EnvMgmt role
        const hostAwsSdk = await this.helper.getAwsSdkForEnvMgmtRole({
            envMgmtRoleArn: envDetails.PROJ.envMgmtRoleArn,
            externalId: envDetails.PROJ.externalId,
            operation: 'Start',
            envType: this._envType
        });
        await hostAwsSdk.clients.sagemaker.startNotebookInstance({ NotebookInstanceName: instanceName });
        // Store env row in DDB
        await this.envService.updateEnvironment(envId, { status: 'STARTING' });
        return { envId, status: 'STARTING' };
    }
    async stop(envId) {
        // Get value from env in DDB
        const envDetails = await this.envService.getEnvironment(envId, true);
        const instanceName = envDetails.instanceId;
        // Assume hosting account EnvMgmt role
        const hostAwsSdk = await this.helper.getAwsSdkForEnvMgmtRole({
            envMgmtRoleArn: envDetails.PROJ.envMgmtRoleArn,
            externalId: envDetails.PROJ.externalId,
            operation: 'Stop',
            envType: this._envType
        });
        await hostAwsSdk.clients.sagemaker.stopNotebookInstance({ NotebookInstanceName: instanceName });
        envDetails.status = 'STOPPING';
        // Store env row in DDB
        await this.envService.updateEnvironment(envId, { status: 'STOPPING' });
        return { envId, status: 'STOPPING' };
    }
}
exports.default = SagemakerNotebookEnvironmentLifecycleService;
//# sourceMappingURL=sagemakerNotebookEnvironmentLifecycleService.js.map