/** ************************************************************************************************
*   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.                             *
*   SPDX-License-Identifier: Apache-2.0                                                            *
 ************************************************************************************************ */

const _ = require('lodash');
const qnabot = require('qnabot/logging');
const qna_settings = require('qnabot/settings');
const handlebars = require('../handlebars');
const kendra_fallback = require('../kendra');
const kendra_retrieve = require('../kendraRetrieve');
const { determineKendraLanguage } = require('../kendraClient');
const open_es = require('../es_query');
const { invokeLambda } = require('./invokeLambda');
const { runKendraQuery } = require('./runKendraQuery');
const { encryptor } = require('./encryptor');
const { runLlmQa } = require('./runLlmQa');
const { updateResWithHit } = require('./updateResWithHit');
const { bedrockRetrieveAndGenerate } = require('../bedrock/bedrockAgents');
const { processKeysForRedact } = require('../redactHelper');

async function runQuery(req, query_params, kendraIndex) {
    query_params.kendraIndex = kendraIndex;
    const onlyES = await open_es.isESonly(req, query_params);
    let response = '';
    // runs kendra query if question supported on Kendra and KENDRA_FAQ_INDEX is set
    if (!onlyES) {
        response = await runKendraQuery(req, query_params);
    } else {
        response = await open_es.run_query_es(req, query_params);
    }
    return response;
}
function getSourceLinksFromPassages(inputText) {
    const sourceLinkPattern = /^\s*Source Link:(.*)$/gm; // NOSONAR - javascript:S5852 - input is user controlled and we have a limit on the number of characters
    let matches;
    const sourceLinks = [];

    while ((matches = sourceLinkPattern.exec(inputText)) !== null) {
        sourceLinks.push(matches[1].trim().replace(/(^")|("$)/g, ''));
    }

    const uniqueLinks = [...new Set(sourceLinks)];
    return uniqueLinks.length > 0 ? `Sources: ${uniqueLinks.join(', ')}` : '';
}

async function kendraFallback(req, res) {
    let hit;
    let errors = [];
    if (req._settings.LLM_QA_ENABLE && req._settings.LLM_QA_USE_KENDRA_RETRIEVAL_API) {
        qnabot.log(`Kendra Fallback using Retrieve API: ${JSON.stringify(req)}`);
        hit = await kendra_retrieve.handler(req, res);
        qnabot.log('Kendra Fallback result: ', JSON.stringify(hit, null, 2));
    }
    // if we still don't have a hit, either retrieval was skipped or failed. Try the Query API
    if (!hit) {
        qnabot.log(`Kendra Fallback using Query API: ${JSON.stringify(req)}`);
        hit = await kendra_fallback.handler({ req, res });
        qnabot.log(`Result from Kendra Fallback ${JSON.stringify(hit)}`);
    }

    if (hit && hit.hit_count != 0) {
        hit.refMarkdown = getSourceLinksFromPassages(hit.alt.markdown);
        // Run any configured QA Summary LLM model options on Kendra results
        [hit, errors] = await runLlmQa(req, hit);
        if (hit) {
            _.set(res, 'answersource', 'Kendra Fallback');
            _.set(res, 'session.qnabot_gotanswer', true);
            _.set(res, 'message', hit.a);
            _.set(req, 'debug', hit.debug);
            res.got_hits = 1;
        }
    }
    return [req, res, hit, errors];
}

function encryptConditionalChainingIfSet(hit) {
    const conditionalChaining = _.get(hit, 'conditionalChaining');
    if (conditionalChaining) {
        qnabot.log('Encrypt conditionalChaining rule to ensure it is tamper proof in session attributes');
        const encrypted = encryptor.encrypt(conditionalChaining);
        _.set(hit, 'conditionalChaining', encrypted);
    }
    return hit;
}

async function getNoHitsResponse(query_params, res, req, errors, KENDRA_FAQ_INDEX) {
    const noHitsQuestion = _.get(req, '_settings.ES_NO_HITS_QUESTION', 'no_hits');
    const errorQuestion = _.get(req, '_settings.ES_ERROR_QUESTION', 'error_msg');

    query_params.question = errors.length > 0 ? errorQuestion : noHitsQuestion;
    query_params.score_text_passage = false;
    query_params.size = 1;
    res.got_hits = 0; // response flag, used in logging / kibana

    qnabot.log(`No hits from query - searching instead for: ${query_params.question}`);

    const response = await runQuery(req, query_params, KENDRA_FAQ_INDEX);
    const noHitsRes = _.get(response, 'hits.hits[0]._source');
    qnabot.log(`No hits response: ${JSON.stringify(noHitsRes)}`);
    return noHitsRes;
}

