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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.util.SystemInfo
import com.intellij.util.Alarm
import com.intellij.util.AlarmFactory
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import org.apache.commons.collections4.queue.CircularFifoQueue
import software.aws.toolkits.jetbrains.core.coroutines.applicationCoroutineScope
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererUnknownLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.telemetry.CodewhispererAutomatedTriggerType
import software.aws.toolkits.telemetry.CodewhispererPreviousSuggestionState
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import java.time.Duration
import java.time.Instant
import kotlin.math.exp

data class ClassifierResult(val shouldTrigger: Boolean, val calculatedResult: Double = 0.0)

class CodeWhispererAutoTriggerService : CodeWhispererAutoTriggerHandler, Disposable {
    private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)
    private val previousUserTriggerDecisions = CircularFifoQueue<CodewhispererPreviousSuggestionState>(5)

    private var lastInvocationTime: Instant? = null
    private var lastInvocationLineNum: Int? = null

    init {
        scheduleReset()
    }

    fun addPreviousDecision(decision: CodewhispererPreviousSuggestionState) {
        previousUserTriggerDecisions.add(decision)
    }

    // a util wrapper
    fun tryInvokeAutoTrigger(editor: Editor, triggerType: CodeWhispererAutomatedTriggerType): Job? {
        // only needed for Classifier group, thus calculate it lazily
        val classifierResult: ClassifierResult by lazy { shouldTriggerClassifier(editor, triggerType.telemetryType) }
        val language = runReadAction {
            FileDocumentManager.getInstance().getFile(editor.document)?.programmingLanguage()
        } ?: CodeWhispererUnknownLanguage.INSTANCE

        // we need classifier result for any type of triggering for classifier group for supported languages
        return if (
            (language.isClassifierSupported() && CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.Classifier) ||
            language.isAllClassifier()
        ) {
            triggerType.calculationResult = classifierResult.calculatedResult

            when (triggerType) {
                // only invoke service if result > threshold for classifier trigger
                is CodeWhispererAutomatedTriggerType.Classifier -> run {
                    if (classifierResult.shouldTrigger) {
                        invoke(editor, triggerType)
                    } else {
                        null
                    }
                }

                // invoke whatever the result is for char / enter based trigger
                else -> run {
                    invoke(editor, triggerType)
                }
            }
        } else {
            invoke(editor, triggerType)
        }
    }

    // real auto trigger logic
    fun invoke(editor: Editor, triggerType: CodeWhispererAutomatedTriggerType): Job? {
        if (!CodeWhispererService.getInstance().canDoInvocation(editor, CodewhispererTriggerType.AutoTrigger)) {
            return null
        }

        lastInvocationTime = Instant.now()
        lastInvocationLineNum = runReadAction { editor.caretModel.visualPosition.line }

        val latencyContext = LatencyContext().apply {
            codewhispererPreprocessingStart = System.nanoTime()
            codewhispererEndToEndStart = System.nanoTime()
        }

        val coroutineScope = applicationCoroutineScope()

        return when (triggerType) {
            is CodeWhispererAutomatedTriggerType.IdleTime -> run {
                coroutineScope.launch {
                    // TODO: potential race condition between hasExistingInvocation and entering edt
                    // but in that case we will just return in performAutomatedTriggerAction
                    while (!CodeWhispererInvocationStatus.getInstance().hasEnoughDelayToInvokeCodeWhisperer() ||
                        CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()
                    ) {
                        if (!isActive) return@launch
                        delay(CodeWhispererConstants.IDLE_TIME_CHECK_INTERVAL)
                    }
                    runInEdt {
                        if (CodeWhispererInvocationStatus.getInstance().isPopupActive()) return@runInEdt
                        performAutomatedTriggerAction(editor, CodeWhispererAutomatedTriggerType.IdleTime(), latencyContext)
                    }
                }
            }

            else -> run {
                coroutineScope.launch {
                    performAutomatedTriggerAction(editor, triggerType, latencyContext)
                }
            }
        }
    }

    private fun scheduleReset() {
        if (!alarm.isDisposed) {
            alarm.addRequest({ resetPreviousStates() }, Duration.ofSeconds(120).toMillis())
        }
    }

    private fun resetPreviousStates() {
        try {
            previousUserTriggerDecisions.clear()
            lastInvocationLineNum = null
            lastInvocationTime = null
        } finally {
            scheduleReset()
        }
    }

    fun shouldTriggerClassifier(
        editor: Editor,
        automatedTriggerType: CodewhispererAutomatedTriggerType = CodewhispererAutomatedTriggerType.Classifier // TODO: need this?
    ): ClassifierResult {
        val caretContext = runReadAction { CodeWhispererEditorUtil.extractCaretContext(editor) }
        val language = runReadAction {
            FileDocumentManager.getInstance().getFile(editor.document)?.programmingLanguage()
        } ?: CodeWhispererUnknownLanguage.INSTANCE
        val caretPosition = runReadAction { CodeWhispererEditorUtil.getCaretPosition(editor) }

        // tryClassifier with only the supported language
        if (!language.isClassifierSupported()) {
            return ClassifierResult(false)
        }

        val leftContextLines = caretContext.leftFileContext.split(Regex("\r?\n"))
        val leftContextAtCurrentLine = caretContext.leftContextOnCurrentLine
        val keyword = leftContextAtCurrentLine.trim().split(" ").let { tokens ->
            if (tokens.size - 1 >= 0) tokens[tokens.size - 1] else ""
        }

        val lengthOfLeftCurrent = leftContextAtCurrentLine.length
        val lengthOfLeftPrev = if (leftContextLines.size - 2 >= 0) {
            leftContextLines[leftContextLines.size - 2].length.toDouble()
        } else {
            0.0
        }

        val rightContext = caretContext.rightFileContext
        val lengthOfRight = rightContext.trim().length
        val triggerTypeCoefficient = CodeWhispererClassifierConstants.triggerTypeCoefficientMap[automatedTriggerType] ?: 0.0

        val osCoefficient: Double = if (SystemInfo.isMac) {
            CodeWhispererClassifierConstants.osMap["Mac OS X"] ?: 0.0
        } else if (SystemInfo.isWindows) {
            val osVersion = SystemInfo.OS_VERSION
            if (osVersion.contains("11", true)) {
                CodeWhispererClassifierConstants.osMap["Windows 10"]
            } else if (osVersion.contains("10", true)) {
                CodeWhispererClassifierConstants.osMap["Windows 10"]
            } else if (osVersion.contains("7", true)) {
                CodeWhispererClassifierConstants.osMap["Windows 7"]
            } else {
                0.0
            }
        } else {
            0.0
        } ?: 0.0

        val lastCharCoefficient = if (leftContextAtCurrentLine.length - 1 >= 0) {
            CodeWhispererClassifierConstants.coefficientsMap[leftContextAtCurrentLine[leftContextAtCurrentLine.length - 1].toString()] ?: 0.0
        } else {
            0.0
        }

        val keywordCoefficient = CodeWhispererClassifierConstants.coefficientsMap[keyword] ?: 0.0
        val languageCoefficient = CodeWhispererClassifierConstants.languageMap[language] ?: 0.0
        val ideCoefficient = 0

        val lineDiff = lastInvocationLineNum?.let { (caretPosition.line.toDouble() - it) } ?: 0.0

        var previousOneAccept: Double = 0.0
        var previousOneReject: Double = 0.0
        var previousOneOther: Double = 0.0
        val previousOneDecision = CodeWhispererTelemetryService.getInstance().previousUserTriggerDecision
        if (previousOneDecision == null) {
            previousOneAccept = 0.0
            previousOneReject = 0.0
            previousOneOther = 0.0
        } else {
            previousOneAccept =
                if (previousOneDecision == CodewhispererPreviousSuggestionState.Accept) CodeWhispererClassifierConstants.prevDecisionAcceptCoefficient else 0.0
            previousOneReject =
                if (previousOneDecision == CodewhispererPreviousSuggestionState.Reject) CodeWhispererClassifierConstants.prevDecisionRejectCoefficient else 0.0
            previousOneOther =
                if (
                    previousOneDecision != CodewhispererPreviousSuggestionState.Accept &&
                    previousOneDecision != CodewhispererPreviousSuggestionState.Reject
                ) {
                    CodeWhispererClassifierConstants.prevDecisionOtherCoefficient
                } else {
                    0.0
                }
        }

        val resultBeforeSigmoid =
            CodeWhispererClassifierConstants.lengthofRightCoefficient * VariableTypeNeedNormalize.LenRight.normalize(lengthOfRight.toDouble()) +
                CodeWhispererClassifierConstants.lengthOfLeftCurrentCoefficient *
                VariableTypeNeedNormalize.LenLeftCur.normalize(lengthOfLeftCurrent.toDouble()) +
                CodeWhispererClassifierConstants.lengthOfLeftPrevCoefficient * VariableTypeNeedNormalize.LenLeftPrev.normalize(lengthOfLeftPrev) +
                CodeWhispererClassifierConstants.lineNumCoefficient * VariableTypeNeedNormalize.LineNum.normalize(caretPosition.line.toDouble()) +
                CodeWhispererClassifierConstants.cursorOffsetCoefficient * VariableTypeNeedNormalize.Cursor.normalize(caretPosition.offset.toDouble()) +
                CodeWhispererClassifierConstants.lineDiffCoefficient * VariableTypeNeedNormalize.LineDiff.normalize(lineDiff) +
                languageCoefficient +
                osCoefficient +
                triggerTypeCoefficient +
                lastCharCoefficient +
                keywordCoefficient +
                ideCoefficient +
                previousOneAccept +
                previousOneReject +
                previousOneOther +
                CodeWhispererClassifierConstants.intercept

        val shouldTrigger = sigmoid(resultBeforeSigmoid) > getThreshold()
        return ClassifierResult(shouldTrigger, sigmoid(resultBeforeSigmoid))
    }

    override fun dispose() {}

    companion object {
        private const val triggerThreshold: Double = 0.4

        fun getInstance(): CodeWhispererAutoTriggerService = service()

        fun getThreshold(): Double = triggerThreshold

        fun sigmoid(x: Double): Double = 1 / (1 + exp(-x))
    }
}

