/*
 * Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license
 * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template
 */
package cloud.cleo.connectgpt;

import cloud.cleo.connectgpt.lang.LangUtil;
import static cloud.cleo.connectgpt.lang.LangUtil.LanguageIds.*;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.events.LexV2Event;
import com.amazonaws.services.lambda.runtime.events.LexV2Event.DialogAction;
import com.amazonaws.services.lambda.runtime.events.LexV2Event.Intent;
import com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState;
import com.amazonaws.services.lambda.runtime.events.LexV2Response;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.service.OpenAiService;
import java.net.SocketTimeoutException;
import java.time.Duration;
import java.time.LocalDate;
import java.time.ZoneId;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable;
import software.amazon.awssdk.enhanced.dynamodb.Key;
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
import software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedTimestampRecordExtension;

/**
 *
 * @author sjensen
 */
public class ChatGPTLambda implements RequestHandler<LexV2Event, LexV2Response> {

    // Initialize the Log4j logger.
    final static Logger log = LogManager.getLogger(ChatGPTLambda.class);

    final static ObjectMapper mapper = new ObjectMapper();

    final static TableSchema<ChatGPTSessionState> schema = TableSchema.fromBean(ChatGPTSessionState.class);

    final static DynamoDbEnhancedClient enhancedClient = DynamoDbEnhancedClient.builder()
            .extensions(AutoGeneratedTimestampRecordExtension.create()).build();

    final static DynamoDbTable<ChatGPTSessionState> sessionState = enhancedClient.table(System.getenv("SESSION_TABLE_NAME"), schema);

    final static OpenAiService open_ai_service = new OpenAiService(System.getenv("OPENAI_API_KEY"), Duration.ofSeconds(20));
    final static String OPENAI_MODEL = System.getenv("OPENAI_MODEL");

    @Override
    public LexV2Response handleRequest(LexV2Event lexRequest, Context cntxt) {
        try {
            log.debug(mapper.valueToTree(lexRequest).toString());
            final var intentName = lexRequest.getSessionState().getIntent().getName();
            log.debug("Intent: " + intentName);
            return processGPT(lexRequest);
        } catch (Exception e) {
            log.error(e);
            // Unhandled Exception
            return buildResponse(lexRequest, new LangUtil(lexRequest.getBot().getLocaleId()).getString(UNHANDLED_EXCEPTION));
        }
    }

    private LexV2Response processGPT(LexV2Event lexRequest) {

        final var input = lexRequest.getInputTranscript();
        final var localId = lexRequest.getBot().getLocaleId();

        final var lang = new LangUtil(localId);
        log.debug("Java Locale is " + lang.getLocale());

        if (input == null || input.isBlank()) {
            log.debug("Got blank input, so just silent or nothing");

            final var attrs = lexRequest.getSessionState().getSessionAttributes();
            var count = Integer.valueOf(attrs.getOrDefault("blankCounter", "0"));
            count++;

            if (count > 2) {
                log.debug("Two blank responses, sending to Quit Intent");
                // Hang up on caller after 2 silience requests
                return buildQuitResponse(lexRequest);
            } else {
                attrs.put("blankCounter", count.toString());
                // If we get slience (timeout without speech), then we get empty string on the transcript
                return buildResponse(lexRequest, lang.getString(BLANK_RESPONSE));
            }
        }

        // When testing in lex console input will be text, so use session ID, for speech we shoud have a phone via Connect
        final var user_id = lexRequest.getSessionId();

        // Key to record in Dynamo
        final var key = Key.builder().partitionValue(user_id).sortValue(LocalDate.now(ZoneId.of("America/Chicago")).toString()).build();

        //  load session state if it exists
        log.debug("Start Retreiving Session State");
        var session = sessionState.getItem(key);
        log.debug("End Retreiving Session State");

        if (session == null) {
            session = new ChatGPTSessionState(user_id);
        }

        // Since we can call and change language during session, always specifiy how we want responses
        session.addSystemMessage(lang.getString(CHATGPT_RESPONSE_LANGUAGE));
       
        // add this request to the session
        session.addUserMessage(input);

        String botResponse;
        try {

            ChatCompletionRequest request = ChatCompletionRequest.builder()
                    .messages(session.getChatMessages())
                    .model(OPENAI_MODEL)
                    .maxTokens(500)
                    .temperature(0.2) // More focused
                    .n(1) // Only return 1 completion
                    .build();

            log.debug("Start API Call to ChatGPT");
            final var completion = open_ai_service.createChatCompletion(request);
            log.debug("End API Call to ChatGPT");
            log.debug(completion);

            botResponse = completion.getChoices().get(0).getMessage().getContent();

            // Add response to session
            session.addAssistantMessage(botResponse);

            // Since we have a valid response, add message asking if there is anything else
            if ( ! "Text".equalsIgnoreCase(lexRequest.getInputMode()) ) {
                // Only add if not text (added to voice response)
                botResponse = botResponse + lang.getString(ANYTHING_ELSE);
            }

            // Save the session to dynamo
            log.debug("Start Saving Session State");
            session.incrementCounter();
            sessionState.putItem(session);
            log.debug("End Saving Session State");
        } catch (RuntimeException rte) {
            if (rte.getCause() != null && rte.getCause() instanceof SocketTimeoutException) {
                log.error("Response timed out", rte);
                    botResponse = lang.getString(OPERATION_TIMED_OUT);
            } else {
                throw rte;
            }
        }

        return buildResponse(lexRequest, botResponse);
    }