function setSessionTopic(hit, res) {
    qnabot.log(`Setting topic for ${JSON.stringify(hit)}`);
    // set res topic from document before running handlebars, so that handlebars can access or overwrite it.
    _.set(res, 'session.topic', _.get(hit, 't'));

    if (_.get(hit, 't')) {
        if (!res._userInfo) {
            res._userInfo = {};
        }
        if (!res._userInfo.recentTopics) {
            res._userInfo.recentTopics = [];
        }
        res._userInfo.recentTopics.push({
            topic: _.get(hit, 't'),
            dateTime: new Date().toISOString(),
        });
    }
    return res;
}

async function invokeLambdaHook(hit, req, res) {
    // Call Lambda Hook with args now & override running as middleware step (old behavior)
    // This results in:
    //  - improved predictability of document chaining behavior.. each doc's lambda is run as it is chained
    //  - autotranslation is now applied to lambda hook responses by default when response is assembled
    // optional setting to turn off this behaviour if it causes problems, and revert to old way
    const lambdaHook = _.get(hit, 'l');
    if (lambdaHook) {
        qnabot.log('Invoking Lambda Hook function: ', lambdaHook);
        const redactEnabled = _.get(req, '_settings.ENABLE_REDACTING');
        const redactComprehendEnabled  =_.get(req, '_settings.ENABLE_REDACTING_WITH_COMPREHEND', false);
        if (lambdaHook.toLowerCase().includes('feedback') && (redactEnabled || redactComprehendEnabled) ) {
            qna_settings.set_environment_variables(req._settings);
            await qnabot.setPIIRedactionEnvironmentVars(
                req._event.inputTranscript,
                _.get(req, '_settings.ENABLE_REDACTING_WITH_COMPREHEND', false),
                _.get(req, '_settings.REDACTING_REGEX', ''),
                _.get(req, '_settings.COMPREHEND_REDACTING_ENTITY_TYPES', ''),
                _.get(req, '_settings.COMPREHEND_REDACTING_CONFIDENCE_SCORE', 0.99),
            );
            processKeysForRedact(req, true)
            processKeysForRedact(res, true)
        }
        [req, res] = await invokeLambda(lambdaHook, req, res);
        // update hit with values returned in res by lambda hook
        _.set(hit, 'a', _.get(res, 'message', ''));
        const markdown = _.get(res, 'session.appContext.altMessages.markdown', '');
        const ssml = _.get(res, 'session.appContext.altMessages.ssml', '');
        const card = _.get(res, 'card', {});
        _.set(hit, 'alt.markdown', markdown);
        _.set(hit, 'alt.ssml', ssml);
        _.set(hit, 'r', card);
    }
    _.set(hit, 'l', '');
    _.set(hit, 'args', []);
    return [req, res, hit];
}

function copyPassageFieldToAnswerFields(hit) {
    if (hit.passage && !hit.a) {
        // Set the answer (a) field to match the text item passage field.
        hit.a = hit.passage;
    }
    if (!_.get(hit, 'alt.markdown')) {
        _.set(hit, 'alt.markdown', hit.a);
    }
    if (!_.get(hit, 'alt.ssml')) {
        _.set(hit, 'alt.ssml', hit.a);
    }

    return hit;
}

async function runFirstQuery(req, query_params, KENDRA_FAQ_INDEX, res) {
    const response = await runQuery(req, query_params, KENDRA_FAQ_INDEX);
    qnabot.debug('Query response: ', JSON.stringify(response, null, 2));
    const hit = _.get(response, 'hits.hits[0]._source');

    _.set(res, 'kendraResultsCached', response.kendraResultsCached);
    if (response.kendraResultsCached) qnabot.debug('kendra results cached in res structure');
    _.set(req, 'session.qnabotcontext.kendra', response.kendra_context);
    if (response.kendra_context) qnabot.log('kendra context set in res session');
    return [req, res, hit];
}

async function opensearchFallback(req, query_params) {
    qnabot.log('OpenSearch Fallback');
    const response = await open_es.run_query_es(req, query_params);
    if (_.get(response, 'hits.hits[0]._source')) {
        _.set(response, 'hits.hits[0]._source.answersource', 'OpenSearch Fallback');
    }
    const osFallbackHit = _.get(response, 'hits.hits[0]._source');
    return osFallbackHit;
}