private enum class VariableTypeNeedNormalize {
    Cursor {
        override fun normalize(value: Double): Double = (value - minn.cursor) / (maxx.cursor - minn.cursor)
    },
    LineNum {
        override fun normalize(value: Double): Double = (value - minn.lineNum) / (maxx.lineNum - minn.lineNum)
    },
    LenLeftCur {
        override fun normalize(value: Double): Double = (value - minn.lenLeftCur) / (maxx.lenLeftCur - minn.lenLeftCur)
    },
    LenLeftPrev {
        override fun normalize(value: Double): Double = (value - minn.lenLeftPrev) / (maxx.lenLeftPrev - minn.lenLeftPrev)
    },
    LenRight {
        override fun normalize(value: Double): Double = (value - minn.lenRight) / (maxx.lenRight - minn.lenRight)
    },
    LineDiff {
        override fun normalize(value: Double): Double = (value - minn.lineDiff) / (maxx.lineDiff - minn.lineDiff)
    };

    abstract fun normalize(value: Double): Double

    data class NormalizedCoefficients(
        val cursor: Double,
        val lineNum: Double,
        val lenLeftCur: Double,
        val lenLeftPrev: Double,
        val lenRight: Double,
        val lineDiff: Double,
    )

    companion object {
        private val maxx = NormalizedCoefficients(
            cursor = 84716.0,
            lineNum = 2033.0,
            lenLeftCur = 157.0,
            lenLeftPrev = 157.0,
            lenRight = 10239.0,
            lineDiff = 270.0,
        )

        private val minn = NormalizedCoefficients(
            cursor = 1.0,
            lineNum = 0.0,
            lenLeftCur = 0.0,
            lenLeftPrev = 0.0,
            lenRight = 0.0,
            lineDiff = -28336.0,
        )
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPlainText
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript
import software.aws.toolkits.telemetry.CodewhispererAutomatedTriggerType

object CodeWhispererClassifierConstants {
    val osMap: Map<String, Double> = mapOf(
        "Mac OS X" to 0.0741,
        "win32" to 0.018,
        "Windows 10" to 0.2765,
        "Windows 7" to 0.0335,
    )

    // these are used for 100% classifier driven auto trigger
    val triggerTypeCoefficientMap: Map<CodewhispererAutomatedTriggerType, Double> = mapOf(
        CodewhispererAutomatedTriggerType.SpecialCharacters to 0.0624,
        CodewhispererAutomatedTriggerType.Enter to 0.207
    )

    val languageMap: Map<CodeWhispererProgrammingLanguage, Double> = mapOf(
        CodeWhispererPython.INSTANCE to -0.2656,
        CodeWhispererJava.INSTANCE to -0.3737,
        CodeWhispererJavaScript.INSTANCE to -0.3611,
        CodeWhispererCsharp.INSTANCE to 0.0,
        CodeWhispererPlainText.INSTANCE to 0.0,
        CodeWhispererTypeScript.INSTANCE to -0.3931,
        CodeWhispererTsx.INSTANCE to -0.393076,
        CodeWhispererJsx.INSTANCE to -0.361063
    )

    // other metadata coefficient
    const val lineNumCoefficient = 2.4507

    const val cursorOffsetCoefficient = -1.9998

    // length of the current line of left_context
    const val lengthOfLeftCurrentCoefficient = -1.0103

    // length of the previous line of left context
    const val lengthOfLeftPrevCoefficient = 0.4099

    // lenght of right_context
    const val lengthofRightCoefficient = -0.426

    const val lineDiffCoefficient = 0.377

    const val prevDecisionAcceptCoefficient = 1.2233

    const val prevDecisionRejectCoefficient = -0.1507

    const val prevDecisionOtherCoefficient = -0.0093

    // intercept of logistic regression classifier
    const val intercept = -0.04756079

    val coefficientsMap = mapOf<String, Double>(
        "False" to -0.1006,
        "None" to 0.1757,
        "True" to -0.1783,
        "abstract" to -0.1583,
        "and" to -0.0246,
        "any" to -0.5444,
        "arguments" to 0.3351,
        "as" to 0.054,
        "assert" to 0.3538,
        "async" to -0.1247,
        "await" to 0.0205,
        "base" to -0.0734,
        "bool" to 0.2111,
        "boolean" to 0.1225,
        "break" to -0.241,
        "byte" to 0.0113,
        "case" to 0.3065,
        "catch" to 1.0111,
        "char" to -0.0324,
        "checked" to 0.12,
        "class" to 0.1572,
        "const" to -0.5621,
        "continue" to -0.7951,
        "debugger" to -0.3898,
        "decimal" to -0.0162,
        "def" to 0.7647,
        "default" to 0.1289,
        "del" to -0.1192,
        "delegate" to 0.2523,
        "delete" to 0.0804,
        "do" to -0.3249,
        "double" to 0.0515,
        "elif" to 1.1261,
        "else" to -0.1974,
        "enum" to -0.2813,
        "eval" to -0.0527,
        "event" to 0.2892,
        "except" to 1.3827,
        "explicit" to 0.0743,
        "export" to -0.4735,
        "extends" to 0.2279,
        "extern" to -0.0482,
        "false" to 0.1996,
        "final" to -1.0326,
        "finally" to 0.0987,
        "fixed" to -0.0101,
        "float" to 0.2225,
        "for" to 0.398,
        "foreach" to 0.1189,
        "from" to 0.1864,
        "function" to 0.6337,
        "get" to -0.337,
        "global" to 0.1947,
        "goto" to 0.0,
        "if" to 0.2856,
        "implements" to 0.0509,
        "implicit" to 0.0,
        "import" to -0.4132,
        "in" to 0.0517,
        "instanceof" to -0.0527,
        "int" to -0.3033,
        "interface" to -0.1086,
        "internal" to 0.1165,
        "is" to 0.1365,
        "lambda" to 0.871,
        "let" to -0.8694,
        "lock" to 0.0728,
        "long" to 0.0014,
        "module" to 0.6073,
        "namespace" to -0.0201,
        "native" to 0.0,
        "new" to 0.3862,
        "nonlocal" to 0.0,
        "not" to 0.1825,
        "null" to 0.2631,
        "number" to 0.3205,
        "object" to 1.0271,
        "operator" to 0.166,
        "or" to 0.7321,
        "out" to 0.0059,
        "override" to 0.0071,
        "package" to -0.0929,
        "params" to -0.1764,
        "pass" to -0.1964,
        "private" to -0.2461,
        "protected" to -0.2725,
        "public" to -0.0519,
        "raise" to 0.7565,
        "readonly" to -0.4363,
        "ref" to 0.3455,
        "return" to 0.1601,
        "sbyte" to 0.0,
        "sealed" to 0.0,
        "short" to 0.0,
        "sizeof" to 0.0,
        "stackalloc" to 0.0,
        "static" to -0.3759,
        "strictfp" to 0.0,
        "string" to 0.3221,
        "struct" to 0.018,
        "super" to -0.105,
        "switch" to 0.3145,
        "synchronized" to 0.0097,
        "this" to 0.6475,
        "throw" to 1.0519,
        "throws" to -0.0815,
        "transient" to 0.0,
        "true" to 0.334,
        "try" to -0.2083,
        "type" to 0.0806,
        "typeof" to 0.3026,
        "uint" to -0.1173,
        "ulong" to 0.0,
        "unchecked" to 0.0,
        "unsafe" to 0.0,
        "ushort" to 0.0,
        "using" to 0.239,
        "var" to -0.9392,
        "virtual" to 0.0,
        "void" to 0.408,
        "volatile" to 0.0,
        "while" to 0.6905,
        "with" to 0.0446,
        "yield" to -0.0593,
        " " to 0.0134,
        "!" to -0.6027,
        "\"" to -0.3134,
        "#" to -1.451,
        "$" to -0.5431,
        "%" to -0.4637,
        "&" to -0.2307,
        "'" to -0.3702,
        "(" to 0.3116,
        ")" to -0.6789,
        "*" to -1.5531,
        "+" to -0.6805,
        "," to -0.2927,
        "-" to -0.9242,
        "." to -0.6001,
        "/" to -1.5167,
        "0" to -1.3701,
        "1" to -1.1693,
        "2" to -1.2146,
        "3" to -0.5654,
        "4" to -1.1667,
        "5" to -1.0519,
        "6" to -1.5824,
        "7" to -1.4413,
        "8" to -1.6181,
        "9" to -1.2799,
        ":" to 0.2868,
        ";" to -1.2036,
        "<" to -0.5011,
        "=" to -0.1256,
        ">" to -0.5585,
        "?" to -0.7477,
        "@" to -0.7144,
        "A" to -0.3267,
        "B" to -0.0695,
        "C" to -0.3239,
        "D" to -0.452,
        "E" to -0.3843,
        "F" to -0.4099,
        "G" to -0.4143,
        "H" to -0.124,
        "I" to -0.2478,
        "J" to -0.4139,
        "K" to -0.5061,
        "L" to -0.245,
        "M" to -0.3303,
        "N" to -0.5471,
        "O" to -0.3772,
        "P" to -0.3847,
        "Q" to -0.1783,
        "R" to -0.3994,
        "S" to -0.2535,
        "T" to -0.3746,
        "U" to -0.6347,
        "V" to -0.2543,
        "W" to -0.4353,
        "X" to -0.105,
        "Y" to -0.6737,
        "Z" to -0.3218,
        "[" to -0.3086,
        "\\" to -1.5857,
        "]" to -0.9263,
        "^" to -0.2029,
        "_" to -0.291,
        "`" to -1.0691,
        "a" to -0.2392,
        "b" to -0.2473,
        "c" to -0.228,
        "d" to -0.13,
        "e" to -0.1277,
        "f" to -0.1625,
        "g" to -0.2912,
        "h" to -0.0463,
        "i" to -0.1384,
        "j" to -0.6122,
        "k" to -0.3229,
        "l" to -0.1911,
        "m" to -0.2133,
        "n" to -0.0712,
        "o" to -0.1545,
        "p" to -0.3171,
        "q" to -0.3805,
        "r" to -0.0471,
        "s" to -0.2272,
        "t" to -0.0944,
        "u" to -0.0051,
        "v" to -0.2343,
        "w" to -0.0309,
        "x" to -0.1762,
        "y" to -0.0288,
        "z" to -0.7588,
        "{" to 0.0843,
        "|" to -0.2508,
        "}" to -0.9324,
        "~" to -0.2903
    )
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.util.messages.Topic
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import java.time.Duration
import java.time.Instant
import java.util.concurrent.atomic.AtomicBoolean

class CodeWhispererInvocationStatus {
    private val isInvokingCodeWhisperer: AtomicBoolean = AtomicBoolean(false)
    private var timeAtLastInvocationComplete: Instant? = null
    var timeAtLastDocumentChanged: Instant = Instant.now()
        private set
    private var isPopupActive: Boolean = false
    private var timeAtLastInvocationStart: Instant? = null
    var popupStartTimestamp: Instant? = null
        private set

    fun checkExistingInvocationAndSet(): Boolean =
        if (isInvokingCodeWhisperer.getAndSet(true)) {
            LOG.debug { "Have existing CodeWhisperer invocation" }
            true
        } else {
            ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(true)
            LOG.debug { "Starting CodeWhisperer invocation" }
            false
        }

    fun hasExistingInvocation(): Boolean = isInvokingCodeWhisperer.get()

    fun finishInvocation() {
        if (isInvokingCodeWhisperer.compareAndSet(true, false)) {
            ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(false)
            LOG.debug { "Ending CodeWhisperer invocation" }
        }
    }

    fun setInvocationComplete() {
        timeAtLastInvocationComplete = Instant.now()
    }

    fun documentChanged() {
        timeAtLastDocumentChanged = Instant.now()
    }

    fun setPopupStartTimestamp() {
        popupStartTimestamp = Instant.now()
    }

    fun getTimeSinceDocumentChanged(): Double {
        val timeSinceDocumentChanged = Duration.between(timeAtLastDocumentChanged, Instant.now())
        val timeInDouble = timeSinceDocumentChanged.toMillis().toDouble()
        return timeInDouble
    }

    fun hasEnoughDelayToShowCodeWhisperer(): Boolean {
        val timeCanShowCodeWhisperer = timeAtLastDocumentChanged.plusMillis(CodeWhispererConstants.POPUP_DELAY)
        return timeCanShowCodeWhisperer.isBefore(Instant.now())
    }

    fun isPopupActive(): Boolean = isPopupActive

    fun setPopupActive(value: Boolean) {
        isPopupActive = value
    }

    fun setInvocationStart() {
        timeAtLastInvocationStart = Instant.now()
    }

    fun hasEnoughDelayToInvokeCodeWhisperer(): Boolean {
        val timeCanShowCodeWhisperer = timeAtLastInvocationStart?.plusMillis(CodeWhispererConstants.INVOCATION_INTERVAL) ?: return true
        return timeCanShowCodeWhisperer.isBefore(Instant.now())
    }

    companion object {
        private val LOG = getLogger<CodeWhispererInvocationStatus>()
        fun getInstance(): CodeWhispererInvocationStatus = service()
        val CODEWHISPERER_INVOCATION_STATE_CHANGED: Topic<CodeWhispererInvocationStateChangeListener> = Topic.create(
            "CodeWhisperer popup state changed",
            CodeWhispererInvocationStateChangeListener::class.java
        )
    }
}

interface CodeWhispererInvocationStateChangeListener {
    fun invocationStateChanged(value: Boolean) {}
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import com.intellij.openapi.components.service
import software.aws.toolkits.jetbrains.utils.runUnderProgressIfNeeded
import software.aws.toolkits.resources.message

class CodeWhispererLicenseInfoManager {
    private val licenseLinks by lazy {
        runUnderProgressIfNeeded(null, message("codewhisperer.loading_licenses"), cancelable = false) {
            this.javaClass.getResourceAsStream("/codewhisperer/licenses.json")?.use { resourceStream ->
                MAPPER.readValue<Map<String, String>>(resourceStream)
            } ?: throw RuntimeException("CodeWhisperer license info not found")
        }
    }

    fun getLicenseLink(code: String) = licenseLinks.getOrDefault(code, "")

    companion object {
        fun getInstance(): CodeWhispererLicenseInfoManager = service()
        private val MAPPER = jacksonObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.util.TextRange
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiFileFactory
import com.intellij.psi.codeStyle.CodeStyleManager
import com.intellij.util.LocalTimeCounter
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.amazon.awssdk.services.codewhispererruntime.model.Span
import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk
import kotlin.math.max

class CodeWhispererRecommendationManager {
    fun reformat(requestContext: RequestContext, recommendation: Completion): Completion {
        val project = requestContext.project
        val editor = requestContext.editor
        val document = editor.document

        // startOffset is the offset at the start of user input since invocation
        val invocationStartOffset = requestContext.caretPosition.offset
        val startOffsetSinceUserInput = editor.caretModel.offset

        // Create a temp file for capturing reformatted text and updated content spans
        val tempPsiFile = PsiDocumentManager.getInstance(project).getPsiFile(document)?.let { psiFile ->
            PsiFileFactory.getInstance(project).createFileFromText(
                "codewhisperer_temp",
                psiFile.fileType,
                document.text,
                LocalTimeCounter.currentTime(),
                true
            )
        }
        val tempDocument = tempPsiFile?.let { psiFile ->
            PsiDocumentManager.getInstance(project).getDocument(psiFile)
        } ?: return recommendation

        val endOffset = invocationStartOffset + recommendation.content().length
        if (startOffsetSinceUserInput > endOffset) return recommendation
        WriteCommandAction.runWriteCommandAction(project) {
            tempDocument.insertString(invocationStartOffset, recommendation.content())
            PsiDocumentManager.getInstance(project).commitDocument(tempDocument)
        }
        val rangeMarkers = mutableMapOf<RangeMarker, Reference>()
        recommendation.references().forEach {
            val referenceStart = invocationStartOffset + it.recommendationContentSpan().start()
            if (referenceStart >= endOffset) return@forEach
            val tempEnd = invocationStartOffset + it.recommendationContentSpan().end()
            val referenceEnd = if (tempEnd <= endOffset) tempEnd else endOffset
            rangeMarkers[
                tempDocument.createRangeMarker(
                    referenceStart,
                    referenceEnd
                )
            ] = it
        }
        val tempRangeMarker = tempDocument.createRangeMarker(invocationStartOffset, endOffset)

        // Currently, only reformat(adjust line indent) starting from user's input
        WriteCommandAction.runWriteCommandAction(project) {
            CodeStyleManager.getInstance(project).adjustLineIndent(tempPsiFile, TextRange(startOffsetSinceUserInput, endOffset))
        }

        val reformattedRecommendation = tempDocument.getText(TextRange(tempRangeMarker.startOffset, tempRangeMarker.endOffset))

        val reformattedReferences = rangeMarkers.map { (rangeMarker, reference) ->
            reformatReference(reference, rangeMarker, invocationStartOffset)
        }
        return Completion.builder()
            .content(reformattedRecommendation)
            .references(reformattedReferences)
            .build()
    }

    /**
     * Build new reference with updated contentSpan(start and end). Since it's reformatted, take the new start and
     * end from the rangeMarker which automatically tracks the range after reformatting
     */
    fun reformatReference(originalReference: Reference, rangeMarker: RangeMarker, invocationStartOffset: Int): Reference {
        rangeMarker.apply {
            val documentContent = document.charsSequence

            // has to plus 1 because right boundary is exclusive
            val spanEndOffset = documentContent.subSequence(0, endOffset).indexOfLast { char -> char != '\n' } + 1
            return originalReference
                .toBuilder()
                .recommendationContentSpan(
                    Span.builder()
                        .start(startOffset - invocationStartOffset)
                        .end(spanEndOffset - invocationStartOffset)
                        .build()
                )
                .build()
        }
    }

    fun buildRecommendationChunks(
        recommendation: String,
        matchingSymbols: List<Pair<Int, Int>>
    ): List<RecommendationChunk> = matchingSymbols
        .dropLast(1)
        .mapIndexed { index, (offset, inlayOffset) ->
            val end = matchingSymbols[index + 1].first - 1
            RecommendationChunk(recommendation.substring(offset, end), offset, inlayOffset)
        }

    fun buildDetailContext(
        requestContext: RequestContext,
        userInput: String,
        recommendations: List<Completion>,
        requestId: String,
    ): List<DetailContext> {
        val seen = mutableSetOf<String>()
        return recommendations.map {
            val isDiscardedByUserInput = !it.content().startsWith(userInput) || it.content() == userInput
            if (isDiscardedByUserInput) {
                return@map DetailContext(requestId, it, it, isDiscarded = true, isTruncatedOnRight = false, rightOverlap = "")
            }

            val overlap = findRightContextOverlap(requestContext, it)
            val overlapIndex = it.content().lastIndexOf(overlap)
            val truncatedContent =
                if (overlap.isNotEmpty() && overlapIndex >= 0) {
                    it.content().substring(0, overlapIndex)
                } else {
                    it.content()
                }
            val truncated = it.toBuilder()
                .content(truncatedContent)
                .build()
            val isDiscardedByUserInputForTruncated = !truncated.content().startsWith(userInput) || truncated.content() == userInput
            if (isDiscardedByUserInputForTruncated) {
                return@map DetailContext(requestId, it, truncated, isDiscarded = true, isTruncatedOnRight = true, rightOverlap = overlap)
            }

            val reformatted = reformat(requestContext, truncated)
            val isDiscardedByRightContextTruncationDedupe = !seen.add(reformatted.content())
            DetailContext(
                requestId,
                it,
                reformatted,
                isDiscardedByRightContextTruncationDedupe,
                truncated.content().length != it.content().length,
                overlap
            )
        }
    }

    private fun findRightContextOverlap(
        requestContext: RequestContext,
        recommendation: Completion
    ): String {
        val document = requestContext.editor.document
        val caret = requestContext.editor.caretModel.primaryCaret
        val rightContext = document.charsSequence.subSequence(caret.offset, document.charsSequence.length).toString()
        val recommendationContent = recommendation.content()
        val rightContextFirstLine = rightContext.substringBefore("\n")
        val overlap =
            if (rightContextFirstLine.isEmpty()) {
                val tempOverlap = overlap(recommendationContent, rightContext)
                if (tempOverlap.isEmpty()) overlap(recommendationContent.trimEnd(), rightContext.trimStart()) else tempOverlap
            } else {
                // this is necessary to prevent display issue if first line of right context is not empty
                var tempOverlap = overlap(recommendationContent, rightContext)
                if (tempOverlap.isEmpty()) {
                    tempOverlap = overlap(recommendationContent.trimEnd(), rightContext.trimStart())
                }
                if (recommendationContent.substring(0, recommendationContent.length - tempOverlap.length).none { it == '\n' }) {
                    tempOverlap
                } else {
                    ""
                }
            }
        return overlap
    }

    fun overlap(first: String, second: String): String {
        for (i in max(0, first.length - second.length) until first.length) {
            val suffix = first.substring(i)
            if (second.startsWith(suffix)) {
                return suffix
            }
        }
        return ""
    }

    companion object {
        fun getInstance(): CodeWhispererRecommendationManager = service()
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.intellij.codeInsight.CodeInsightSettings
import com.intellij.codeInsight.hint.HintManager
import com.intellij.openapi.application.ApplicationInfo
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.VisualPosition
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.openapi.util.Disposer
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiFile
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.messages.Topic
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.delay
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import software.amazon.awssdk.core.exception.SdkServiceException
import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.FileContext
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.ProgrammingLanguage
import software.amazon.awssdk.services.codewhispererruntime.model.RecommendationsWithReferencesPreference
import software.amazon.awssdk.services.codewhispererruntime.model.SupplementalContext
import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.getCaretPosition
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.WorkerContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.checkCompletionType
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.checkEmptyRecommendations
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
import software.aws.toolkits.jetbrains.utils.isInjectedText
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.CodewhispererCompletionType
import software.aws.toolkits.telemetry.CodewhispererSuggestionState
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import java.util.concurrent.TimeUnit

class CodeWhispererService {
    fun showRecommendationsInPopup(
        editor: Editor,
        triggerTypeInfo: TriggerTypeInfo,
        latencyContext: LatencyContext
    ) {
        val project = editor.project ?: return
        if (!isCodeWhispererEnabled(project)) return

        latencyContext.credentialFetchingStart = System.nanoTime()

        if (promptReAuth(project)) return

        latencyContext.credentialFetchingEnd = System.nanoTime()
        val psiFile = runReadAction { PsiDocumentManager.getInstance(project).getPsiFile(editor.document) }

        if (psiFile == null) {
            LOG.debug { "No PSI file for the current document" }
            if (triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
                showCodeWhispererInfoHint(editor, message("codewhisperer.trigger.document.unsupported"))
            }
            return
        }
        val isInjectedFile = runReadAction { psiFile.isInjectedText() }
        if (isInjectedFile) return

        val requestContext = try {
            getRequestContext(triggerTypeInfo, editor, project, psiFile, latencyContext)
        } catch (e: Exception) {
            LOG.debug { e.message.toString() }
            CodeWhispererTelemetryService.getInstance().sendFailedServiceInvocationEvent(project, e::class.simpleName)
            return
        }

        val language = requestContext.fileContextInfo.programmingLanguage
        if (!language.isCodeCompletionSupported()) {
            LOG.debug { "Programming language $language is not supported by CodeWhisperer" }
            if (triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
                showCodeWhispererInfoHint(
                    requestContext.editor,
                    message("codewhisperer.language.error", psiFile.fileType.name)
                )
            }
            return
        }

        LOG.debug {
            "Calling CodeWhisperer service, trigger type: ${triggerTypeInfo.triggerType}" +
                if (triggerTypeInfo.triggerType == CodewhispererTriggerType.AutoTrigger) {
                    ", auto-trigger type: ${triggerTypeInfo.automatedTriggerType}"
                } else {
                    ""
                }
        }

        val invocationStatus = CodeWhispererInvocationStatus.getInstance()
        if (invocationStatus.checkExistingInvocationAndSet()) {
            return
        }

        invokeCodeWhispererInBackground(requestContext)
    }

    private fun invokeCodeWhispererInBackground(requestContext: RequestContext) {
        val popup = CodeWhispererPopupManager.getInstance().initPopup()
        Disposer.register(popup) { CodeWhispererInvocationStatus.getInstance().finishInvocation() }

        val workerContexts = mutableListOf<WorkerContext>()
        // When popup is disposed we will cancel this coroutine. The only places popup can get disposed should be
        // from CodeWhispererPopupManager.cancelPopup() and CodeWhispererPopupManager.closePopup().
        // It's possible and ok that coroutine will keep running until the next time we check it's state.
        // As long as we don't show to the user extra info we are good.
        val coroutineScope = disposableCoroutineScope(popup)

        var states: InvocationContext? = null
        var lastRecommendationIndex = -1

        val responseIterable = CodeWhispererClientAdaptor.getInstance(requestContext.project).generateCompletionsPaginator(
            buildCodeWhispererRequest(requestContext.fileContextInfo, requestContext.supplementalContext)
        )
        coroutineScope.launch {
            try {
                var startTime = System.nanoTime()
                requestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime()
                requestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime()
                CodeWhispererInvocationStatus.getInstance().setInvocationStart()
                var requestCount = 0
                for (response in responseIterable) {
                    requestCount++
                    val endTime = System.nanoTime()
                    val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
                    startTime = endTime
                    val requestId = response.responseMetadata().requestId()
                    if (requestCount == 1) {
                        requestContext.latencyContext.codewhispererPostprocessingStart = System.nanoTime()
                        requestContext.latencyContext.paginationFirstCompletionTime = latency
                        requestContext.latencyContext.firstRequestId = requestId
                    }
                    if (response.nextToken().isEmpty()) {
                        requestContext.latencyContext.paginationAllCompletionsEnd = System.nanoTime()
                    }
                    val sessionId = response.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
                    val emptyRecommendations = checkEmptyRecommendations(response.completions())
                    val completionType = checkCompletionType(response.completions(), emptyRecommendations)
                    val responseContext = ResponseContext(sessionId, completionType)
                    logServiceInvocation(requestId, requestContext, responseContext, response.completions(), latency, null)
                    lastRecommendationIndex += response.completions().size
                    ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_CODE_COMPLETION_PERFORMED)
                        .onSuccess(requestContext.fileContextInfo)
                    CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
                        requestId,
                        requestContext,
                        responseContext,
                        lastRecommendationIndex,
                        true,
                        latency,
                        null
                    )

                    val validatedResponse = validateResponse(response)

                    runInEdt {
                        // If delay is not met, add them to the worker queue and process them later.
                        // On first response, workers queue must be empty. If there's enough delay before showing,
                        // process CodeWhisperer UI rendering and workers queue will remain empty throughout this
                        // CodeWhisperer session. If there's not enough delay before showing, the CodeWhisperer UI rendering task
                        // will be added to the workers queue.
                        // On subsequent responses, if they see workers queue is not empty, it means the first worker
                        // task hasn't been finished yet, in this case simply add another task to the queue. If they
                        // see worker queue is empty, the previous tasks must have been finished before this. In this
                        // case render CodeWhisperer UI directly.
                        val workerContext = WorkerContext(requestContext, responseContext, validatedResponse, popup)
                        if (workerContexts.isNotEmpty()) {
                            workerContexts.add(workerContext)
                        } else {
                            if (states == null && !popup.isDisposed &&
                                !CodeWhispererInvocationStatus.getInstance().hasEnoughDelayToShowCodeWhisperer()
                            ) {
                                // It's the first response, and no enough delay before showing
                                projectCoroutineScope(requestContext.project).launch {
                                    while (!CodeWhispererInvocationStatus.getInstance().hasEnoughDelayToShowCodeWhisperer()) {
                                        delay(CodeWhispererConstants.POPUP_DELAY_CHECK_INTERVAL)
                                    }
                                    runInEdt {
                                        workerContexts.forEach {
                                            states = processCodeWhispererUI(it, states)
                                        }
                                        workerContexts.clear()
                                    }
                                }
                                workerContexts.add(workerContext)
                            } else {
                                // Have enough delay before showing for the first response, or it's subsequent responses
                                states = processCodeWhispererUI(workerContext, states)
                            }
                        }
                    }
                    if (!isActive) {
                        // If job is cancelled before we do another request, don't bother making
                        // another API call to save resources
                        LOG.debug { "Skipping sending remaining requests on CodeWhisperer session exit" }
                        break
                    }
                }
            } catch (e: Exception) {
                val requestId: String
                val sessionId: String
                val displayMessage: String
                if (e is CodeWhispererException) {
                    requestId = e.requestId() ?: ""
                    sessionId = e.awsErrorDetails().sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
                    displayMessage = e.awsErrorDetails().errorMessage() ?: message("codewhisperer.trigger.error.server_side")
                } else if (e is software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException) {
                    requestId = e.requestId() ?: ""
                    sessionId = e.awsErrorDetails().sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
                    displayMessage = e.awsErrorDetails().errorMessage() ?: message("codewhisperer.trigger.error.server_side")
                } else {
                    requestId = ""
                    sessionId = ""
                    val statusCode = if (e is SdkServiceException) e.statusCode() else 0
                    displayMessage =
                        if (statusCode >= 500) {
                            message("codewhisperer.trigger.error.server_side")
                        } else {
                            message("codewhisperer.trigger.error.client_side")
                        }
                    if (statusCode < 500) {
                        LOG.debug(e) { "Error invoking CodeWhisperer service" }
                    }
                }
                val exceptionType = e::class.simpleName
                val responseContext = ResponseContext(sessionId, CodewhispererCompletionType.Unknown)
                logServiceInvocation(requestId, requestContext, responseContext, emptyList(), null, exceptionType)
                CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
                    requestId,
                    requestContext,
                    responseContext,
                    lastRecommendationIndex,
                    false,
                    0.0,
                    exceptionType
                )

                if (e is ThrottlingException &&
                    e.message == CodeWhispererConstants.THROTTLING_MESSAGE
                ) {
                    CodeWhispererExplorerActionManager.getInstance().setSuspended(requestContext.project)
                    if (requestContext.triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
                        notifyErrorCodeWhispererUsageLimit(requestContext.project)
                    }
                } else {
                    if (requestContext.triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
                        // We should only show error hint when CodeWhisperer popup is not visible,
                        // and make it silent if CodeWhisperer popup is showing.
                        runInEdt {
                            if (!CodeWhispererInvocationStatus.getInstance().isPopupActive()) {
                                showCodeWhispererErrorHint(requestContext.editor, displayMessage)
                            }
                        }
                    }
                }
                CodeWhispererInvocationStatus.getInstance().finishInvocation()
                runInEdt {
                    states?.let {
                        CodeWhispererPopupManager.getInstance().updatePopupPanel(
                            it,
                            CodeWhispererPopupManager.getInstance().sessionContext
                        )
                    }
                }
            } finally {
                CodeWhispererInvocationStatus.getInstance().setInvocationComplete()
            }
        }
    }

    @RequiresEdt
    private fun processCodeWhispererUI(workerContext: WorkerContext, currStates: InvocationContext?): InvocationContext? {
        val requestContext = workerContext.requestContext
        val responseContext = workerContext.responseContext
        val response = workerContext.response
        val popup = workerContext.popup
        val requestId = response.responseMetadata().requestId()

        // At this point when we are in EDT, the state of the popup will be thread-safe
        // across this thread execution, so if popup is disposed, we will stop here.
        // This extra check is needed because there's a time between when we get the response and
        // when we enter the EDT.
        if (popup.isDisposed) {
            LOG.debug { "Stop showing CodeWhisperer recommendations on CodeWhisperer session exit. RequestId: $requestId" }
            return null
        }

        if (requestContext.editor.isDisposed) {
            LOG.debug { "Stop showing CodeWhisperer recommendations since editor is disposed. RequestId: $requestId" }
            return null
        }

        if (response.nextToken().isEmpty()) {
            CodeWhispererInvocationStatus.getInstance().finishInvocation()
        }

        val caretMovement = CodeWhispererEditorManager.getInstance().getCaretMovement(
            requestContext.editor,
            requestContext.caretPosition
        )
        val isPopupShowing: Boolean
        val nextStates: InvocationContext?
        if (currStates == null) {
            // first response
            nextStates = initStates(requestContext, responseContext, response, caretMovement, popup)
            isPopupShowing = false

            // receiving a null state means caret has moved backward or there's a conflict with
            // Intellisense popup, so we are going to cancel the job
            if (nextStates == null) {
                LOG.debug { "Cancelling popup and exiting CodeWhisperer session. RequestId: $requestId" }
                CodeWhispererPopupManager.getInstance().cancelPopup(popup)
                return null
            }
        } else {
            // subsequent responses
            nextStates = updateStates(currStates, response)
            isPopupShowing = checkRecommendationsValidity(currStates, false)
        }

        val hasAtLeastOneValid = checkRecommendationsValidity(nextStates, response.nextToken().isEmpty())

        // If there are no recommendations at all in this session, we need to manually send the user decision event here
        // since it won't be sent automatically later
        if (nextStates.recommendationContext.details.isEmpty() && response.nextToken().isEmpty()) {
            LOG.debug { "Received just an empty list from this session, requestId: $requestId" }
            CodeWhispererTelemetryService.getInstance().sendUserDecisionEvent(
                requestId,
                requestContext,
                responseContext,
                Completion.builder().build(),
                -1,
                CodewhispererSuggestionState.Empty,
                nextStates.recommendationContext.details.size
            )
        }
        if (!hasAtLeastOneValid) {
            if (response.nextToken().isEmpty()) {
                LOG.debug { "None of the recommendations are valid, exiting CodeWhisperer session" }
                CodeWhispererPopupManager.getInstance().cancelPopup(popup)
                return null
            }
        } else {
            updateCodeWhisperer(nextStates, isPopupShowing)
        }
        return nextStates
    }

    private fun initStates(
        requestContext: RequestContext,
        responseContext: ResponseContext,
        response: GenerateCompletionsResponse,
        caretMovement: CaretMovement,
        popup: JBPopup
    ): InvocationContext? {
        val requestId = response.responseMetadata().requestId()
        val recommendations = response.completions()
        val visualPosition = requestContext.editor.caretModel.visualPosition

        if (CodeWhispererPopupManager.getInstance().hasConflictingPopups(requestContext.editor)) {
            LOG.debug { "Detect conflicting popup window with CodeWhisperer popup, not showing CodeWhisperer popup" }
            sendDiscardedUserDecisionEventForAll(requestContext, responseContext, recommendations)
            return null
        }
        if (caretMovement == CaretMovement.MOVE_BACKWARD) {
            LOG.debug { "Caret moved backward, discarding all of the recommendations. Request ID: $requestId" }
            sendDiscardedUserDecisionEventForAll(requestContext, responseContext, recommendations)
            return null
        }
        val userInputOriginal = CodeWhispererEditorManager.getInstance().getUserInputSinceInvocation(
            requestContext.editor,
            requestContext.caretPosition.offset
        )
        val userInput =
            if (caretMovement == CaretMovement.NO_CHANGE) {
                LOG.debug { "Caret position not changed since invocation. Request ID: $requestId" }
                ""
            } else {
                userInputOriginal.trimStart().also {
                    LOG.debug {
                        "Caret position moved forward since invocation. Request ID: $requestId, " +
                            "user input since invocation: $userInputOriginal, " +
                            "user input without leading spaces: $it"
                    }
                }
            }
        val detailContexts = CodeWhispererRecommendationManager.getInstance().buildDetailContext(
            requestContext,
            userInput,
            recommendations,
            requestId
        )
        val recommendationContext = RecommendationContext(detailContexts, userInputOriginal, userInput, visualPosition)
        return buildInvocationContext(requestContext, responseContext, recommendationContext, popup)
    }

    private fun updateStates(
        states: InvocationContext,
        response: GenerateCompletionsResponse
    ): InvocationContext {
        val recommendationContext = states.recommendationContext
        val details = recommendationContext.details
        val newDetailContexts = CodeWhispererRecommendationManager.getInstance().buildDetailContext(
            states.requestContext,
            recommendationContext.userInputSinceInvocation,
            response.completions(),
            response.responseMetadata().requestId()
        )
        Disposer.dispose(states)

        val updatedStates = states.copy(
            recommendationContext = recommendationContext.copy(details = details + newDetailContexts)
        )
        Disposer.register(states.popup, updatedStates)
        CodeWhispererPopupManager.getInstance().initPopupListener(updatedStates)
        return updatedStates
    }

    private fun checkRecommendationsValidity(states: InvocationContext, showHint: Boolean): Boolean {
        val details = states.recommendationContext.details

        // set to true when at least one is not discarded or empty
        val hasAtLeastOneValid = details.any { !it.isDiscarded && it.recommendation.content().isNotEmpty() }

        if (!hasAtLeastOneValid && showHint && states.requestContext.triggerTypeInfo.triggerType == CodewhispererTriggerType.OnDemand) {
            showCodeWhispererInfoHint(
                states.requestContext.editor,
                message("codewhisperer.popup.no_recommendations")
            )
        }
        return hasAtLeastOneValid
    }

    private fun updateCodeWhisperer(states: InvocationContext, recommendationAdded: Boolean) {
        CodeWhispererPopupManager.getInstance().changeStates(states, 0, "", true, recommendationAdded)
    }

    private fun sendDiscardedUserDecisionEventForAll(
        requestContext: RequestContext,
        responseContext: ResponseContext,
        recommendations: List<Completion>
    ) {
        val detailContexts = recommendations.map { DetailContext("", it, it, true, false) }
        val recommendationContext = RecommendationContext(detailContexts, "", "", VisualPosition(0, 0))

        CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
            requestContext,
            responseContext,
            recommendationContext,
            SessionContext(),
            false
        )
    }

    fun getRequestContext(
        triggerTypeInfo: TriggerTypeInfo,
        editor: Editor,
        project: Project,
        psiFile: PsiFile,
        latencyContext: LatencyContext
    ): RequestContext {
        // 1. file context
        val fileContext: FileContextInfo = runReadAction { FileContextProvider.getInstance(project).extractFileContext(editor, psiFile) }

        // the upper bound for supplemental context duration is 50ms
        // 2. supplemental context
        val startFetchingTimestamp = System.currentTimeMillis()
        val isTstFile = FileContextProvider.getInstance(project).isTestFile(psiFile)
        val supplementalContext = if (CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.CrossFile && !isTstFile) {
            runBlocking {
                try {
                    withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
                        FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext)
                    }
                } catch (e: Exception) {
                    if (e is TimeoutCancellationException) {
                        LOG.debug {
                            "Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms"
                        }
                        SupplementalContextInfo(
                            isUtg = isTstFile,
                            contents = emptyList(),
                            latency = System.currentTimeMillis() - startFetchingTimestamp,
                            targetFileName = fileContext.filename
                        )
                    } else {
                        LOG.error(e) { "Run into unexpected error when fetching supplemental context" }
                        null
                    }
                }
            }
        } else {
            null
        }

        // 3. caret position
        val caretPosition = runReadAction { getCaretPosition(editor) }

        // 4. connection
        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
        return RequestContext(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, latencyContext)
    }

    fun validateResponse(response: GenerateCompletionsResponse): GenerateCompletionsResponse {
        // If contentSpans in reference are not consistent with content(recommendations),
        // remove the incorrect references.
        val validatedRecommendations = response.completions().map {
            val validReferences = it.hasReferences() && it.references().isNotEmpty() &&
                it.references().none { reference ->
                    val span = reference.recommendationContentSpan()
                    span.start() > span.end() || span.start() < 0 || span.end() > it.content().length
                }
            if (validReferences) {
                it
            } else {
                it.toBuilder().references(DefaultSdkAutoConstructList.getInstance()).build()
            }
        }

        return response.toBuilder().completions(validatedRecommendations).build()
    }

    private fun buildInvocationContext(
        requestContext: RequestContext,
        responseContext: ResponseContext,
        recommendationContext: RecommendationContext,
        popup: JBPopup
    ): InvocationContext {
        addPopupChildDisposables(popup)
        // Creating a disposable for managing all listeners lifecycle attached to the popup.
        // previously(before pagination) we use popup as the parent disposable.
        // After pagination, listeners need to be updated as states are updated, for the same popup,
        // so disposable chain becomes popup -> disposable -> listeners updates, and disposable gets replaced on every
        // state update.
        val states = InvocationContext(requestContext, responseContext, recommendationContext, popup)
        Disposer.register(popup, states)
        CodeWhispererPopupManager.getInstance().initPopupListener(states)
        return states
    }

    private fun addPopupChildDisposables(popup: JBPopup) {
        val originalTabExitsBracketsAndQuotes = CodeInsightSettings.getInstance().TAB_EXITS_BRACKETS_AND_QUOTES
        CodeInsightSettings.getInstance().TAB_EXITS_BRACKETS_AND_QUOTES = false
        Disposer.register(popup) {
            CodeInsightSettings.getInstance().TAB_EXITS_BRACKETS_AND_QUOTES = originalTabExitsBracketsAndQuotes
        }
        val originalAutoPopupCompletionLookup = CodeInsightSettings.getInstance().AUTO_POPUP_COMPLETION_LOOKUP
        CodeInsightSettings.getInstance().AUTO_POPUP_COMPLETION_LOOKUP = false
        Disposer.register(popup) {
            CodeInsightSettings.getInstance().AUTO_POPUP_COMPLETION_LOOKUP = originalAutoPopupCompletionLookup
        }
        Disposer.register(popup) {
            CodeWhispererPopupManager.getInstance().reset()
        }
    }

    private fun logServiceInvocation(
        requestId: String,
        requestContext: RequestContext,
        responseContext: ResponseContext,
        recommendations: List<Completion>,
        latency: Double?,
        exceptionType: String?
    ) {
        val recommendationLogs = recommendations.map { it.content().trimEnd() }
            .reduceIndexedOrNull { index, acc, recommendation -> "$acc\n[${index + 1}]\n$recommendation" }
        LOG.info {
            "SessionId: ${responseContext.sessionId}, " +
                "RequestId: $requestId, " +
                "Jetbrains IDE: ${ApplicationInfo.getInstance().fullApplicationName}, " +
                "IDE version: ${ApplicationInfo.getInstance().apiVersion}, " +
                "Filename: ${requestContext.fileContextInfo.filename}, " +
                "Left context of current line: ${requestContext.fileContextInfo.caretContext.leftContextOnCurrentLine}, " +
                "Cursor line: ${requestContext.caretPosition.line}, " +
                "Caret offset: ${requestContext.caretPosition.offset}, " +
                (latency?.let { "Latency: $latency, " } ?: "") +
                (exceptionType?.let { "Exception Type: $it, " } ?: "") +
                "Recommendations: \n${recommendationLogs ?: "None"}"
        }
    }

    fun canDoInvocation(editor: Editor, type: CodewhispererTriggerType): Boolean {
        editor.project?.let {
            if (!isCodeWhispererEnabled(it)) {
                return false
            }
        }

        if (type == CodewhispererTriggerType.AutoTrigger && !CodeWhispererExplorerActionManager.getInstance().isAutoEnabled()) {
            LOG.debug { "CodeWhisperer auto-trigger is disabled, not invoking service" }
            return false
        }

        if (CodeWhispererPopupManager.getInstance().hasConflictingPopups(editor)) {
            LOG.debug { "Find other active popup windows before triggering CodeWhisperer, not invoking service" }
            return false
        }

        if (CodeWhispererInvocationStatus.getInstance().isPopupActive()) {
            LOG.debug { "Find an existing CodeWhisperer popup window before triggering CodeWhisperer, not invoking service" }
            return false
        }
        return true
    }

    fun showCodeWhispererInfoHint(editor: Editor, message: String) {
        HintManager.getInstance().showInformationHint(editor, message, HintManager.UNDER)
    }

    fun showCodeWhispererErrorHint(editor: Editor, message: String) {
        HintManager.getInstance().showErrorHint(editor, message, HintManager.UNDER)
    }

    companion object {
        private val LOG = getLogger<CodeWhispererService>()
        val CODEWHISPERER_CODE_COMPLETION_PERFORMED: Topic<CodeWhispererCodeCompletionServiceListener> = Topic.create(
            "CodeWhisperer code completion service invoked",
            CodeWhispererCodeCompletionServiceListener::class.java
        )

        fun getInstance(): CodeWhispererService = service()
        const val KET_SESSION_ID = "x-amzn-SessionId"
        private var reAuthPromptShown = false

        fun markReAuthPromptShown() {
            reAuthPromptShown = true
        }

        fun hasReAuthPromptBeenShown() = reAuthPromptShown

        fun buildCodeWhispererRequest(
            fileContextInfo: FileContextInfo,
            supplementalContext: SupplementalContextInfo?
        ): GenerateCompletionsRequest {
            val programmingLanguage = ProgrammingLanguage.builder()
                .languageName(fileContextInfo.programmingLanguage.toCodeWhispererRuntimeLanguage().languageId)
                .build()
            val fileContext = FileContext.builder()
                .leftFileContent(fileContextInfo.caretContext.leftFileContext)
                .rightFileContent(fileContextInfo.caretContext.rightFileContext)
                .filename(fileContextInfo.filename)
                .programmingLanguage(programmingLanguage)
                .build()
            val supplementalContexts = supplementalContext?.contents?.map {
                SupplementalContext.builder()
                    .content(it.content)
                    .filePath(it.path)
                    .build()
            }.orEmpty()
            val includeCodeWithReference = if (CodeWhispererSettings.getInstance().isIncludeCodeWithReference()) {
                RecommendationsWithReferencesPreference.ALLOW
            } else {
                RecommendationsWithReferencesPreference.BLOCK
            }

            return GenerateCompletionsRequest.builder()
                .fileContext(fileContext)
                .supplementalContexts(supplementalContexts)
                .referenceTrackerConfiguration { it.recommendationsWithReferences(includeCodeWithReference) }
                .build()
        }
    }
}

data class RequestContext(
    val project: Project,
    val editor: Editor,
    val triggerTypeInfo: TriggerTypeInfo,
    val caretPosition: CaretPosition,
    val fileContextInfo: FileContextInfo,
    val supplementalContext: SupplementalContextInfo?,
    val connection: ToolkitConnection?,
    val latencyContext: LatencyContext
)

data class ResponseContext(
    val sessionId: String,
    val completionType: CodewhispererCompletionType
)

interface CodeWhispererCodeCompletionServiceListener {
    fun onSuccess(fileContextInfo: FileContextInfo) {}
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.service

import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.components.service
import org.jetbrains.annotations.VisibleForTesting
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.AwsToolkit
import java.util.concurrent.ConcurrentHashMap
import kotlin.reflect.KClass

/**
 * Component controlling codewhisperer user group settings
 */
@State(name = "codewhispererUserGroupSettings", storages = [Storage("aws.xml")])
class CodeWhispererUserGroupSettings : PersistentStateComponent<CodeWhispererUserGroupStates> {
    private var version: String? = null

    private val settings = ConcurrentHashMap<String, String>()
    override fun getState() = CodeWhispererUserGroupStates(
        version,
        settings
    )

    override fun loadState(state: CodeWhispererUserGroupStates) {
        version = state.version

        settings.clear()
        settings.putAll(state.settings)
    }

    inline fun <reified T : CodeWhispererGroup> getGroup(): T? = getGroup(T::class)

    fun <T : CodeWhispererGroup> getGroup(clazz: KClass<T>): T? {
        @Suppress("UNCHECKED_CAST")
        return when (clazz) {
            CodeWhispererUserGroup::class -> {
                tryOrNull {
                    settings[USER_GROUP_KEY]?.let {
                        CodeWhispererUserGroup.valueOf(it)
                    }
                } as T?
            }
            CodeWhispererExpThresholdGroup::class -> {
                tryOrNull {
                    settings[EXP_THRESHOLD_KEY]?.let {
                        CodeWhispererExpThresholdGroup.valueOf(it)
                    }
                } as T?
            }
            else -> null
        }
    }

    fun getUserGroup(): CodeWhispererUserGroup {
        if (version != AwsToolkit.PLUGIN_VERSION) {
            resetGroupSettings()
        }

        return getGroup<CodeWhispererUserGroup>() ?: determineUserGroup()
    }

    fun isExpThreshold(): Boolean {
        if (version != AwsToolkit.PLUGIN_VERSION) {
            resetGroupSettings()
        }

        val group = getGroup<CodeWhispererExpThresholdGroup>()
        return (group ?: determineThresholdGroup()) == CodeWhispererExpThresholdGroup.Exp
    }

    @VisibleForTesting
    fun getVersion() = version

    @VisibleForTesting
    fun determineUserGroup(): CodeWhispererUserGroup {
        val randomNum = Math.random()
        val group = if (randomNum < 1 / 2.0) {
            CodeWhispererUserGroup.Control
        } else {
            CodeWhispererUserGroup.CrossFile
        }

        settings[USER_GROUP_KEY] = group.name
        version = AwsToolkit.PLUGIN_VERSION

        return group
    }

    private fun determineThresholdGroup(): CodeWhispererExpThresholdGroup {
        val randomNum = Math.random()
        val group = if (randomNum < 1 / 2.0) {
            CodeWhispererExpThresholdGroup.Control
        } else {
            CodeWhispererExpThresholdGroup.Exp
        }

        settings[EXP_THRESHOLD_KEY] = group.name
        version = AwsToolkit.PLUGIN_VERSION

        return group
    }

    private fun resetGroupSettings() {
        version = null
        settings.clear()
    }

    companion object {
        fun getInstance(): CodeWhispererUserGroupSettings = service()

        // TODO: add into CodeWhispererGroup interface
        const val USER_GROUP_KEY = "userGroup"

        const val EXP_THRESHOLD_KEY = "expThreshold"
    }
}

data class CodeWhispererUserGroupStates(
    var version: String? = null,
    var settings: Map<String, String> = emptyMap()
)

interface CodeWhispererGroup

enum class CodeWhispererUserGroup : CodeWhispererGroup {
    Control,
    CrossFile,
    Classifier
}

enum class CodeWhispererExpThresholdGroup : CodeWhispererGroup {
    Control,
    Exp
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.settings

import com.intellij.icons.AllIcons
import com.intellij.openapi.options.BoundConfigurable
import com.intellij.openapi.options.SearchableConfigurable
import com.intellij.openapi.project.Project
import com.intellij.ui.dsl.builder.bindSelected
import com.intellij.ui.dsl.builder.panel
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.resources.message

//  As the connection is project-level, we need to make this project-level too (we have different config for Sono vs SSO users)
class CodeWhispererConfigurable(private val project: Project) :
    BoundConfigurable(message("aws.settings.codewhisperer.configurable.title")),
    SearchableConfigurable {
    private val codeWhispererSettings
        get() = CodeWhispererSettings.getInstance()

    private val isSso: Boolean
        get() = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.SSO

    override fun getId() = "aws.codewhisperer"

    override fun createPanel() = panel {
        val connect = project.messageBus.connect(disposable ?: error("disposable wasn't initialized by framework"))
        val invoke = isCodeWhispererEnabled(project)

        // TODO: can we remove message bus subscribe and solely use visible(boolean) / enabled(boolean), consider multi project cases
        row {
            label(message("aws.settings.codewhisperer.warning")).apply {
                component.icon = AllIcons.General.Warning
            }.apply {
                visible(!invoke)
                connect.subscribe(
                    ToolkitConnectionManagerListener.TOPIC,
                    object : ToolkitConnectionManagerListener {
                        override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                            visible(!isCodeWhispererEnabled(project))
                        }
                    }
                )
            }
        }

        row {
            checkBox(message("aws.settings.codewhisperer.include_code_with_reference")).apply {
                connect.subscribe(
                    ToolkitConnectionManagerListener.TOPIC,
                    object : ToolkitConnectionManagerListener {
                        override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                            enabled(isCodeWhispererEnabled(project) && !isSso)
                        }
                    }
                )
                enabled(invoke && !isSso)
                bindSelected(codeWhispererSettings::isIncludeCodeWithReference, codeWhispererSettings::toggleIncludeCodeWithReference)
            }
        }.rowComment(message("aws.settings.codewhisperer.include_code_with_reference.tooltip"))

        row {
            checkBox(message("aws.settings.codewhisperer.automatic_import_adder")).apply {
                connect.subscribe(
                    ToolkitConnectionManagerListener.TOPIC,
                    object : ToolkitConnectionManagerListener {
                        override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                            enabled(isCodeWhispererEnabled(project))
                        }
                    }
                )
                enabled(invoke)
                bindSelected(codeWhispererSettings::isImportAdderEnabled, codeWhispererSettings::toggleImportAdder)
            }
        }.rowComment(message("aws.settings.codewhisperer.automatic_import_adder.tooltip"))

        row {
            checkBox(message("aws.settings.codewhisperer.configurable.opt_out.title")).apply {
                connect.subscribe(
                    ToolkitConnectionManagerListener.TOPIC,
                    object : ToolkitConnectionManagerListener {
                        override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                            enabled(isCodeWhispererEnabled(project) && !isSso)
                        }
                    }
                )

                enabled(invoke && !isSso)

                if (isSso) {
                    bindSelected({ false }, {})
                } else {
                    bindSelected(codeWhispererSettings::isMetricOptIn, codeWhispererSettings::toggleMetricOptIn)
                }
            }
        }.rowComment(message("aws.settings.codewhisperer.configurable.opt_out.tooltip"))

        row {
            comment(message("aws.settings.codewhisperer.configurable.iam_identity_center.warning"))
        }.visible(isSso)
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.settings

import com.intellij.openapi.components.BaseState
import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.components.service
import com.intellij.util.xmlb.annotations.Property

@State(name = "codewhispererSettings", storages = [Storage("aws.xml")])
class CodeWhispererSettings : PersistentStateComponent<CodeWhispererConfiguration> {
    private val state = CodeWhispererConfiguration()

    fun toggleIncludeCodeWithReference(value: Boolean) {
        state.value[CodeWhispererConfigurationType.IsIncludeCodeWithReference] = value
    }

    fun isIncludeCodeWithReference() = state.value.getOrDefault(
        CodeWhispererConfigurationType.IsIncludeCodeWithReference,
        false
    )

    fun toggleImportAdder(value: Boolean) {
        state.value[CodeWhispererConfigurationType.IsImportAdderEnabled] = value
    }

    fun isImportAdderEnabled() = state.value.getOrDefault(
        CodeWhispererConfigurationType.IsImportAdderEnabled,
        true
    )

    fun toggleMetricOptIn(value: Boolean) {
        state.value[CodeWhispererConfigurationType.OptInSendingMetric] = value
    }

    fun isMetricOptIn() = state.value.getOrDefault(
        CodeWhispererConfigurationType.OptInSendingMetric,
        true
    )

    companion object {
        fun getInstance(): CodeWhispererSettings = service()
    }

    override fun getState(): CodeWhispererConfiguration = CodeWhispererConfiguration().apply { value.putAll(state.value) }

    override fun loadState(state: CodeWhispererConfiguration) {
        this.state.value.clear()
        this.state.value.putAll(state.value)
    }
}

class CodeWhispererConfiguration : BaseState() {
    @get:Property
    val value by map<CodeWhispererConfigurationType, Boolean>()
}

enum class CodeWhispererConfigurationType {
    IsIncludeCodeWithReference,
    OptInSendingMetric,
    IsImportAdderEnabled
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.startup

import com.intellij.codeInsight.lookup.Lookup
import com.intellij.codeInsight.lookup.LookupEvent
import com.intellij.codeInsight.lookup.LookupListener
import com.intellij.codeInsight.lookup.LookupManagerListener
import com.intellij.codeInsight.lookup.impl.LookupImpl
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutoTriggerService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType

object CodeWhispererIntelliSenseAutoTriggerListener : LookupManagerListener {
    override fun activeLookupChanged(oldLookup: Lookup?, newLookup: Lookup?) {
        if (oldLookup != null || newLookup == null) return

        newLookup.addLookupListener(object : LookupListener {
            override fun itemSelected(event: LookupEvent) {
                val editor = event.lookup.editor
                if (!(event.lookup as LookupImpl).isShown) {
                    cleanup()
                    return
                }

                // Classifier
                CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.IntelliSense())
                cleanup()
            }
            override fun lookupCanceled(event: LookupEvent) {
                cleanup()
            }

            private fun cleanup() {
                newLookup.removeLookupListener(this)
            }
        })
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.startup

import com.intellij.codeInsight.lookup.LookupManagerListener
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.invokeLater
import com.intellij.openapi.project.Project
import com.intellij.openapi.startup.StartupActivity
import software.aws.toolkits.jetbrains.core.explorer.refreshDevToolTree
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.importadder.CodeWhispererImportAdderListener
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager.Companion.CODEWHISPERER_USER_ACTION_PERFORMED
import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarManager
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorAccountless
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyWarnAccountless
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
import java.time.LocalDateTime
import java.util.Date
import java.util.Timer
import kotlin.concurrent.schedule

// TODO: add logics to check if we want to remove recommendation suspension date when user open the IDE
class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
    private var runOnce = false

    /**
     * Should be invoked when
     * (1) new users accept CodeWhisperer ToS (have to be triggered manually))
     * (2) existing users open the IDE (automatically triggered)
     */
    override fun runActivity(project: Project) {
        if (!ApplicationManager.getApplication().isUnitTestMode) {
            CodeWhispererStatusBarManager.getInstance(project).updateWidget()
        }
        if (!isCodeWhispererEnabled(project)) return
        if (runOnce) return

        // Reconnect CodeWhisperer on startup
        promptReAuth(project, isPluginStarting = true)

        // install intellsense autotrigger listener, this only need to be executed 1 time
        project.messageBus.connect().subscribe(LookupManagerListener.TOPIC, CodeWhispererIntelliSenseAutoTriggerListener)
        project.messageBus.connect().subscribe(CODEWHISPERER_USER_ACTION_PERFORMED, CodeWhispererImportAdderListener)

        // show notification to accountless users
        showAccountlessNotificationIfNeeded(project)
        runOnce = true
    }

    private fun showAccountlessNotificationIfNeeded(project: Project) {
        if (CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Accountless) {
            // simply show a notification when user login with Accountless, and it's still supported by CodeWhisperer
            if (!isExpired()) {
                // don't show warn notification if user selected Don't show again or if notification was shown less than a week ago
                if (!timeToShowAccessTokenWarn() || CodeWhispererExplorerActionManager.getInstance().getDoNotShowAgainWarn()) {
                    return
                }
                notifyWarnAccountless()
                CodeWhispererExplorerActionManager.getInstance().setAccountlessNotificationWarnTimestamp()

                // to handle the case when user open the IDE when Accountless not yet expired but expire soon e.g. 30min etc.
                Timer().schedule(CodeWhispererConstants.EXPIRE_DATE) { notifyErrorAndDisableAccountless(project) }
            } else {
                if (!timeToShowAccessTokenError() || CodeWhispererExplorerActionManager.getInstance().getDoNotShowAgainError()) {
                    return
                }
                CodeWhispererExplorerActionManager.getInstance().setAccountlessNotificationErrorTimestamp()
                notifyErrorAndDisableAccountless(project)
            }
        } else if (CodeWhispererExplorerActionManager.getInstance().getAccountlessNullified()) {
            if (!timeToShowAccessTokenError() || CodeWhispererExplorerActionManager.getInstance().getDoNotShowAgainError()) {
                return
            }
            CodeWhispererExplorerActionManager.getInstance().setAccountlessNotificationErrorTimestamp()
            notifyErrorAndDisableAccountless(project)
        }
    }

    private fun notifyErrorAndDisableAccountless(project: Project) {
        // show an error and deactivate CW when user login with Accountless, and it already expired
        notifyErrorAccountless()
        CodeWhispererExplorerActionManager.getInstance().nullifyAccountlessCredentialIfNeeded()
        invokeLater { project.refreshDevToolTree() }
    }

    private fun timeToShowAccessTokenWarn(): Boolean {
        val lastShown = CodeWhispererExplorerActionManager.getInstance().getAccountlessWarnNotificationTimestamp()
        return lastShown?.let {
            val parsedLastShown = LocalDateTime.parse(lastShown, CodeWhispererConstants.TIMESTAMP_FORMATTER)
            parsedLastShown.plusDays(7) <= LocalDateTime.now()
        } ?: true
    }

    private fun timeToShowAccessTokenError(): Boolean {
        val lastShown = CodeWhispererExplorerActionManager.getInstance().getAccountlessErrorNotificationTimestamp()
        return lastShown?.let {
            val parsedLastShown = LocalDateTime.parse(lastShown, CodeWhispererConstants.TIMESTAMP_FORMATTER)
            parsedLastShown.plusDays(7) <= LocalDateTime.now()
        } ?: true
    }
}

// TODO: do we have time zone issue with Date?
private fun isExpired() = CodeWhispererConstants.EXPIRE_DATE.before(Date())


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

package software.aws.toolkits.jetbrains.services.codewhisperer.startup

import com.intellij.analysis.problemsView.toolWindow.ProblemsView
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.openapi.wm.ToolWindow
import com.intellij.openapi.wm.ex.ToolWindowManagerListener
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetsManager
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererActivationChangedListener
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidgetFactory
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager

class CodeWhispererProjectStartupSettingsListener(private val project: Project) :
    CodeWhispererActivationChangedListener,
    ToolWindowManagerListener,
    ToolkitConnectionManagerListener {
    override fun activationChanged(value: Boolean) {
        project.service<StatusBarWidgetsManager>().updateWidget(CodeWhispererStatusBarWidgetFactory::class.java)
        CodeWhispererCodeReferenceManager.getInstance(project).toolWindow?.isAvailable = value
        if (value) {
            CodeWhispererSettings.getInstance().toggleIncludeCodeWithReference(true)
            CodeWhispererCodeScanManager.getInstance(project).addCodeScanUI()
        } else {
            CodeWhispererCodeScanManager.getInstance(project).removeCodeScanUI()
        }
    }

    override fun toolWindowShown(toolWindow: ToolWindow) {
        super.toolWindowShown(toolWindow)
        if (toolWindow.id != ProblemsView.ID) return
        if (!isCodeWhispererEnabled(project)) return
        CodeWhispererCodeScanManager.getInstance(project).addCodeScanUI()
    }

    override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
        // For now we have the assumption that any connection change will include CW
        // will need to change if we separate CW connections and CodeCatalyst connections
        runInEdt {
            CodeWhispererCodeReferenceManager.getInstance(project).toolWindow?.isAvailable = newConnection != null
        }
        if (newConnection != null) {
            CodeWhispererCodeScanManager.getInstance(project).addCodeScanUI()
        } else {
            CodeWhispererCodeScanManager.getInstance(project).removeCodeScanUI()
        }
    }
}


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

package software.aws.toolkits.jetbrains.services.codewhisperer.status

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.project.Project
import com.intellij.openapi.wm.StatusBarWidgetFactory
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetSettings
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetsManager
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.core.credentials.pinning.ConnectionPinningManagerListener
import software.aws.toolkits.jetbrains.core.credentials.pinning.FeatureWithPinnedConnection
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled

/**
 * Manager visibility of CodeWhisperer status bar widget, only display it when CodeWhisperer is connected
 */
@Service(Service.Level.PROJECT)
class CodeWhispererStatusBarManager(private val project: Project) : Disposable {
    private val widgetsManager = project.getService(StatusBarWidgetsManager::class.java)
    private val settings = ApplicationManager.getApplication().getService(StatusBarWidgetSettings::class.java)

    init {
        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
            ToolkitConnectionManagerListener.TOPIC,
            object : ToolkitConnectionManagerListener {
                override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                    updateWidget()
                }
            }
        )

        project.messageBus.connect().subscribe(
            ConnectionPinningManagerListener.TOPIC,
            object : ConnectionPinningManagerListener {
                override fun pinnedConnectionChanged(feature: FeatureWithPinnedConnection, newConnection: ToolkitConnection?) {
                    if (feature !is CodeWhispererConnection) return
                    updateWidget()
                }
            }
        )
    }

    fun updateWidget() {
        ExtensionPointName<StatusBarWidgetFactory>("com.intellij.statusBarWidgetFactory").extensionList.find {
            it.id == CodeWhispererStatusBarWidgetFactory.ID
        }?.let {
            settings.setEnabled(it, isCodeWhispererEnabled(project))
            widgetsManager.updateWidget(it)
        }
    }

    override fun dispose() {}

    companion object {
        fun getInstance(project: Project): CodeWhispererStatusBarManager = project.service()
    }
}