    /**
     * Response that sends you to the Quit intent so the call can be ended
     *
     * @param lexRequest
     * @param response
     * @return
     */
    private LexV2Response buildQuitResponse(LexV2Event lexRequest) {

        // State to return
        final var ss = SessionState.builder()
                // Retain the current session attributes
                .withSessionAttributes(lexRequest.getSessionState().getSessionAttributes())
                // Send back Quit Intent
                .withIntent(Intent.builder().withName("Quit").withState("ReadyForFulfillment").build())
                // Indicate the state is Delegate
                .withDialogAction(DialogAction.builder().withType("Delegate").build())
                .build();

        final var lexV2Res = LexV2Response.builder()
                .withSessionState(ss)
                .build();
        log.debug("Response is " + mapper.valueToTree(lexV2Res));
        return lexV2Res;
    }

    /**
     * General Response used to send back a message and Elicit Intent again at LEX
     *
     * @param lexRequest
     * @param response
     * @return
     */
    private LexV2Response buildResponse(LexV2Event lexRequest, String response) {

        // State to return
        final var ss = SessionState.builder()
                // Retain the current session attributes
                .withSessionAttributes(lexRequest.getSessionState().getSessionAttributes())
                // Always ElictIntent, so you're back at the LEX Bot looking for more input
                .withDialogAction(DialogAction.builder().withType("ElicitIntent").build())
                .build();

        final var lexV2Res = LexV2Response.builder()
                .withSessionState(ss)
                // We are using plain text responses
                .withMessages(new LexV2Response.Message[]{new LexV2Response.Message("PlainText", response, null)})
                .build();
        log.debug("Response is " + mapper.valueToTree(lexV2Res));
        return lexV2Res;
    }
}

com.theokanning.openai.completion.chat.ChatCompletionRequest.builder
[((1987, 2099), 'software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient.builder'), ((1987, 2091), 'software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient.builder'), ((4370, 4481), 'software.amazon.awssdk.enhanced.dynamodb.Key.builder'), ((4370, 4473), 'software.amazon.awssdk.enhanced.dynamodb.Key.builder'), ((4370, 4407), 'software.amazon.awssdk.enhanced.dynamodb.Key.builder'), ((4418, 4472), 'java.time.LocalDate.now'), ((5138, 5440), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((5138, 5383), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((5138, 5341), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((5138, 5303), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((5138, 5267), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((5138, 5226), 'com.theokanning.openai.completion.chat.ChatCompletionRequest.builder'), ((7069, 7547), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((7069, 7522), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((7069, 7385), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((7069, 7240), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((7310, 7384), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.Intent.builder'), ((7310, 7376), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.Intent.builder'), ((7310, 7343), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.Intent.builder'), ((7470, 7521), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.DialogAction.builder'), ((7470, 7513), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.DialogAction.builder'), ((7579, 7665), 'com.amazonaws.services.lambda.runtime.events.LexV2Response.builder'), ((7579, 7640), 'com.amazonaws.services.lambda.runtime.events.LexV2Response.builder'), ((8067, 8446), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((8067, 8421), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((8067, 8238), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.SessionState.builder'), ((8365, 8420), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.DialogAction.builder'), ((8365, 8412), 'com.amazonaws.services.lambda.runtime.events.LexV2Event.DialogAction.builder'), ((8478, 8732), 'com.amazonaws.services.lambda.runtime.events.LexV2Response.builder'), ((8478, 8707), 'com.amazonaws.services.lambda.runtime.events.LexV2Response.builder'), ((8478, 8539), 'com.amazonaws.services.lambda.runtime.events.LexV2Response.builder')]