async function kendraRedirect(hit, req, res, ALT_SEARCH_KENDRA_FALLBACK_CONFIDENCE_SCORE, query_params) {
    let redirect = {
        kendraRedirectQueryText: _.get(hit, 'kendraRedirectQueryText'),
        kendraRedirectQueryArgs: _.get(hit, 'kendraRedirectQueryArgs', []),
    };
    // process any handlebars before running Kendra redirect query
    qnabot.log('Kendra redirect query: Process with handlebars before redirecting.');
    redirect = await handlebars(req, res, redirect);
    const kendraRedirectQueryText = _.get(redirect, 'kendraRedirectQueryText');
    const kendraRedirectQueryArgs = _.get(redirect, 'kendraRedirectQueryArgs', []);
    const kendraRedirectQueryConfidenceThreshold = _.get(
        hit,
        'kendraRedirectQueryConfidenceThreshold',
        ALT_SEARCH_KENDRA_FALLBACK_CONFIDENCE_SCORE,
    );
    qnabot.log(`Kendra redirect query: '${kendraRedirectQueryText}' - Args = '${kendraRedirectQueryArgs}'`);
    qnabot.log(`Kendra redirect query confidence threshold: '${kendraRedirectQueryConfidenceThreshold}'`);
    req.question = kendraRedirectQueryText;
    req.kendraQueryArgs = kendraRedirectQueryArgs;
    req.kendraRedirect = true;
    req._settings.ALT_SEARCH_KENDRA_FALLBACK_CONFIDENCE_SCORE = kendraRedirectQueryConfidenceThreshold;
    // remove any cached results from FAQ query
    delete res.kendraResultsCached;
    const redirect_hit = await kendra_fallback.handler({ req, res });
    if (redirect_hit) {
        qnabot.log(`Result from Kendra Redirect query: ${JSON.stringify(redirect_hit)}`);
        hit.answersource = 'KENDRA REDIRECT';
        hit.a = _.get(redirect_hit, 'a');
        hit.alt = _.get(redirect_hit, 'alt');
    } else {
        qnabot.log('Kendra Redirect query returned no hits. Disable Kendra fallback query.');
        query_params.kendra_indexes = '';
        hit = null;
    }

    return [req, res, hit];
}

function getQuestion(req) {
    let { question } = req;
    const qid = _.get(req, 'qid');
    if (qid) {
        question = `QID::${qid}`;
        qnabot.log(`*** QID identified in request: ${qid}`);
    }
    return question;
}

const useFallbackMethod = async (req, res, query_params) => {
    const { KNOWLEDGE_BASE_ID, KNOWLEDGE_BASE_MODEL_ID, FALLBACK_ORDER } = req._settings;

    let errors = [];
    let hit;

    const useFallbackOptions = determineFallbackOptions(KNOWLEDGE_BASE_ID, KNOWLEDGE_BASE_MODEL_ID, query_params);
    const fallbackOrder = FALLBACK_ORDER || 'KNOWLEDGEBASE-FIRST'; // 'KENDRA-FIRST' or 'KNOWLEDGEBASE-FIRST'

    const invokeKendra = async () => {
        [req, res, hit, errors] = await kendraFallback(req, res);
        return hit != null;
    };

    const invokeKnowledgeBase = async () => {
        try {
            [res, hit] = await bedrockRetrieveAndGenerate(req, res);
            return hit != null;
        } catch (e) {
            const errMsg = `Bedrock Knowledgebase exception: ${e.name} ${e.message.substring(0, 500)}`;
            qnabot.log(`BEDROCK KNOWLEDGEBASE EXCEPTION: ${errMsg}`);

            if (!errors.includes(errMsg)) {
                errors.push(errMsg);
            }
        }
    };

    await fallbackOrchaestration(useFallbackOptions, invokeKendra, invokeKnowledgeBase, fallbackOrder);

    return [req, res, hit, errors];
};

const fallbackOrchaestration = async (useFallbackOptions, invokeKendra, invokeKnowledgeBase, fallbackOrder) => {

    if (useFallbackOptions === 'KENDRA') {
        qnabot.log('Using Kendra as a fallback option');
        await invokeKendra();

    } else if (useFallbackOptions === 'KNOWLEDGEBASE') {

        qnabot.log('Using KnowledgeBase as a fallback option');
        await invokeKnowledgeBase();

    } else if (useFallbackOptions === 'MULTIPLE') {
        if (fallbackOrder === 'KENDRA-FIRST') {

            qnabot.log('Using Kendra as a first fallback option');
            if (!(await invokeKendra())) {
                qnabot.log('Using KnowledgeBase as a second fallback option');
                await invokeKnowledgeBase();
            }

        } else {
            qnabot.log('Using KnowledgeBase as a first fallback option');
            if (!(await invokeKnowledgeBase())) {

                qnabot.log('Using Kendra as a second fallback option');
                await invokeKendra();
            }
        }
    } else {
        qnabot.log('Skipping fallback as there are no fallback data sources configured');
    }
};

