"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.ConversationResolverGenerator = void 0;
const tools_1 = require("../utils/tools");
const graphql_transformer_core_1 = require("@aws-amplify/graphql-transformer-core");
const graphql_transformer_common_1 = require("graphql-transformer-common");
const cdk = __importStar(require("aws-cdk-lib"));
const ai_constructs_1 = require("@aws-amplify/ai-constructs");
const graphql_transformer_core_2 = require("@aws-amplify/graphql-transformer-core");
const init_resolver_1 = require("../resolvers/init-resolver");
const auth_resolver_1 = require("../resolvers/auth-resolver");
const verify_session_owner_resolver_1 = require("../resolvers/verify-session-owner-resolver");
const write_message_to_table_resolver_1 = require("../resolvers/write-message-to-table-resolver");
const invoke_lambda_resolver_1 = require("../resolvers/invoke-lambda-resolver");
const assistant_mutation_resolver_1 = require("../resolvers/assistant-mutation-resolver");
const assistant_messages_subscription_resolver_1 = require("../resolvers/assistant-messages-subscription-resolver");
const graphql_index_transformer_1 = require("@aws-amplify/graphql-index-transformer");
const pluralize_1 = __importDefault(require("pluralize"));
const list_messages_init_resolver_1 = require("../resolvers/list-messages-init-resolver");
const assistant_streaming_mutation_resolver_1 = require("../resolvers/assistant-streaming-mutation-resolver");
class ConversationResolverGenerator {
    constructor(functionNameMap) {
        this.functionNameMap = functionNameMap;
    }
    generateResolvers(directives, ctx) {
        for (const directive of directives) {
            this.processToolsForDirective(directive, ctx);
            this.generateResolversForDirective(directive, ctx);
            this.addInitSlotToListMessagesPipeline(ctx, directive);
        }
    }
    processToolsForDirective(directive, ctx) {
        const tools = (0, tools_1.processTools)(directive.tools, ctx);
        if (tools) {
            directive.toolSpec = tools;
        }
    }
    generateResolversForDirective(directive, ctx) {
        const { parent, field } = directive;
        const parentName = parent.name.value;
        const capitalizedFieldName = (0, graphql_transformer_common_1.toUpper)(field.name.value);
        const fieldName = field.name.value;
        const functionStack = this.createFunctionStack(ctx, capitalizedFieldName);
        const { functionDataSourceId, referencedFunction } = this.setupFunctionDataSource(directive, functionStack, capitalizedFieldName);
        const functionDataSource = this.addLambdaDataSource(ctx, functionDataSourceId, referencedFunction, capitalizedFieldName);
        const invokeLambdaFunction = (0, invoke_lambda_resolver_1.invokeLambdaMappingTemplate)(directive);
        this.setupMessageTableIndex(ctx, directive);
        const initResolverFunction = (0, init_resolver_1.initMappingTemplate)(ctx);
        const authResolverFunction = (0, auth_resolver_1.authMappingTemplate)(directive);
        const verifySessionOwnerSendMessageResolverFunction = (0, verify_session_owner_resolver_1.verifySessionOwnerSendMessageMappingTemplate)(directive);
        const verifySessionOwnerAssistantResponseResolverFunction = (0, verify_session_owner_resolver_1.verifySessionOwnerAssistantResponseMappingTemplate)(directive);
        const writeMessageToTableFunction = (0, write_message_to_table_resolver_1.writeMessageToTableMappingTemplate)(directive);
        this.createConversationPipelineResolver(ctx, parentName, fieldName, capitalizedFieldName, functionDataSource, invokeLambdaFunction, initResolverFunction, authResolverFunction, verifySessionOwnerSendMessageResolverFunction, writeMessageToTableFunction);
        this.createAssistantResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerAssistantResponseResolverFunction);
        this.createAssistantStreamingResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerAssistantResponseResolverFunction);
        this.createAssistantResponseSubscriptionResolver(ctx, directive, capitalizedFieldName);
    }
    createFunctionStack(ctx, capitalizedFieldName) {
        return ctx.stackManager.createStack(`${capitalizedFieldName}ConversationDirectiveLambdaStack`);
    }
    setupFunctionDataSource(directive, functionStack, capitalizedFieldName) {
        if (directive.functionName) {
            return this.setupExistingFunctionDataSource(directive.functionName);
        }
        else {
            return this.setupDefaultConversationHandler(functionStack, capitalizedFieldName, directive.aiModel);
        }
    }
    setupExistingFunctionDataSource(functionName) {
        const functionDataSourceId = graphql_transformer_common_1.FunctionResourceIDs.FunctionDataSourceID(functionName);
        if (!this.functionNameMap) {
            throw new Error('Function name map is not provided');
        }
        const referencedFunction = this.functionNameMap[functionName];
        if (!referencedFunction) {
            throw new Error(`Function ${functionName} not found in function name map`);
        }
        return { functionDataSourceId, referencedFunction };
    }
    setupDefaultConversationHandler(functionStack, capitalizedFieldName, aiModel) {
        const defaultConversationHandler = new ai_constructs_1.conversation.ConversationHandlerFunction(functionStack, `${capitalizedFieldName}DefaultConversationHandler`, {
            models: [
                {
                    modelId: aiModel,
                },
            ],
        });
        const functionDataSourceId = graphql_transformer_common_1.FunctionResourceIDs.FunctionDataSourceID(`${capitalizedFieldName}DefaultConversationHandler`);
        const referencedFunction = defaultConversationHandler.resources.lambda;
        return { functionDataSourceId, referencedFunction };
    }
    createConversationPipelineResolver(ctx, parentName, fieldName, capitalizedFieldName, functionDataSource, invokeLambdaFunction, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction, writeMessageToTableFunction) {
        const resolverResourceId = graphql_transformer_common_1.ResolverResourceIDs.ResolverResourceID(parentName, fieldName);
        const runtime = graphql_transformer_core_1.APPSYNC_JS_RUNTIME;
        const conversationPipelineResolver = new graphql_transformer_core_1.TransformerResolver(parentName, fieldName, resolverResourceId, { codeMappingTemplate: invokeLambdaFunction }, ['init', 'auth', 'verifySessionOwner', 'writeMessageToTable'], ['handleLambdaResponse', 'finish'], functionDataSource, runtime);
        this.addPipelineResolverFunctions(ctx, conversationPipelineResolver, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction, writeMessageToTableFunction);
        ctx.resolvers.addResolver(parentName, fieldName, conversationPipelineResolver);
    }
    addPipelineResolverFunctions(ctx, resolver, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction, writeMessageToTableFunction) {
        resolver.addJsFunctionToSlot('init', initResolverFunction);
        resolver.addJsFunctionToSlot('auth', authResolverFunction);
        const sessionModelName = `Conversation${capitalizedFieldName}`;
        const sessionModelDDBDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, sessionModelName);
        const conversationSessionDDBDataSource = ctx.api.host.getDataSource(sessionModelDDBDataSourceName);
        resolver.addJsFunctionToSlot('verifySessionOwner', verifySessionOwnerResolverFunction, conversationSessionDDBDataSource);
        const messageModelName = `ConversationMessage${capitalizedFieldName}`;
        const messageModelDDBDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, messageModelName);
        const messageDDBDataSource = ctx.api.host.getDataSource(messageModelDDBDataSourceName);
        resolver.addJsFunctionToSlot('writeMessageToTable', writeMessageToTableFunction, messageDDBDataSource);
    }
    createAssistantResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction) {
        const assistantResponseResolverResourceId = graphql_transformer_common_1.ResolverResourceIDs.ResolverResourceID('Mutation', directive.responseMutationName);
        const assistantResponseResolverFunction = (0, assistant_mutation_resolver_1.assistantMutationResolver)(directive);
        const conversationMessageDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, `ConversationMessage${capitalizedFieldName}`);
        const conversationMessageDataSource = ctx.api.host.getDataSource(conversationMessageDataSourceName);
        const resolver = new graphql_transformer_core_1.TransformerResolver('Mutation', directive.responseMutationName, assistantResponseResolverResourceId, { codeMappingTemplate: assistantResponseResolverFunction }, ['init', 'auth', 'verifySessionOwner'], [], conversationMessageDataSource, graphql_transformer_core_1.APPSYNC_JS_RUNTIME);
        resolver.addJsFunctionToSlot('init', initResolverFunction);
        resolver.addJsFunctionToSlot('auth', authResolverFunction);
        const sessionModelName = `Conversation${capitalizedFieldName}`;
        const sessionModelDDBDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, sessionModelName);
        const conversationSessionDDBDataSource = ctx.api.host.getDataSource(sessionModelDDBDataSourceName);
        resolver.addJsFunctionToSlot('verifySessionOwner', verifySessionOwnerResolverFunction, conversationSessionDDBDataSource);
        ctx.resolvers.addResolver('Mutation', directive.responseMutationName, resolver);
    }
    createAssistantStreamingResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction) {
        const assistantResponseResolverResourceId = graphql_transformer_common_1.ResolverResourceIDs.ResolverResourceID('Mutation', directive.messageModel.assistantStreamingMutationField.name.value);
        const assistantResponseResolverFunction = (0, assistant_streaming_mutation_resolver_1.assistantStreamingMutationResolver)(directive);
        const conversationMessageDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, `ConversationMessage${capitalizedFieldName}`);
        const conversationMessageDataSource = ctx.api.host.getDataSource(conversationMessageDataSourceName);
        const resolver = new graphql_transformer_core_1.TransformerResolver('Mutation', directive.messageModel.assistantStreamingMutationField.name.value, assistantResponseResolverResourceId, { codeMappingTemplate: assistantResponseResolverFunction }, ['init', 'auth', 'verifySessionOwner'], [], conversationMessageDataSource, graphql_transformer_core_1.APPSYNC_JS_RUNTIME);
        resolver.addJsFunctionToSlot('init', initResolverFunction);
        resolver.addJsFunctionToSlot('auth', authResolverFunction);
        const sessionModelName = `Conversation${capitalizedFieldName}`;
        const sessionModelDDBDataSourceName = (0, graphql_transformer_core_2.getModelDataSourceNameForTypeName)(ctx, sessionModelName);
        const conversationSessionDDBDataSource = ctx.api.host.getDataSource(sessionModelDDBDataSourceName);
        resolver.addJsFunctionToSlot('verifySessionOwner', verifySessionOwnerResolverFunction, conversationSessionDDBDataSource);
        ctx.resolvers.addResolver('Mutation', directive.messageModel.assistantStreamingMutationField.name.value, resolver);
    }
    createAssistantResponseSubscriptionResolver(ctx, directive, capitalizedFieldName) {
        const onAssistantResponseSubscriptionFieldName = `onCreateAssistantResponse${capitalizedFieldName}`;
        const onAssistantResponseSubscriptionResolverResourceId = graphql_transformer_common_1.ResolverResourceIDs.ResolverResourceID('Subscription', onAssistantResponseSubscriptionFieldName);
        const onAssistantResponseSubscriptionResolverFunction = (0, assistant_messages_subscription_resolver_1.conversationMessageSubscriptionMappingTamplate)(directive);
        const mappingTemplate = {
            codeMappingTemplate: onAssistantResponseSubscriptionResolverFunction,
        };
        const onAssistantResponseSubscriptionResolver = new graphql_transformer_core_1.TransformerResolver('Subscription', onAssistantResponseSubscriptionFieldName, onAssistantResponseSubscriptionResolverResourceId, mappingTemplate, [], [], undefined, graphql_transformer_core_1.APPSYNC_JS_RUNTIME);
        ctx.resolvers.addResolver('Subscription', onAssistantResponseSubscriptionFieldName, onAssistantResponseSubscriptionResolver);
    }
    addLambdaDataSource(ctx, functionDataSourceId, referencedFunction, capitalizedFieldName) {
        const functionDataSourceScope = ctx.stackManager.getScopeFor(functionDataSourceId, `${capitalizedFieldName}ConversationDirectiveLambdaStack`);
        return ctx.api.host.addLambdaDataSource(functionDataSourceId, referencedFunction, {}, functionDataSourceScope);
    }
    addInitSlotToListMessagesPipeline(ctx, directive) {
        const messageModelName = directive.messageModel.messageModel.name.value;
        const pluralized = (0, pluralize_1.default)(messageModelName);
        const listMessagesResolver = ctx.resolvers.getResolver('Query', `list${pluralized}`);
        const initResolverFn = (0, list_messages_init_resolver_1.listMessageInitMappingTemplate)(directive);
        listMessagesResolver.addJsFunctionToSlot('init', initResolverFn);
    }
    setupMessageTableIndex(ctx, directive) {
        const messageModelName = directive.messageModel.messageModel.name.value;
        const referenceFieldName = 'conversationId';
        const messageModel = directive.messageModel.messageModel;
        const conversationMessagesTable = (0, graphql_transformer_core_2.getTable)(ctx, messageModel);
        const gsiPartitionKeyName = referenceFieldName;
        const gsiPartitionKeyType = 'S';
        const gsiSortKeyName = 'createdAt';
        const gsiSortKeyType = 'S';
        const indexName = 'gsi-ConversationMessage.conversationId.createdAt';
        this.addGlobalSecondaryIndex(conversationMessagesTable, indexName, { name: gsiPartitionKeyName, type: gsiPartitionKeyType }, { name: gsiSortKeyName, type: gsiSortKeyType }, ctx, messageModelName);
    }
    addGlobalSecondaryIndex(table, indexName, partitionKey, sortKey, ctx, typeName) {
        table.addGlobalSecondaryIndex({
            indexName,
            projectionType: 'ALL',
            partitionKey,
            sortKey,
            readCapacity: cdk.Fn.ref(graphql_transformer_common_1.ResourceConstants.PARAMETERS.DynamoDBModelTableReadIOPS),
            writeCapacity: cdk.Fn.ref(graphql_transformer_common_1.ResourceConstants.PARAMETERS.DynamoDBModelTableWriteIOPS),
        });
        const gsi = table.globalSecondaryIndexes.find((g) => g.indexName === indexName);
        const newIndex = {
            indexName,
            keySchema: gsi.keySchema,
            projection: { projectionType: 'ALL' },
            provisionedThroughput: cdk.Fn.conditionIf(graphql_transformer_common_1.ResourceConstants.CONDITIONS.ShouldUsePayPerRequestBilling, cdk.Fn.ref('AWS::NoValue'), {
                ReadCapacityUnits: cdk.Fn.ref(graphql_transformer_common_1.ResourceConstants.PARAMETERS.DynamoDBModelTableReadIOPS),
                WriteCapacityUnits: cdk.Fn.ref(graphql_transformer_common_1.ResourceConstants.PARAMETERS.DynamoDBModelTableWriteIOPS),
            }),
        };
        (0, graphql_index_transformer_1.overrideIndexAtCfnLevel)(ctx, typeName, table, newIndex);
    }
}
exports.ConversationResolverGenerator = ConversationResolverGenerator;
//# sourceMappingURL=conversation-resolver-generator.js.map