const determineFallbackOptions = (KNOWLEDGE_BASE_ID, KNOWLEDGE_BASE_MODEL_ID, query_params) => {
    const fallbackOptions = [];

    // Check if Kendra is configured
    if (query_params.kendra_indexes.length !== 0) {
        fallbackOptions.push('KENDRA');
    }

    // Check if Knowledge Base is configured
    if (KNOWLEDGE_BASE_ID && KNOWLEDGE_BASE_MODEL_ID) {
        fallbackOptions.push('KNOWLEDGEBASE');
    }

    // Check if Multiple Options are configured
    if (fallbackOptions.length > 1) {
        fallbackOptions.push('MULTIPLE');
    }

    // Check if no fallback options are configured, default to DISABLED
    if (fallbackOptions.length === 0) {
        fallbackOptions.push('DISABLED');
    }

    return fallbackOptions[fallbackOptions.length - 1];
};
async function getHit(req, res) {
    const question = getQuestion(req);
    let size = 1;

    if (open_es.isQuestionAllStopwords(question)) {
        qnabot.log(`Question '${question}' contains only stop words. Forcing no hits.`);
        size = 0;
    }
    let errors = [];

    const {
        ES_MINIMUM_SHOULD_MATCH,
        ES_PHRASE_BOOST,
        ES_USE_KEYWORD_FILTERS,
        ES_KEYWORD_SYNTAX_TYPES,
        ES_SYNTAX_CONFIDENCE_LIMIT,
        ES_USE_FUZZY_MATCH,
        ES_EXPAND_CONTRACTIONS,
        ALT_SEARCH_KENDRA_INDEXES,
        ALT_SEARCH_KENDRA_FAQ_CONFIDENCE_SCORE,
        KENDRA_FAQ_INDEX,
        KENDRA_FAQ_ES_FALLBACK,
        ALT_SEARCH_KENDRA_FALLBACK_CONFIDENCE_SCORE,
        RUN_LAMBDAHOOK_FROM_QUERY_STEP,
    } = req._settings;

    const language = determineKendraLanguage(req);

    const query_params = {
        question,
        topic: _.get(req, 'session.topic', ''),
        from: 0,
        size,
        minimum_should_match: ES_MINIMUM_SHOULD_MATCH,
        phrase_boost: ES_PHRASE_BOOST,
        use_keyword_filters: ES_USE_KEYWORD_FILTERS,
        keyword_syntax_types: ES_KEYWORD_SYNTAX_TYPES,
        syntax_confidence_limit: ES_SYNTAX_CONFIDENCE_LIMIT,
        fuzziness: ES_USE_FUZZY_MATCH,
        es_expand_contractions: ES_EXPAND_CONTRACTIONS,
        kendra_indexes: ALT_SEARCH_KENDRA_INDEXES,
        minimum_confidence_score: ALT_SEARCH_KENDRA_FAQ_CONFIDENCE_SCORE,
        qnaClientFilter: _.get(req, 'session.QNAClientFilter'),
        settings: req._settings,
        locale: _.get(req, '_locale.localeIdentified'),
        translation: _.get(req, '_translation.QuestionInBackupLanguage'),
        language: language,
    };

    let hit;
    [req, res, hit] = await runFirstQuery(req, query_params, KENDRA_FAQ_INDEX, res);

    // ES fallback if KendraFAQ fails
    if (!hit && KENDRA_FAQ_INDEX && KENDRA_FAQ_ES_FALLBACK) {
        hit = await opensearchFallback(req, query_params);
    }

    // Check if item contains redirects to a targeted Kendra query
    if (hit && _.get(hit, 'kendraRedirectQueryText')) {
        [req, res, hit] = await kendraRedirect(hit, req, res, ALT_SEARCH_KENDRA_FALLBACK_CONFIDENCE_SCORE, query_params);
    }

    if (hit) {
        res.got_hits = 1; // response flag, used in logging / kibana
        if (!hit.debug) {
            hit.debug = [];
        }

        if (hit.type === 'text') {
            hit = copyPassageFieldToAnswerFields(hit);
            // Run any configured QA Summary options on the text passage result
            [hit, errors] = await runLlmQa(req, hit);
        }
    } else {
        [req, res, hit, errors] = await useFallbackMethod(req, res, query_params);
    }

    if (!hit) {
        hit = await getNoHitsResponse(query_params, res, req, errors, KENDRA_FAQ_INDEX);
    }

    if (hit) {
        res = setSessionTopic(hit, res);
        // run handlebars template processing
        hit = await handlebars(req, res, hit);

        hit = encryptConditionalChainingIfSet(hit);

        // update the res object with the hit results
        res = updateResWithHit(req, res, hit);

        if (RUN_LAMBDAHOOK_FROM_QUERY_STEP) {
            [req, res, hit] = await invokeLambdaHook(hit, req, res);
        }
    }
    return [req, res, hit, errors];
}
exports.getHit = getHit;