// 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.icons.AllIcons
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopupFactory
import com.intellij.openapi.ui.popup.ListPopup
import com.intellij.openapi.wm.StatusBar
import com.intellij.openapi.wm.StatusBarWidget
import com.intellij.openapi.wm.impl.status.EditorBasedWidget
import com.intellij.ui.AnimatedIcon
import com.intellij.util.Consumer
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererExpired
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStateChangeListener
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.reconnectCodeWhisperer
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent
import javax.swing.Icon

class CodeWhispererStatusBarWidget(project: Project) :
    EditorBasedWidget(project),
    StatusBarWidget.MultipleTextValuesPresentation {

    override fun install(statusBar: StatusBar) {
        super.install(statusBar)
        project.messageBus.connect(this).subscribe(
            CodeWhispererInvocationStatus.CODEWHISPERER_INVOCATION_STATE_CHANGED,
            object : CodeWhispererInvocationStateChangeListener {
                override fun invocationStateChanged(value: Boolean) {
                    statusBar.updateWidget(ID)
                }
            }
        )
        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
            BearerTokenProviderListener.TOPIC,
            object : BearerTokenProviderListener {
                override fun onChange(providerId: String) {
                    statusBar.updateWidget(ID)
                }
            }
        )
    }

    override fun ID(): String = ID

    override fun getPresentation(): StatusBarWidget.WidgetPresentation = this

    override fun getTooltipText(): String = message("codewhisperer.statusbar.tooltip")

    override fun getClickConsumer(): Consumer<MouseEvent>? = null

    override fun getPopupStep(): ListPopup? =
        if (isCodeWhispererExpired(project)) {
            JBPopupFactory.getInstance().createConfirmation(message("codewhisperer.statusbar.popup.title"), { reconnectCodeWhisperer(project) }, 0)
        } else {
            null
        }

    override fun getSelectedValue(): String = message("codewhisperer.statusbar.display_name")

    override fun getIcon(): Icon =
        if (isCodeWhispererExpired(project)) {
            AllIcons.General.BalloonWarning
        } else if (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
            AnimatedIcon.Default()
        } else {
            AllIcons.Actions.Commit
        }

    companion object {
        const val ID = "aws.codewhisperer.statusWidget"
    }
}


// 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.project.Project
import com.intellij.openapi.util.Disposer
import com.intellij.openapi.wm.StatusBar
import com.intellij.openapi.wm.StatusBarWidget
import com.intellij.openapi.wm.impl.status.widget.StatusBarEditorBasedWidgetFactory
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.resources.message

class CodeWhispererStatusBarWidgetFactory : StatusBarEditorBasedWidgetFactory() {
    override fun getId(): String = ID

    override fun getDisplayName(): String = message("codewhisperer.statusbar.display_name")

    override fun isAvailable(project: Project): Boolean =
        isCodeWhispererEnabled(project)

    override fun createWidget(project: Project): StatusBarWidget = CodeWhispererStatusBarWidget(project)

    override fun disposeWidget(widget: StatusBarWidget) {
        Disposer.dispose(widget)
    }

    override fun canBeEnabledOn(statusBar: StatusBar) = true

    companion object {
        const val ID = "aws.codewhisperer"
    }
}


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

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

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.editor.Document
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.editor.event.DocumentEvent
import com.intellij.openapi.util.Key
import com.intellij.refactoring.suggested.range
import com.intellij.util.Alarm
import com.intellij.util.AlarmFactory
import info.debatty.java.stringsimilarity.Levenshtein
import org.jetbrains.annotations.TestOnly
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
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.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererCodeCompletionServiceListener
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_SECONDS_IN_MINUTE
import software.aws.toolkits.telemetry.CodewhispererTelemetry
import java.time.Duration
import java.time.Instant
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import kotlin.math.roundToInt

abstract class CodeWhispererCodeCoverageTracker(
    private val timeWindowInSec: Long,
    private val language: CodeWhispererProgrammingLanguage,
    private val rangeMarkers: MutableList<RangeMarker>,
    private val fileToTokens: MutableMap<Document, CodeCoverageTokens>,
    private val myServiceInvocationCount: AtomicInteger
) : Disposable {
    val percentage: Int?
        get() = if (totalTokensSize != 0) calculatePercentage(acceptedTokensSize, totalTokensSize) else null
    val acceptedTokensSize: Int
        get() = fileToTokens.map {
            it.value.acceptedTokens.get()
        }.fold(0) { acc, next ->
            acc + next
        }
    val totalTokensSize: Int
        get() = fileToTokens.map {
            it.value.totalTokens.get()
        }.fold(0) { acc, next ->
            acc + next
        }
    val acceptedRecommendationsCount: Int
        get() = rangeMarkers.size
    val serviceInvocationCount: Int
        get() = myServiceInvocationCount.get()
    private val isActive: AtomicBoolean = AtomicBoolean(false)
    private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)
    private val isShuttingDown = AtomicBoolean(false)
    private var startTime: Instant = Instant.now()

    @Synchronized
    fun activateTrackerIfNotActive() {
        // tracker will only be activated if and only if IsTelemetryEnabled = true && isActive = false
        if (!isTelemetryEnabled() || isTrackerActive()) return

        val conn = ApplicationManager.getApplication().messageBus.connect()
        conn.subscribe(
            CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED,
            object : CodeWhispererUserActionListener {
                override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {
                    if (states.requestContext.fileContextInfo.programmingLanguage != language) return
                    rangeMarkers.add(rangeMarker)
                    val originalRecommendation = extractRangeMarkerString(rangeMarker)
                    originalRecommendation?.let {
                        rangeMarker.putUserData(KEY_REMAINING_RECOMMENDATION, it)
                    }
                }
            }
        )

        conn.subscribe(
            CodeWhispererService.CODEWHISPERER_CODE_COMPLETION_PERFORMED,
            object : CodeWhispererCodeCompletionServiceListener {
                override fun onSuccess(fileContextInfo: FileContextInfo) {
                    if (language == fileContextInfo.programmingLanguage) {
                        myServiceInvocationCount.getAndIncrement()
                    }
                }
            }
        )
        startTime = Instant.now()
        isActive.set(true)
        scheduleCodeWhispererCodeCoverageTracker()
    }

    fun isTrackerActive() = isActive.get()

    internal fun documentChanged(event: DocumentEvent) {
        // When open a file for the first time, IDE will also emit DocumentEvent for loading with `isWholeTextReplaced = true`
        // Added this condition to filter out those events
        if (event.isWholeTextReplaced) {
            LOG.debug { "event with isWholeTextReplaced flag: $event" }
            if (event.oldTimeStamp == 0L) return
        }
        // This case capture IDE reformatting the document, which will be blank string
        if (isDocumentEventFromReformatting(event)) return
        incrementTotalTokens(event.document, event.newLength - event.oldLength)
    }

    internal fun extractRangeMarkerString(rangeMarker: RangeMarker): String? = runReadAction {
        rangeMarker.range?.let { myRange -> rangeMarker.document.getText(myRange) }
    }

    // With edit distance, complicate usermodification can be considered as simple edit(add, delete, replace),
    // and thus the unmodified part of recommendation length can be deducted/approximated
    // ex. (modified > original): originalRecom: foo -> modifiedRecom: fobarbarbaro, distance = 9, delta = 12 - 9 = 3
    // ex. (modified == original): originalRecom: helloworld -> modifiedRecom: HelloWorld, distance = 2, delta = 10 - 2 = 8
    // ex. (modified < original): originalRecom: CodeWhisperer -> modifiedRecom: CODE, distance = 12, delta = 13 - 12 = 1
    internal fun getAcceptedTokensDelta(originalRecommendation: String, modifiedRecommendation: String): Int {
        val editDistance = getEditDistance(modifiedRecommendation, originalRecommendation).toInt()
        return maxOf(originalRecommendation.length, modifiedRecommendation.length) - editDistance
    }

    protected open fun getEditDistance(modifiedString: String, originalString: String): Double =
        levenshteinChecker.distance(modifiedString, originalString)

    private fun flush() {
        try {
            if (isTelemetryEnabled()) emitCodeWhispererCodeContribution()
        } finally {
            reset()
            scheduleCodeWhispererCodeCoverageTracker()
        }
    }

    private fun scheduleCodeWhispererCodeCoverageTracker() {
        if (!alarm.isDisposed && !isShuttingDown.get()) {
            alarm.addRequest({ flush() }, Duration.ofSeconds(timeWindowInSec).toMillis())
        }
    }

    private fun incrementAcceptedTokens(document: Document, delta: Int) {
        var tokens = fileToTokens[document]
        if (tokens == null) {
            tokens = CodeCoverageTokens()
            fileToTokens[document] = tokens
        }
        tokens.acceptedTokens.addAndGet(delta)
    }

    private fun incrementTotalTokens(document: Document, delta: Int) {
        var tokens = fileToTokens[document]
        if (tokens == null) {
            tokens = CodeCoverageTokens()
            fileToTokens[document] = tokens
        }
        tokens.apply {
            totalTokens.addAndGet(delta)
            if (totalTokens.get() < 0) totalTokens.set(0)
        }
    }

    private fun isDocumentEventFromReformatting(event: DocumentEvent): Boolean =
        (event.newFragment.toString().isBlank() && event.oldFragment.toString().isBlank()) && (event.oldLength == 0 || event.newLength == 0)

    private fun reset() {
        startTime = Instant.now()
        rangeMarkers.clear()
        fileToTokens.clear()
        myServiceInvocationCount.set(0)
    }

    internal fun emitCodeWhispererCodeContribution() {
        rangeMarkers.forEach { rangeMarker ->
            if (!rangeMarker.isValid) return@forEach
            // if users add more code upon the recommendation generated from CodeWhisperer, we consider those added part as userToken but not CwsprTokens
            val originalRecommendation = rangeMarker.getUserData(KEY_REMAINING_RECOMMENDATION)
            val modifiedRecommendation = extractRangeMarkerString(rangeMarker)
            if (originalRecommendation == null || modifiedRecommendation == null) {
                LOG.debug {
                    "failed to get accepted recommendation. " +
                        "OriginalRecommendation is null: ${originalRecommendation == null}; " +
                        "ModifiedRecommendation is null: ${modifiedRecommendation == null}"
                }
                return@forEach
            }
            val delta = getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
            runReadAction {
                incrementAcceptedTokens(rangeMarker.document, delta)
            }
        }

        // percentage == null means totalTokens == 0 and users are not editing the document, thus we shouldn't emit telemetry for this
        percentage?.let { percentage ->
            CodewhispererTelemetry.codePercentage(
                project = null,
                acceptedTokensSize,
                language.toTelemetryType(),
                percentage,
                totalTokensSize,
                successCount = myServiceInvocationCount.get(),
                codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
            )
        }
    }

    @TestOnly
    fun forceTrackerFlush() {
        alarm.drainRequestsInTest()
    }

    @TestOnly
    fun activeRequestCount() = alarm.activeRequestCount

    override fun dispose() {
        if (isShuttingDown.getAndSet(true)) {
            return
        }
        flush()
    }

    companion object {
        @JvmStatic
        protected val levenshteinChecker = Levenshtein()
        private const val REMAINING_RECOMMENDATION = "remainingRecommendation"
        private val KEY_REMAINING_RECOMMENDATION = Key<String>(REMAINING_RECOMMENDATION)
        private val LOG = getLogger<CodeWhispererCodeCoverageTracker>()
        private val instances: MutableMap<CodeWhispererProgrammingLanguage, CodeWhispererCodeCoverageTracker> = mutableMapOf()

        fun calculatePercentage(acceptedTokens: Int, totalTokens: Int): Int = ((acceptedTokens.toDouble() * 100) / totalTokens).roundToInt()
        fun getInstance(language: CodeWhispererProgrammingLanguage): CodeWhispererCodeCoverageTracker = when (val instance = instances[language]) {
            null -> {
                val newTracker = DefaultCodeWhispererCodeCoverageTracker(language)
                instances[language] = newTracker
                newTracker
            }
            else -> instance
        }

        @TestOnly
        fun getInstancesMap(): MutableMap<CodeWhispererProgrammingLanguage, CodeWhispererCodeCoverageTracker> {
            assert(ApplicationManager.getApplication().isUnitTestMode)
            return instances
        }
    }
}

class DefaultCodeWhispererCodeCoverageTracker(language: CodeWhispererProgrammingLanguage) : CodeWhispererCodeCoverageTracker(
    5 * TOTAL_SECONDS_IN_MINUTE,
    language,
    mutableListOf(),
    mutableMapOf(),
    AtomicInteger(0)
)

class CodeCoverageTokens(totalTokens: Int = 0, acceptedTokens: Int = 0) {
    val totalTokens: AtomicInteger
    val acceptedTokens: AtomicInteger
    init {
        this.totalTokens = AtomicInteger(totalTokens)
        this.acceptedTokens = AtomicInteger(acceptedTokens)
    }
}


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

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

import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import org.apache.commons.collections4.queue.CircularFifoQueue
import org.jetbrains.annotations.TestOnly
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanTelemetryEvent
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
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.service.CodeWhispererAutoTriggerService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getConnectionStartUrl
import software.aws.toolkits.jetbrains.settings.AwsSettings
import software.aws.toolkits.telemetry.CodewhispererLanguage
import software.aws.toolkits.telemetry.CodewhispererPreviousSuggestionState
import software.aws.toolkits.telemetry.CodewhispererSuggestionState
import software.aws.toolkits.telemetry.CodewhispererTelemetry
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import java.time.Duration
import java.time.Instant
import java.util.Queue

class CodeWhispererTelemetryService {
    // store previous 5 userTrigger decisions
    private val previousUserTriggerDecisions = CircularFifoQueue<CodewhispererPreviousSuggestionState>(5)

    private var previousUserTriggerDecisionTimestamp: Instant? = null

    private val codewhispererTimeSinceLastUserDecision: Double?
        get() {
            return previousUserTriggerDecisionTimestamp?.let {
                Duration.between(it, Instant.now()).toMillis().toDouble()
            }
        }

    val previousUserTriggerDecision: CodewhispererPreviousSuggestionState?
        get() = if (previousUserTriggerDecisions.isNotEmpty()) previousUserTriggerDecisions.last() else null

    companion object {
        fun getInstance(): CodeWhispererTelemetryService = service()
        val LOG = getLogger<CodeWhispererTelemetryService>()
    }

    fun sendFailedServiceInvocationEvent(project: Project, exceptionType: String?) {
        CodewhispererTelemetry.serviceInvocation(
            project = project,
            codewhispererCursorOffset = 0,
            codewhispererLanguage = CodewhispererLanguage.Unknown,
            codewhispererLastSuggestionIndex = -1,
            codewhispererLineNumber = 0,
            codewhispererTriggerType = CodewhispererTriggerType.Unknown,
            duration = 0.0,
            reason = exceptionType,
            success = false,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun sendServiceInvocationEvent(
        requestId: String,
        requestContext: RequestContext,
        responseContext: ResponseContext,
        lastRecommendationIndex: Int,
        invocationSuccess: Boolean,
        latency: Double,
        exceptionType: String?
    ) {
        val (project, _, triggerTypeInfo, caretPosition) = requestContext
        val (triggerType, automatedTriggerType) = triggerTypeInfo
        val (offset, line) = caretPosition

        // since python now only supports UTG but not cross file context
        val supContext = if (requestContext.fileContextInfo.programmingLanguage.isUTGSupported() &&
            requestContext.supplementalContext?.isUtg == true
        ) {
            requestContext.supplementalContext
        } else if (requestContext.fileContextInfo.programmingLanguage.isSupplementalContextSupported() &&
            requestContext.supplementalContext?.isUtg == false
        ) {
            requestContext.supplementalContext
        } else {
            null
        }

        val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
        val startUrl = getConnectionStartUrl(requestContext.connection)
        CodewhispererTelemetry.serviceInvocation(
            project = project,
            codewhispererAutomatedTriggerType = automatedTriggerType.telemetryType,
            codewhispererCompletionType = responseContext.completionType,
            codewhispererCursorOffset = offset,
            codewhispererLanguage = codewhispererLanguage,
            codewhispererLastSuggestionIndex = lastRecommendationIndex,
            codewhispererLineNumber = line,
            codewhispererRequestId = requestId,
            codewhispererSessionId = responseContext.sessionId,
            codewhispererTriggerType = triggerType,
            duration = latency,
            reason = exceptionType,
            success = invocationSuccess,
            credentialStartUrl = startUrl,
            codewhispererImportRecommendationEnabled = CodeWhispererSettings.getInstance().isImportAdderEnabled(),
            codewhispererSupplementalContextTimeout = supContext?.isProcessTimeout,
            codewhispererSupplementalContextIsUtg = supContext?.isUtg,
            codewhispererSupplementalContextLatency = supContext?.latency?.toDouble(),
            codewhispererSupplementalContextLength = supContext?.contentLength,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun sendUserDecisionEvent(
        requestId: String,
        requestContext: RequestContext,
        responseContext: ResponseContext,
        detail: Completion,
        index: Int,
        suggestionState: CodewhispererSuggestionState,
        numOfRecommendations: Int
    ) {
        val (project, _, triggerTypeInfo) = requestContext
        val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
        val supplementalContext = requestContext.supplementalContext

        LOG.debug {
            "Recording user decisions of recommendation. " +
                "Index: $index, " +
                "State: $suggestionState, " +
                "Request ID: $requestId, " +
                "Recommendation: ${detail.content()}"
        }
        val startUrl = getConnectionStartUrl(requestContext.connection)
        val importEnabled = CodeWhispererSettings.getInstance().isImportAdderEnabled()
        CodewhispererTelemetry.userDecision(
            project = project,
            codewhispererCompletionType = responseContext.completionType,
            codewhispererLanguage = codewhispererLanguage,
            codewhispererPaginationProgress = numOfRecommendations,
            codewhispererRequestId = requestId,
            codewhispererSessionId = responseContext.sessionId,
            codewhispererSuggestionIndex = index,
            codewhispererSuggestionReferenceCount = detail.references().size,
            codewhispererSuggestionReferences = jacksonObjectMapper().writeValueAsString(detail.references().map { it.licenseName() }.toSet().toList()),
            codewhispererSuggestionImportCount = if (importEnabled) detail.mostRelevantMissingImports().size else null,
            codewhispererSuggestionState = suggestionState,
            codewhispererTriggerType = triggerTypeInfo.triggerType,
            credentialStartUrl = startUrl,
            codewhispererSupplementalContextIsUtg = supplementalContext?.isUtg,
            codewhispererSupplementalContextLength = supplementalContext?.contentLength,
            codewhispererSupplementalContextTimeout = supplementalContext?.isProcessTimeout,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun sendUserTriggerDecisionEvent(
        requestContext: RequestContext,
        responseContext: ResponseContext,
        recommendationContext: RecommendationContext,
        suggestionState: CodewhispererSuggestionState,
        popupShownTime: Duration?
    ) {
        val totalImportCount = recommendationContext.details.fold(0) { grandTotal, detail ->
            grandTotal + detail.recommendation.mostRelevantMissingImports().size
        }

        val automatedTriggerType = requestContext.triggerTypeInfo.automatedTriggerType
        val triggerChar = if (automatedTriggerType is CodeWhispererAutomatedTriggerType.SpecialChar) {
            automatedTriggerType.specialChar.toString()
        } else {
            null
        }

        val language = requestContext.fileContextInfo.programmingLanguage

        val shouldIncludeClassifier = language.isAllClassifier() ||
            (language.isClassifierSupported() && CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.Classifier)

        val classifierResult = if (shouldIncludeClassifier) {
            requestContext.triggerTypeInfo.automatedTriggerType.calculationResult
        } else {
            null
        }

        val classifierThreshold = if (shouldIncludeClassifier) {
            CodeWhispererAutoTriggerService.getThreshold()
        } else {
            null
        }

        val supplementalContext = requestContext.supplementalContext

        CodewhispererTelemetry.userTriggerDecision(
            project = requestContext.project,
            codewhispererSessionId = responseContext.sessionId,
            codewhispererFirstRequestId = requestContext.latencyContext.firstRequestId,
            credentialStartUrl = getConnectionStartUrl(requestContext.connection),
            codewhispererIsPartialAcceptance = null,
            codewhispererPartialAcceptanceCount = null,
            codewhispererCharactersAccepted = null,
            codewhispererCharactersRecommended = null,
            codewhispererCompletionType = responseContext.completionType,
            codewhispererLanguage = language.toTelemetryType(),
            codewhispererTriggerType = requestContext.triggerTypeInfo.triggerType,
            codewhispererAutomatedTriggerType = automatedTriggerType.telemetryType,
            codewhispererLineNumber = requestContext.caretPosition.line,
            codewhispererCursorOffset = requestContext.caretPosition.offset,
            codewhispererSuggestionCount = recommendationContext.details.size,
            codewhispererSuggestionImportCount = totalImportCount,
            codewhispererTotalShownTime = popupShownTime?.toMillis()?.toDouble(),
            codewhispererTriggerCharacter = triggerChar,
            codewhispererTypeaheadLength = recommendationContext.userInputSinceInvocation.length,
            codewhispererTimeSinceLastDocumentChange = CodeWhispererInvocationStatus.getInstance().getTimeSinceDocumentChanged(),
            codewhispererTimeSinceLastUserDecision = codewhispererTimeSinceLastUserDecision,
            codewhispererTimeToFirstRecommendation = requestContext.latencyContext.paginationFirstCompletionTime,
            codewhispererPreviousSuggestionState = previousUserTriggerDecision,
            codewhispererSuggestionState = suggestionState,
            codewhispererClassifierResult = classifierResult,
            codewhispererClassifierThreshold = classifierThreshold,
            codewhispererSupplementalContextIsUtg = supplementalContext?.isUtg,
            codewhispererSupplementalContextLength = supplementalContext?.contentLength,
            codewhispererSupplementalContextTimeout = supplementalContext?.isProcessTimeout,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun sendSecurityScanEvent(codeScanEvent: CodeScanTelemetryEvent, project: Project? = null) {
        val (payloadContext, serviceInvocationContext, codeScanJobId, totalIssues, reason) = codeScanEvent.codeScanResponseContext
        val startUrl = getConnectionStartUrl(codeScanEvent.connection)

        LOG.debug {
            "Recording code security scan event. \n" +
                "Total number of security scan issues found: $totalIssues, \n" +
                "Language: ${payloadContext.language}, \n" +
                "Uncompressed source payload size in bytes: ${payloadContext.srcPayloadSize}, \n" +
                "Uncompressed build payload size in bytes: ${payloadContext.buildPayloadSize}, \n" +
                "Compressed source zip file size in bytes: ${payloadContext.srcZipFileSize}, \n" +
                "Total project size in bytes: ${codeScanEvent.totalProjectSizeInBytes}, \n" +
                "Total duration of the security scan job in milliseconds: ${codeScanEvent.duration}, \n" +
                "Context truncation duration in milliseconds: ${payloadContext.totalTimeInMilliseconds}, \n" +
                "Artifacts upload duration in milliseconds: ${serviceInvocationContext.artifactsUploadDuration}, \n" +
                "Service invocation duration in milliseconds: ${serviceInvocationContext.serviceInvocationDuration}, \n" +
                "Total number of lines scanned: ${payloadContext.totalLines}, \n" +
                "Reason: $reason \n"
        }
        CodewhispererTelemetry.securityScan(
            project = project,
            codewhispererCodeScanLines = payloadContext.totalLines.toInt(),
            codewhispererCodeScanJobId = codeScanJobId,
            codewhispererCodeScanProjectBytes = codeScanEvent.totalProjectSizeInBytes,
            codewhispererCodeScanSrcPayloadBytes = payloadContext.srcPayloadSize.toInt(),
            codewhispererCodeScanBuildPayloadBytes = payloadContext.buildPayloadSize?.toInt(),
            codewhispererCodeScanSrcZipFileBytes = payloadContext.srcZipFileSize.toInt(),
            codewhispererCodeScanTotalIssues = totalIssues,
            codewhispererLanguage = payloadContext.language,
            duration = codeScanEvent.duration,
            contextTruncationDuration = payloadContext.totalTimeInMilliseconds.toInt(),
            artifactsUploadDuration = serviceInvocationContext.artifactsUploadDuration.toInt(),
            codeScanServiceInvocationsDuration = serviceInvocationContext.serviceInvocationDuration.toInt(),
            reason = reason,
            result = codeScanEvent.result,
            credentialStartUrl = startUrl
        )
    }

    fun enqueueAcceptedSuggestionEntry(
        requestId: String,
        requestContext: RequestContext,
        responseContext: ResponseContext,
        time: Instant,
        vFile: VirtualFile?,
        range: RangeMarker,
        suggestion: String,
        selectedIndex: Int
    ) {
        val (sessionId, completionType) = responseContext
        val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
        CodeWhispererUserModificationTracker.getInstance(requestContext.project).enqueue(
            AcceptedSuggestionEntry(
                time, vFile, range, suggestion, sessionId, requestId, selectedIndex,
                requestContext.triggerTypeInfo.triggerType, completionType,
                codewhispererLanguage, null, null,
                requestContext.connection
            )
        )
    }

    fun sendUserDecisionEventForAll(
        requestContext: RequestContext,
        responseContext: ResponseContext,
        recommendationContext: RecommendationContext,
        sessionContext: SessionContext,
        hasUserAccepted: Boolean,
        popupShownTime: Duration? = null
    ) {
        val detailContexts = recommendationContext.details
        val decisions = mutableListOf<CodewhispererSuggestionState>()

        detailContexts.forEachIndexed { index, detailContext ->
            val (requestId, detail, _, isDiscarded) = detailContext
            val suggestionState = recordSuggestionState(
                index,
                sessionContext.selectedIndex,
                sessionContext.seen.contains(index),
                hasUserAccepted,
                isDiscarded,
                detail.content().isEmpty()
            )
            sendUserDecisionEvent(requestId, requestContext, responseContext, detail, index, suggestionState, detailContexts.size)

            decisions.add(suggestionState)
        }

        with(aggregateUserDecision(decisions)) {
            // the order of the following matters
            // step 1, send out current decision
            previousUserTriggerDecisionTimestamp = Instant.now()
            sendUserTriggerDecisionEvent(
                requestContext,
                responseContext,
                recommendationContext,
                CodewhispererSuggestionState.from(this.toString()),
                popupShownTime,
            )

            // step 2, put current decision into queue for later reference
            previousUserTriggerDecisions.add(this)
            // we need this as well because AutotriggerService will reset the queue periodically
            CodeWhispererAutoTriggerService.getInstance().addPreviousDecision(this)
        }
    }

    /**
     * Aggregate recommendation level user decision to trigger level user decision based on the following rule
     * - Accept if there is an Accept
     * - Reject if there is a Reject
     * - Empty if all decisions are Empty
     * - Discard otherwise
     */
    fun aggregateUserDecision(decisions: List<CodewhispererSuggestionState>): CodewhispererPreviousSuggestionState {
        var isEmpty = true

        for (decision in decisions) {
            if (decision == CodewhispererSuggestionState.Accept) {
                return CodewhispererPreviousSuggestionState.Accept
            } else if (decision == CodewhispererSuggestionState.Reject) {
                return CodewhispererPreviousSuggestionState.Reject
            } else if (decision != CodewhispererSuggestionState.Empty) {
                isEmpty = false
            }
        }

        return if (isEmpty) {
            CodewhispererPreviousSuggestionState.Empty
        } else CodewhispererPreviousSuggestionState.Discard
    }

    fun sendPerceivedLatencyEvent(
        requestId: String,
        requestContext: RequestContext,
        responseContext: ResponseContext,
        latency: Double,
    ) {
        val (project, _, triggerTypeInfo) = requestContext
        val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
        val startUrl = getConnectionStartUrl(requestContext.connection)
        CodewhispererTelemetry.perceivedLatency(
            project = project,
            codewhispererCompletionType = responseContext.completionType,
            codewhispererLanguage = codewhispererLanguage,
            codewhispererRequestId = requestId,
            codewhispererSessionId = responseContext.sessionId,
            codewhispererTriggerType = triggerTypeInfo.triggerType,
            duration = latency,
            passive = true,
            credentialStartUrl = startUrl,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun sendClientComponentLatencyEvent(states: InvocationContext) {
        val requestContext = states.requestContext
        val responseContext = states.responseContext
        val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
        val startUrl = getConnectionStartUrl(requestContext.connection)
        CodewhispererTelemetry.clientComponentLatency(
            project = requestContext.project,
            codewhispererSessionId = responseContext.sessionId,
            codewhispererRequestId = requestContext.latencyContext.firstRequestId,
            codewhispererFirstCompletionLatency = requestContext.latencyContext.paginationFirstCompletionTime,
            codewhispererPreprocessingLatency = requestContext.latencyContext.getCodeWhispererPreprocessingLatency(),
            codewhispererEndToEndLatency = requestContext.latencyContext.getCodeWhispererEndToEndLatency(),
            codewhispererAllCompletionsLatency = requestContext.latencyContext.getCodeWhispererAllCompletionsLatency(),
            codewhispererPostprocessingLatency = requestContext.latencyContext.getCodeWhispererPostprocessingLatency(),
            codewhispererCredentialFetchingLatency = requestContext.latencyContext.getCodeWhispererCredentialFetchingLatency(),
            codewhispererTriggerType = requestContext.triggerTypeInfo.triggerType,
            codewhispererCompletionType = responseContext.completionType,
            codewhispererLanguage = codewhispererLanguage,
            credentialStartUrl = startUrl,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    fun recordSuggestionState(
        index: Int,
        selectedIndex: Int,
        hasSeen: Boolean,
        hasUserAccepted: Boolean,
        isDiscarded: Boolean,
        isEmpty: Boolean
    ): CodewhispererSuggestionState =
        if (isEmpty) {
            CodewhispererSuggestionState.Empty
        } else if (isDiscarded) {
            CodewhispererSuggestionState.Discard
        } else if (!hasSeen) {
            CodewhispererSuggestionState.Unseen
        } else if (hasUserAccepted) {
            if (selectedIndex == index) {
                CodewhispererSuggestionState.Accept
            } else {
                CodewhispererSuggestionState.Ignore
            }
        } else {
            CodewhispererSuggestionState.Reject
        }

    @TestOnly
    fun previousDecisions(): Queue<CodewhispererPreviousSuggestionState> {
        assert(ApplicationManager.getApplication().isUnitTestMode)
        return this.previousUserTriggerDecisions
    }
}

fun isTelemetryEnabled(): Boolean = AwsSettings.getInstance().isTelemetryEnabled


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

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

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.RangeMarker
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.util.Alarm
import com.intellij.util.AlarmFactory
import info.debatty.java.stringsimilarity.Levenshtein
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getConnectionStartUrl
import software.aws.toolkits.jetbrains.settings.AwsSettings
import software.aws.toolkits.telemetry.CodewhispererCompletionType
import software.aws.toolkits.telemetry.CodewhispererLanguage
import software.aws.toolkits.telemetry.CodewhispererRuntime
import software.aws.toolkits.telemetry.CodewhispererTelemetry
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import java.time.Duration
import java.time.Instant
import java.util.concurrent.LinkedBlockingDeque
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.math.min

data class AcceptedSuggestionEntry(
    val time: Instant,
    val vFile: VirtualFile?,
    val range: RangeMarker,
    val suggestion: String,
    val sessionId: String,
    val requestId: String,
    val index: Int,
    val triggerType: CodewhispererTriggerType,
    val completionType: CodewhispererCompletionType,
    val codewhispererLanguage: CodewhispererLanguage,
    val codewhispererRuntime: CodewhispererRuntime?,
    val codewhispererRuntimeSource: String?,
    val connection: ToolkitConnection?
)

class CodeWhispererUserModificationTracker(private val project: Project) : Disposable {
    private val acceptedSuggestions = LinkedBlockingDeque<AcceptedSuggestionEntry>(DEFAULT_MAX_QUEUE_SIZE)
    private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)

    private val isShuttingDown = AtomicBoolean(false)

    init {
        scheduleCodeWhispererTracker()
    }

    private fun scheduleCodeWhispererTracker() {
        if (!alarm.isDisposed && !isShuttingDown.get()) {
            alarm.addRequest({ flush() }, DEFAULT_CHECK_INTERVAL.toMillis())
        }
    }

    private fun isTelemetryEnabled(): Boolean = TELEMETRY_ENABLED and AwsSettings.getInstance().isTelemetryEnabled

    fun enqueue(event: AcceptedSuggestionEntry) {
        if (!isTelemetryEnabled()) {
            return
        }

        acceptedSuggestions.add(event)
        LOG.debug { "Enqueue Accepted Suggestion on line $event.lineNumber in $event.filePath" }
    }

    private fun flush() {
        try {
            if (!isTelemetryEnabled()) {
                acceptedSuggestions.clear()
                return
            }

            val copyList = LinkedBlockingDeque<AcceptedSuggestionEntry>()

            val currentTime = Instant.now()
            for (acceptedSuggestion in acceptedSuggestions) {
                if (Duration.between(acceptedSuggestion.time, currentTime).seconds > DEFAULT_MODIFICATION_INTERVAL_IN_SECONDS) {
                    LOG.debug { "Passed $DEFAULT_MODIFICATION_INTERVAL_IN_SECONDS for $acceptedSuggestion" }
                    emitTelemetryOnSuggestion(acceptedSuggestion)
                } else {
                    copyList.add(acceptedSuggestion)
                }
            }

            acceptedSuggestions.clear()
            acceptedSuggestions.addAll(copyList)
        } finally {
            scheduleCodeWhispererTracker()
        }
    }

    private fun emitTelemetryOnSuggestion(acceptedSuggestion: AcceptedSuggestionEntry) {
        val file = acceptedSuggestion.vFile
        if (file == null || (!file.isValid)) {
            sendModificationTelemetry(acceptedSuggestion, 1.0)
        } else {
            try {
                /**
                 * this try-catch is to check if the offsets are valid since the method does not return null
                 */
                val document = runReadAction {
                    FileDocumentManager.getInstance().getDocument(file)
                }
                val currentString = document?.getText(
                    TextRange(acceptedSuggestion.range.startOffset, acceptedSuggestion.range.endOffset)
                )
                sendModificationTelemetry(acceptedSuggestion, checkDiff(currentString?.trim(), acceptedSuggestion.suggestion.trim()))
            } catch (e: Exception) {
                sendModificationTelemetry(acceptedSuggestion, 1.0)
            }
        }
    }

    /**
     * Use Levenshtein distance to check how
     * Levenshtein distance was preferred over Jaroâ€“Winkler distance for simplicity
     */
    private fun checkDiff(currString: String?, acceptedString: String?): Double {
        if (currString == null || acceptedString == null || acceptedString.isEmpty() || currString.isEmpty()) {
            return 1.0
        }

        val diff = checker.distance(currString, acceptedString)
        val percentage = diff / acceptedString.length

        return min(1.0, percentage)
    }

    private fun sendModificationTelemetry(suggestion: AcceptedSuggestionEntry, percentage: Double) {
        LOG.debug { "Sending user modification telemetry. Request Id: ${suggestion.requestId}" }
        val startUrl = getConnectionStartUrl(suggestion.connection)
        CodewhispererTelemetry.userModification(
            project = project,
            codewhispererCompletionType = suggestion.completionType,
            codewhispererLanguage = suggestion.codewhispererLanguage,
            codewhispererModificationPercentage = percentage,
            codewhispererRequestId = suggestion.requestId,
            codewhispererRuntime = suggestion.codewhispererRuntime,
            codewhispererRuntimeSource = suggestion.codewhispererRuntimeSource,
            codewhispererSessionId = suggestion.sessionId,
            codewhispererSuggestionIndex = suggestion.index,
            codewhispererTriggerType = suggestion.triggerType,
            credentialStartUrl = startUrl,
            codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
        )
    }

    companion object {
        private val DEFAULT_CHECK_INTERVAL = Duration.ofMinutes(1)
        private const val DEFAULT_MAX_QUEUE_SIZE = 10000
        private const val DEFAULT_MODIFICATION_INTERVAL_IN_SECONDS = 300 // 5 minutes
        private const val TELEMETRY_KEY = "aws.toolkits.enableTelemetry"
        private val checker = Levenshtein()
        private val TELEMETRY_ENABLED = System.getProperty(TELEMETRY_KEY)?.toBoolean() ?: true

        private val LOG = getLogger<CodeWhispererUserModificationTracker>()

        fun getInstance(project: Project) = project.service<CodeWhispererUserModificationTracker>()
    }

    override fun dispose() {
        if (isShuttingDown.getAndSet(true)) {
            return
        }

        flush()
    }
}


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

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

import com.intellij.openapi.editor.RangeMarker
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener

class CodeWhispererCodeReferenceActionListener : CodeWhispererUserActionListener {
    override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {
        val (project, editor) = states.requestContext
        val manager = CodeWhispererCodeReferenceManager.getInstance(project)
        manager.insertCodeReference(states, sessionContext.selectedIndex)
        manager.addListeners(editor)
    }
}


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

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

import com.intellij.ide.BrowserUtil
import com.intellij.openapi.editor.colors.EditorColorsManager
import com.intellij.openapi.options.ShowSettingsUtil
import com.intellij.openapi.project.Project
import com.intellij.ui.components.ActionLink
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
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.layout.CodeWhispererLayoutConfig.addHorizontalGlue
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.addVerticalGlue
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.horizontalPanelConstraints
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.inlineLabelConstraints
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererLicenseInfoManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.TOOLWINDOW_BACKGROUND
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.TOOLWINDOW_CODE
import software.aws.toolkits.resources.message
import java.awt.Font
import java.awt.GridBagLayout
import java.time.LocalTime
import java.time.format.DateTimeFormatter
import javax.swing.BorderFactory
import javax.swing.JComponent
import javax.swing.JLabel
import javax.swing.JPanel

class CodeWhispererCodeReferenceComponents(private val project: Project) {
    private val settingsLabelPrefixText = JLabel().apply {
        text = message("codewhisperer.toolwindow.settings.prefix")
    }.asCodeReferencePanelFont()

    private val settingsLabelLink = ActionLink().apply {
        text = message("codewhisperer.toolwindow.settings")
        addActionListener {
            ShowSettingsUtil.getInstance().showSettingsDialog(project, CodeWhispererConfigurable::class.java)
        }
    }.asCodeReferencePanelFont()

    private val settingsPanel = JPanel(GridBagLayout()).apply {
        background = TOOLWINDOW_BACKGROUND
        border = BorderFactory.createEmptyBorder(0, 0, 17, 0)
        add(settingsLabelPrefixText, inlineLabelConstraints)
        add(settingsLabelLink, inlineLabelConstraints)
        add(JLabel("."), inlineLabelConstraints)
        addHorizontalGlue()
    }
    val contentPanel = JPanel(GridBagLayout()).apply {
        background = TOOLWINDOW_BACKGROUND
        border = BorderFactory.createEmptyBorder(7, 14, 0, 0)
        add(settingsPanel, horizontalPanelConstraints)
        addVerticalGlue()
    }

    private val codeReferenceTimeFormatter = DateTimeFormatter.ofPattern("HH:mm:ss")
    private val acceptRecommendationPrefixText
        get() = JLabel().apply {
            text = message("codewhisperer.toolwindow.entry.prefix", LocalTime.now().format(codeReferenceTimeFormatter))
        }.asCodeReferencePanelFont()

    init {
        repaint(project)

        // set the reference panel text different for SSO users vs AWS Builder ID / Accless users
        project.messageBus.connect().subscribe(
            ToolkitConnectionManagerListener.TOPIC,
            object : ToolkitConnectionManagerListener {
                override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                    repaint(project)
                }
            }
        )
    }

    // TODO: figure out how to have a different view for SSO user in a cleaner way, maybe have 2 sets of components stored in [ReferenceManager]?
    private fun repaint(project: Project) {
        val loginType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project)
        settingsLabelPrefixText as JLabel
        settingsLabelLink as ActionLink
        if (loginType == CodeWhispererLoginType.SSO) {
            settingsLabelPrefixText.text = message("codewhisperer.toolwindow.settings.prefix_sso")
            settingsLabelLink.isVisible = false
        } else {
            settingsLabelPrefixText.text = message("codewhisperer.toolwindow.settings.prefix")
            settingsLabelLink.isVisible = true
        }
    }

    private fun licenseNameLink(licenseName: String) = ActionLink(licenseName) {
        BrowserUtil.browse(CodeWhispererLicenseInfoManager.getInstance().getLicenseLink(licenseName))
    }.asCodeReferencePanelFont()

    private fun repoNameLink(repo: String, url: String) = ActionLink(repo) {
        BrowserUtil.browse(url)
    }.asCodeReferencePanelFont()

    private fun acceptRecommendationSuffixText(path: String?, line: String) = JLabel().apply {
        val choice = if (path != null) 1 else 0
        text = message("codewhisperer.toolwindow.entry.suffix", path ?: "", choice, line)
    }.asCodeReferencePanelFont()

    fun codeReferenceRecordPanel(ref: Reference, relativePath: String?, lineNums: String) = JPanel(GridBagLayout()).apply {
        background = EditorColorsManager.getInstance().globalScheme.defaultBackground
        border = BorderFactory.createEmptyBorder(5, 0, 0, 0)
        add(acceptRecommendationPrefixText, inlineLabelConstraints)

        // if url to source package/repo is missing, the UX remains the same as we have for now
        // if url to source package/repo is present, the url pointing to the source will be present and remove the hyperlink to SPDX
        if (ref.url().isNullOrEmpty()) {
            add(
                licenseNameLink(ref.licenseName()).apply {
                    font = font.deriveFont(Font.ITALIC + Font.BOLD)
                },
                inlineLabelConstraints
            )
            add(JLabel(" from ").asCodeReferencePanelFont(), inlineLabelConstraints)
            add(JLabel(ref.repository()), inlineLabelConstraints)
        } else {
            add(
                JLabel(ref.licenseName()).apply {
                    font = font.deriveFont(Font.ITALIC + Font.BOLD)
                },
                inlineLabelConstraints
            )
            add(JLabel(" from ").asCodeReferencePanelFont(), inlineLabelConstraints)
            add(repoNameLink(ref.repository(), ref.url()), inlineLabelConstraints)
        }

        add(acceptRecommendationSuffixText(relativePath, lineNums), inlineLabelConstraints)
        addHorizontalGlue()
    }

    fun codeContentLine(line: String) = JLabel(line).apply {
        foreground = TOOLWINDOW_CODE
    }.asCodeReferencePanelFont()

    fun codeContentPanel(line: String) = JPanel(GridBagLayout()).apply {
        background = EditorColorsManager.getInstance().globalScheme.defaultBackground
        if (line == "") {
            add(codeContentLine(" "), inlineLabelConstraints)
        } else {
            add(codeContentLine(line), inlineLabelConstraints)
        }
        addHorizontalGlue()
    }

    private fun JComponent.asCodeReferencePanelFont(): JComponent {
        font = Font("JetBrains mono", font.style, font.size)
        return 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.toolwindow

import com.intellij.openapi.components.service
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.colors.EditorColorsListener
import com.intellij.openapi.editor.colors.EditorColorsManager
import com.intellij.openapi.editor.event.EditorMouseEvent
import com.intellij.openapi.editor.event.EditorMouseEventArea
import com.intellij.openapi.editor.event.EditorMouseMotionListener
import com.intellij.openapi.editor.ex.RangeHighlighterEx
import com.intellij.openapi.editor.impl.EditorImpl
import com.intellij.openapi.editor.markup.HighlighterLayer
import com.intellij.openapi.editor.markup.HighlighterTargetArea
import com.intellij.openapi.editor.markup.TextAttributes
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.openapi.ui.popup.JBPopupFactory
import com.intellij.openapi.util.Key
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.wm.ToolWindowManager
import com.intellij.ui.awt.RelativePoint
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.getPopupPositionAboveText
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.getRelativePathToContentRoot
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.horizontalPanelConstraints
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.EDITOR_CODE_REFERENCE_HOVER
import software.aws.toolkits.resources.message
import javax.swing.JLabel
import javax.swing.JPanel

class CodeWhispererCodeReferenceManager(private val project: Project) {
    val codeReferenceComponents = CodeWhispererCodeReferenceComponents(project)
    val toolWindow = ToolWindowManager.getInstance(project).getToolWindow(CodeWhispererCodeReferenceToolWindowFactory.id)
    val highlighters = mutableListOf<CodeReferenceHighLightContext>()
    var currentHighLightPopupContext: CodeReferenceHighLightPopupContext? = null
    private val referenceTextAttribute = TextAttributes().apply {
        effectColor = EDITOR_CODE_REFERENCE_HOVER
    }

    init {
        // Listen for global scheme changes
        project.messageBus.connect().subscribe(
            EditorColorsManager.TOPIC,
            EditorColorsListener { scheme ->
                if (scheme == null) return@EditorColorsListener
                codeReferenceComponents.apply {
                    contentPanel.background = scheme.defaultBackground
                    contentPanel.components.forEach {
                        it.background = scheme.defaultBackground
                    }
                }
            }
        )
    }

    fun showCodeReferencePanel() {
        toolWindow?.show()
    }

    fun insertCodeReference(states: InvocationContext, selectedIndex: Int) {
        val (requestContext, _, recommendationContext) = states
        val (_, editor, _, caretPosition) = requestContext
        val (_, detail, reformattedDetail) = recommendationContext.details[selectedIndex]
        val startOffset = caretPosition.offset
        val relativePath = getRelativePathToContentRoot(editor)
        reformattedDetail.references().forEachIndexed { i, reference ->
            val start = startOffset + reference.recommendationContentSpan().start()
            val end = startOffset + reference.recommendationContentSpan().end()
            val lineNums = getReferenceLineNums(editor, start, end)

            // There is an unformatted recommendation(directly from response) and reformatted one. We want to get
            // the line number, start/end offset of the reformatted one because it's the one inserted to the editor.
            // However, the one that shows in the tool window record should show the original recommendation, as below.
            val originalContentLines = getOriginalContentLines(detail, i)

            codeReferenceComponents.contentPanel.apply {
                add(
                    codeReferenceComponents.codeReferenceRecordPanel(reference, relativePath, lineNums),
                    horizontalPanelConstraints,
                    components.size - 1
                )

                // add each line of the original reference a JPanel to the tool window content panel
                originalContentLines.forEach { line ->
                    if (line.isEmpty()) return@forEach
                    add(
                        codeReferenceComponents.codeContentPanel(line),
                        horizontalPanelConstraints,
                        components.size - 1
                    )
                }
            }

            insertHighLightContext(editor, start, end, reference)
        }
    }

    fun getReferenceLineNums(editor: Editor, start: Int, end: Int): String {
        val startLine = editor.document.getLineNumber(start)
        val endLine = editor.document.getLineNumber(end)
        val lineNums = if (startLine == endLine) {
            (startLine + 1).toString()
        } else {
            "${startLine + 1} to ${endLine + 1}"
        }
        return lineNums
    }

    fun getOriginalContentLines(detail: Completion, i: Int): List<String> {
        val originalSpan = detail.references()[i].recommendationContentSpan()
        return detail.content()
            .substring(originalSpan.start(), originalSpan.end())
            .split("\n")
    }

    private fun insertHighLightContext(editor: Editor, start: Int, end: Int, reference: Reference) {
        val codeContent = editor.document.getText(TextRange.create(start, end))
        val referenceContent = message(
            "codewhisperer.toolwindow.popup.text",
            reference.licenseName(),
            reference.repository()
        )
        val highlighter = editor.markupModel.addRangeHighlighter(
            start,
            end,
            HighlighterLayer.LAST + 1,
            null,
            HighlighterTargetArea.EXACT_RANGE
        ) as RangeHighlighterEx
        highlighters.add(CodeReferenceHighLightContext(editor, highlighter, codeContent, referenceContent))
    }

    private fun showPopup() {
        currentHighLightPopupContext?.let {
            val (context, popup) = it
            val (editor, highlighter) = context
            val point = getPopupPositionAboveText(editor, popup, highlighter.startOffset)
            popup.show(RelativePoint(point))
            highlighter.textAttributes = referenceTextAttribute
        }
    }

    private fun hidePopup() {
        currentHighLightPopupContext?.let {
            val (context, popup) = it
            popup.cancel()
            context.highlighter.textAttributes = null
        }
        currentHighLightPopupContext = null
    }

    fun addListeners(editor: Editor) {
        // If there already is a listener attached to the editor, don't add a duplicated one
        if (editor.getUserData(listenerDisposable) != null) {
            return
        }
        editor.addEditorMouseMotionListener(
            object : EditorMouseMotionListener {
                override fun mouseMoved(e: EditorMouseEvent) {
                    // If not hover on text, hide any highlight if applicable and exit
                    if (e.area != EditorMouseEventArea.EDITING_AREA || !e.isOverText) {
                        hidePopup()
                        return
                    }
                    currentHighLightPopupContext?.let {
                        // There's an active highlight in the editor
                        val (localEditor, highlighter, codeContent) = it.context

                        // Check if mouse is in the range and the content matches the original one, else hide
                        // this highlight
                        if (highlighter.contains(e.offset)) {
                            if (!isHighlighterRangeMatchCodeContent(localEditor, highlighter, codeContent)) {
                                hidePopup()
                            }
                            return
                        }
                        hidePopup()
                    }

                    // Remove invalid highlighters when we are about to iterate them
                    val toRemove = highlighters.filter { !it.highlighter.isValid }
                    toRemove.forEach {
                        highlighters.remove(it)
                    }

                    // No highlight should be visible at this point, we still need to find if mouse is hover on any text
                    // with references
                    highlighters.forEach {
                        val (localEditor, highlighter, codeContent, referenceContent) = it
                        if (highlighter.contains(e.offset)) {
                            if (!isHighlighterRangeMatchCodeContent(localEditor, highlighter, codeContent)) {
                                return
                            }

                            // find a valid range to highlight, show the highlight and popup, no need to check others
                            val popup = JBPopupFactory.getInstance().createComponentPopupBuilder(
                                JPanel().apply {
                                    add(JLabel(referenceContent))
                                },
                                null
                            ).createPopup()
                            currentHighLightPopupContext = CodeReferenceHighLightPopupContext(it, popup)
                            showPopup()
                            return
                        }
                    }
                }
            },
            (editor as EditorImpl).disposable
        )
        editor.putUserData(listenerDisposable, true)
    }

    private fun RangeHighlighterEx.contains(offset: Int) = this.startOffset <= offset && offset < this.endOffset

    private fun isHighlighterRangeMatchCodeContent(
        editor: Editor,
        highlighter: RangeHighlighterEx,
        codeContent: String
    ): Boolean =
        highlighter.isValid && highlighter.endOffset <= editor.document.textLength &&
            editor.document.getText(TextRange.create(highlighter.startOffset, highlighter.endOffset)) == codeContent

    companion object {
        fun getInstance(project: Project): CodeWhispererCodeReferenceManager = project.service()
        private val listenerDisposable = Key.create<Boolean>("codewhisperer.reference.listener.disposable")
    }
}

data class CodeReferenceHighLightContext(
    val editor: Editor,
    val highlighter: RangeHighlighterEx,
    val codeContent: String,
    val referenceContent: String
)

data class CodeReferenceHighLightPopupContext(
    val context: CodeReferenceHighLightContext,
    val popup: JBPopup
)


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

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

import com.intellij.openapi.project.DumbAware
import com.intellij.openapi.project.Project
import com.intellij.openapi.wm.ToolWindow
import com.intellij.openapi.wm.ToolWindowFactory
import com.intellij.ui.components.JBScrollPane
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled

class CodeWhispererCodeReferenceToolWindowFactory : ToolWindowFactory, DumbAware {
    override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) {
        val toolWindowContent = toolWindow.contentManager.factory.createContent(
            JBScrollPane(CodeWhispererCodeReferenceManager.getInstance(project).codeReferenceComponents.contentPanel, 20, 30),
            null,
            false
        )

        toolWindowContent.isCloseable = false
        toolWindow.contentManager.addContent(toolWindowContent)
    }

    override fun shouldBeAvailable(project: Project): Boolean = isCodeWhispererEnabled(project)

    companion object {
        const val id = "aws.codewhisperer.codereference"
    }
}


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

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

import kotlin.math.ln

private val WORD_REGEX = """\w+""".toRegex()

// TODO: we still need NLTK tokenizer instead of this naive one
// tokenize given string and filter out non-english word
/**
 * Equivalent to the following:
 * String.tokenize(): List<String> {
 *     val s = this.split(" ")
 *     val res = mutableListOf<String>()
 *
 *     s.forEach {
 *         val temp = WORD_REGEX.findAll(it)
 *         res.addAll(temp.map { it.value })
 *     }
 *
 *     return res
 * }
 */
private fun String.tokenize(): List<String> = this.split(" ").map { str ->
    WORD_REGEX.findAll(str)
        .map { it.value }
        .toList()
}.flatten()

data class BM25Result(
    val docString: String,
    val score: Double
) : Comparable<BM25Result> {
    override fun compareTo(other: BM25Result): Int = compareValuesBy(other, this) { it.score }
}

// Kotlin implementation based on python library rank_bm25, refer to https://github.com/dorianbrown/rank_bm25 for more detail
abstract class BM25(val corpus: List<String>, tokenizer: (String) -> List<String>) {
    abstract val k1: Double
    abstract val b: Double
    abstract val epsilon: Double

    protected val corpusSize: Int = corpus.size
    protected val avgdl: Double
    protected val idf = mutableMapOf<String, Double>()
    protected val docLen = mutableListOf<Int>()
    protected val docFreqs = mutableListOf<Map<String, Int>>()

    protected val nd = mutableMapOf<String, Int>()

    protected val tokenize: (String) -> List<String> = tokenizer

    init {
        var numDoc = 0
        corpus
            .map { tokenize(it) }
            .forEach { document ->
                docLen.add(document.size)
                numDoc += document.size

                val frequencies = mutableMapOf<String, Int>()
                document.forEach { word ->
                    frequencies[word] = 1 + (frequencies[word] ?: 0)
                }
                docFreqs.add(frequencies)

                frequencies.forEach { (word, freq) ->
                    nd[word] = 1 + (nd[word] ?: 0)
                }
            }

        avgdl = numDoc.toDouble() / corpusSize

        calIdf(nd)
    }

    abstract fun calIdf(nd: Map<String, Int>)

    abstract fun score(query: String): List<BM25Result>

    fun topN(query: String, n: Int = 3): List<BM25Result> {
        val notSorted = score(query)
        val sorted = notSorted.sorted()

        return sorted.take(n)
    }
}

class BM250kapi(documentSets: List<String>, tokenizer: (String) -> List<String> = String::tokenize) : BM25(documentSets, tokenizer) {
    override val k1: Double
        get() = 1.5
    override val b: Double
        get() = 0.75
    override val epsilon: Double
        get() = 0.25

    override fun calIdf(nd: Map<String, Int>) {
        // collect idf sum to calculate an average idf for epsilon value
        var idfSum = 0.0
        // collect words with negative idf to set them a special epsilon value
        // idf can be negative if word is contained in mroe than half of documents
        val negativeIdfs = nd.mapNotNull { (word, freq) ->
            val idf = ln(this.corpusSize - freq + 0.5) - ln(freq + 0.5)
            this.idf[word] = idf
            idfSum += idf

            return@mapNotNull if (idf < 0) {
                word
            } else {
                null
            }
        }

        val averageIdf = idfSum / this.idf.size
        val eps = this.epsilon * averageIdf
        negativeIdfs.forEach { word ->
            this.idf[word] = eps
        }
    }

    override fun score(query: String): List<BM25Result> = this.docFreqs.mapIndexed { index, docFreq ->
        val score = tokenize(query).fold(0.0) { currScore, queryWord ->
            val queryWordFreqForDocument = docFreq[queryWord] ?: 0
            currScore + (idf[queryWord] ?: 0.0) * queryWordFreqForDocument * (k1 + 1) / (queryWordFreqForDocument + k1 * (1 - b + b * docLen[index] / avgdl))
        }

        BM25Result(corpus[index], score)
    }
}


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

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

import com.intellij.openapi.editor.colors.EditorColorsManager
import com.intellij.ui.Gray
import com.intellij.ui.JBColor
import com.intellij.util.ui.UIUtil
import java.awt.Color

object CodeWhispererColorUtil {
    val POPUP_HOVER = JBColor(Gray.xC0, Gray.xFF)
    val POPUP_BUTTON_BORDER = JBColor(Gray.x32, Gray.x64)
    val POPUP_PANEL_SEPARATOR = JBColor.border()
    val POPUP_DIM_HEX = JBColor.GRAY.getHexString()
    val POPUP_REF_NOTICE_HEX = JBColor(0x2097F6, 0x2097F6).getHexString()
    val POPUP_REF_INFO = Gray.x8C
    val TOOLWINDOW_BACKGROUND = EditorColorsManager.getInstance().globalScheme.defaultBackground
    val TOOLWINDOW_CODE = JBColor(0x629623, 0x629623)
    val EDITOR_CODE_REFERENCE_HOVER = JBColor(0x4B4D4D, 0x4B4D4D)
    val INACTIVE_TEXT_COLOR = UIUtil.getInactiveTextColor().getHexString()

    fun Color.getHexString() = String.format("#%02x%02x%02x", this.red, this.green, this.blue)
}


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

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

import com.intellij.openapi.editor.markup.EffectType
import com.intellij.openapi.editor.markup.TextAttributes
import com.intellij.ui.JBColor
import software.amazon.awssdk.regions.Region
import java.awt.Font
import java.text.SimpleDateFormat
import java.time.format.DateTimeFormatter

object CodeWhispererConstants {
    const val CHARACTERS_LIMIT = 10240
    const val BEGINNING_OF_FILE = 0
    const val FILENAME_CHARS_LIMIT = 1024
    const val INVOCATION_KEY_INTERVAL_THRESHOLD = 15
    val SPECIAL_CHARACTERS_LIST = listOf("{", "[", "(", ":")
    val PAIRED_BRACKETS = mapOf('{' to '}', '(' to ')', '[' to ']', '<' to '>')
    val PAIRED_QUOTES = setOf('"', '\'', '`')
    const val INVOCATION_TIME_INTERVAL_THRESHOLD = 2
    const val LEFT_CONTEXT_ON_CURRENT_LINE = 50
    const val POPUP_INFO_TEXT_SIZE = 11f
    const val POPUP_BUTTON_TEXT_SIZE = 12f
    const val POPUP_DELAY: Long = 250
    const val POPUP_DELAY_CHECK_INTERVAL: Long = 25
    const val IDLE_TIME_CHECK_INTERVAL: Long = 25
    const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 50L

    // TODO: this is currently set to 2050 to account for the server side 0.5 TPS and and extra 50 ms buffer to
    // avoid ThrottlingException as much as possible.
    const val INVOCATION_INTERVAL: Long = 2050

    const val CODEWHISPERER_LEARN_MORE_URI = "https://aws.amazon.com/codewhisperer"
    const val CODEWHISPERER_SSO_LEARN_MORE_URI = "https://docs.aws.amazon.com/toolkit-for-jetbrains/latest/userguide/codewhisperer-auth.html"
    const val CODEWHISPERER_LOGIN_LEARN_MORE_URI = "https://docs.aws.amazon.com/toolkit-for-jetbrains/latest/userguide/codewhisper-setup-general.html"
    const val CODEWHISPERER_LOGIN_HELP_URI = "https://docs.aws.amazon.com/toolkit-for-jetbrains/latest/userguide/setup-credentials.html"

    const val THROTTLING_MESSAGE = "Maximum recommendation count reached for this month."

    // Code scan feature constants
    val ISSUE_HIGHLIGHT_TEXT_ATTRIBUTES = TextAttributes(null, null, JBColor.YELLOW, EffectType.WAVE_UNDERSCORE, Font.PLAIN)
    const val JAVA_CODE_SCAN_TIMEOUT_IN_SECONDS: Long = 60
    const val JAVA_PAYLOAD_LIMIT_IN_BYTES = 1024 * 1024 // 1MB
    const val PYTHON_CODE_SCAN_TIMEOUT_IN_SECONDS: Long = 60
    const val PYTHON_PAYLOAD_LIMIT_IN_BYTES = 1024 * 200 // 200KB
    const val JS_CODE_SCAN_TIMEOUT_IN_SECONDS: Long = 60
    const val JS_PAYLOAD_LIMIT_IN_BYTES = 1024 * 200 // 200KB
    const val CODE_SCAN_POLLING_INTERVAL_IN_SECONDS: Long = 5
    const val CODE_SCAN_CREATE_PAYLOAD_TIMEOUT_IN_SECONDS: Long = 10
    const val TOTAL_BYTES_IN_KB = 1024
    const val TOTAL_BYTES_IN_MB = 1024 * 1024
    const val TOTAL_MILLIS_IN_SECOND = 1000
    const val TOTAL_SECONDS_IN_MINUTE: Long = 60L
    const val ACCOUNTLESS_START_URL = "accountless"

    // Date when Accountless is not supported
    val EXPIRE_DATE = SimpleDateFormat("yyyy-MM-dd").parse("2023-01-31")

    // Formatter for timestamp on accountless warn notification
    val TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm")

    object AutoSuggestion {
        const val SETTING_ID = "codewhisperer_autoSuggestionActivation"
        const val ACTIVATED = "Activated"
        const val DEACTIVATED = "Deactivated"
    }

    object Config {
        const val CODEWHISPERER_ENDPOINT = "https://codewhisperer.us-east-1.amazonaws.com/"

        const val CODEWHISPERER_IDPOOL_ID = "us-east-1:70717e99-906f-4add-908c-bd9074a2f5b9"

        val Sigv4ClientRegion = Region.US_EAST_1
        val BearerClientRegion = Region.US_EAST_1
    }
}


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

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

import com.intellij.openapi.application.ApplicationManager
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider
import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration
import software.amazon.awssdk.core.interceptor.Context
import software.amazon.awssdk.core.interceptor.ExecutionAttributes
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor
import software.amazon.awssdk.core.retry.RetryPolicy
import software.amazon.awssdk.http.SdkHttpRequest
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClientBuilder
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClientBuilder
import software.amazon.awssdk.services.cognitoidentity.CognitoIdentityClient
import software.aws.toolkits.core.ToolkitClientCustomizer
import software.aws.toolkits.jetbrains.core.AwsSdkClient
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.telemetry.AwsCognitoCredentialsProvider
import java.net.URI

// TODO: move this file to package /client
class CodeWhispererEndpointCustomizer : ToolkitClientCustomizer {

    override fun customize(
        credentialProvider: AwsCredentialsProvider?,
        tokenProvider: SdkTokenProvider?,
        regionId: String,
        builder: AwsClientBuilder<*, *>,
        clientOverrideConfiguration: ClientOverrideConfiguration.Builder
    ) {
        if (builder is CodeWhispererRuntimeClientBuilder) {
            builder
                .endpointOverride(URI.create(CodeWhispererConstants.Config.CODEWHISPERER_ENDPOINT))
                .region(CodeWhispererConstants.Config.BearerClientRegion)
            clientOverrideConfiguration.retryPolicy(RetryPolicy.none())
            clientOverrideConfiguration.addExecutionInterceptor(
                object : ExecutionInterceptor {
                    override fun modifyHttpRequest(context: Context.ModifyHttpRequest, executionAttributes: ExecutionAttributes): SdkHttpRequest {
                        val requestBuilder = context.httpRequest().toBuilder()
                        executionAttributes.attributes.forEach { (k, v) ->
                            if (k.toString() != "OperationName") return@forEach
                            val isMetricOptIn = CodeWhispererSettings.getInstance().isMetricOptIn()
                            if (v == "GenerateCompletions") {
                                requestBuilder.putHeader(OPTOUT_KEY_NAME, (!isMetricOptIn).toString())
                            }
                            return requestBuilder.build()
                        }
                        return context.httpRequest()
                    }
                }
            )
        } else if (builder is CodeWhispererClientBuilder) {
            clientOverrideConfiguration.addExecutionInterceptor(
                object : ExecutionInterceptor {
                    override fun modifyHttpRequest(context: Context.ModifyHttpRequest, executionAttributes: ExecutionAttributes): SdkHttpRequest {
                        val requestBuilder = context.httpRequest().toBuilder()
                        executionAttributes.attributes.forEach { (k, v) ->
                            if (k.toString() != "OperationName") return@forEach
                            if (v == "GetAccessToken") return requestBuilder.build()
                            val token = CodeWhispererExplorerActionManager.getInstance().resolveAccessToken() ?: return requestBuilder.build()
                            requestBuilder.putHeader(TOKEN_KEY_NAME, token)

                            val isMetricOptIn = CodeWhispererSettings.getInstance().isMetricOptIn()
                            if (v == "ListRecommendations") {
                                requestBuilder.putHeader(OPTOUT_KEY_NAME, (!isMetricOptIn).toString())
                            }
                            return requestBuilder.build()
                        }
                        return context.httpRequest()
                    }
                }
            )

            builder
                .region(CodeWhispererConstants.Config.Sigv4ClientRegion)

            if (!ApplicationManager.getApplication().isUnitTestMode) {
                builder
                    .credentialsProvider(
                        AwsCognitoCredentialsProvider(
                            CodeWhispererConstants.Config.CODEWHISPERER_IDPOOL_ID,
                            CognitoIdentityClient.builder()
                                .credentialsProvider(AnonymousCredentialsProvider.create())
                                .region(CodeWhispererConstants.Config.Sigv4ClientRegion)
                                .httpClient(AwsSdkClient.getInstance().sharedSdkClient())
                                .build()
                        )
                    )
            }
        }
    }

    companion object {
        internal const val TOKEN_KEY_NAME = "x-amzn-codewhisperer-token"
        internal const val OPTOUT_KEY_NAME = "x-amzn-codewhisperer-optout"
    }
}


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

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

import com.intellij.ide.actions.CopyContentRootPathProvider
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.TestSourcesFilter
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiFile
import com.intellij.util.gist.GistManager
import com.intellij.util.io.DataExternalizer
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.yield
import org.jetbrains.annotations.VisibleForTesting
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
import java.io.DataInput
import java.io.DataOutput
import java.util.Collections

private val contentRootPathProvider = CopyContentRootPathProvider()

private val codewhispererCodeChunksIndex = GistManager.getInstance()
    .newPsiFileGist("psi to code chunk index", 0, CodeWhispererCodeChunkExternalizer) { psiFile ->
        runBlocking {
            val fileCrawler = getFileCrawlerForLanguage(psiFile.programmingLanguage())
            val fileProducers = listOf<suspend (PsiFile) -> List<VirtualFile>> { psiFile -> fileCrawler.listRelevantFilesInEditors(psiFile) }
            FileContextProvider.getInstance(psiFile.project).extractCodeChunksFromFiles(psiFile, fileProducers)
        }
    }

private fun getFileCrawlerForLanguage(programmingLanguage: CodeWhispererProgrammingLanguage) = when (programmingLanguage) {
    is CodeWhispererJava -> JavaCodeWhispererFileCrawler
    is CodeWhispererPython -> PythonCodeWhispererFileCrawler
    else -> NoOpFileCrawler()
}

private object CodeWhispererCodeChunkExternalizer : DataExternalizer<List<Chunk>> {
    override fun save(out: DataOutput, value: List<Chunk>) {
        out.writeInt(value.size)
        value.forEach { chunk ->
            out.writeUTF(chunk.path)
            out.writeUTF(chunk.content)
            out.writeUTF(chunk.nextChunk)
        }
    }

    override fun read(`in`: DataInput): List<Chunk> {
        val result = mutableListOf<Chunk>()
        val size = `in`.readInt()
        repeat(size) {
            result.add(
                Chunk(
                    path = `in`.readUTF(),
                    content = `in`.readUTF(),
                    nextChunk = `in`.readUTF()
                )
            )
        }

        return result
    }
}

/**
 * [extractFileContext] will extract the context from a psi file provided
 * [extractSupplementalFileContext] supplemental means file context extracted from files other than the provided one
 */
interface FileContextProvider {
    fun extractFileContext(editor: Editor, psiFile: PsiFile): FileContextInfo

    suspend fun extractSupplementalFileContext(psiFile: PsiFile, fileContext: FileContextInfo): SupplementalContextInfo?

    suspend fun extractCodeChunksFromFiles(psiFile: PsiFile, fileProducers: List<suspend (PsiFile) -> List<VirtualFile>>): List<Chunk>

    fun isTestFile(psiFile: PsiFile): Boolean

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

class DefaultCodeWhispererFileContextProvider(private val project: Project) : FileContextProvider {
    override fun extractFileContext(editor: Editor, psiFile: PsiFile): FileContextInfo = CodeWhispererEditorUtil.getFileContextInfo(editor, psiFile)

    /**
     * codewhisperer extract the supplemental context with 2 different approaches depending on what type of file the target file is.
     * 1. source file -> explore files/classes imported from the target file + files within the same project root
     * 2. test file -> explore "focal file" if applicable, otherwise fall back to most "relevant" file.
     * for focal files, e.g. "MainTest.java" -> "Main.java", "test_main.py" -> "main.py"
     * for the most relevant file -> we extract "keywords" from files opened in editor then get the one with the highest similarity with target file
     */
    override suspend fun extractSupplementalFileContext(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo? {
        val startFetchingTimestamp = System.currentTimeMillis()
        val isTst = isTestFile(psiFile)

        val chunks = if (isTst && targetContext.programmingLanguage.isUTGSupported()) {
            extractSupplementalFileContextForTst(psiFile, targetContext)
        } else if (!isTst && targetContext.programmingLanguage.isSupplementalContextSupported()) {
            extractSupplementalFileContextForSrc(psiFile, targetContext)
        } else {
            LOG.debug { "${ if (isTst) "UTG" else "CrossFile" } not supported for ${targetContext.programmingLanguage.languageId}" }
            null
        }

        return chunks?.let {
            if (it.isNotEmpty()) {
                LOG.info { "Successfully fetched supplemental context." }
                it.forEachIndexed { index, chunk ->
                    LOG.info {
                        """
                            |---------------------------------------------------------------
                            | Chunk $index:
                            |    path = ${chunk.path},
                            |    score = ${chunk.score},
                            |    content = ${chunk.content}
                            |----------------------------------------------------------------
                        """.trimMargin()
                    }
                }
            } else {
                LOG.warn { "Failed to fetch supplemental context, empty list." }
            }

            SupplementalContextInfo(
                isUtg = isTst,
                contents = it,
                latency = System.currentTimeMillis() - startFetchingTimestamp,
                targetFileName = targetContext.filename
            )
        }
    }

    override suspend fun extractCodeChunksFromFiles(psiFile: PsiFile, fileProducers: List<suspend (PsiFile) -> List<VirtualFile>>): List<Chunk> {
        val parseFilesStart = System.currentTimeMillis()
        val hasUsed = Collections.synchronizedSet(mutableSetOf<VirtualFile>())
        val chunks = mutableListOf<Chunk>()

        for (fileProducer in fileProducers) {
            yield()
            val files = fileProducer(psiFile)
            files.forEach { file ->
                yield()
                if (hasUsed.contains(file)) {
                    return@forEach
                }
                val relativePath = runReadAction { contentRootPathProvider.getPathToElement(project, file, null) ?: file.path }
                chunks.addAll(file.toCodeChunk(relativePath))
                hasUsed.add(file)
                if (chunks.size > CHUNK_SIZE) {
                    LOG.debug { "finish fetching 60 chunks in ${System.currentTimeMillis() - parseFilesStart} ms" }
                    return chunks.take(CHUNK_SIZE)
                }
            }
        }

        LOG.debug { "finish fetching 60 chunks in ${System.currentTimeMillis() - parseFilesStart} ms" }
        return chunks.take(CHUNK_SIZE)
    }

    override fun isTestFile(psiFile: PsiFile) = when (psiFile.programmingLanguage()) {
        is CodeWhispererJava -> TestSourcesFilter.isTestSources(psiFile.virtualFile, project)
        is CodeWhispererPython -> PythonCodeWhispererFileCrawler.testFilenamePattern.matches(psiFile.name)
        else -> true
    }

    @VisibleForTesting
    suspend fun extractSupplementalFileContextForSrc(psiFile: PsiFile, targetContext: FileContextInfo): List<Chunk> {
        if (!targetContext.programmingLanguage.isSupplementalContextSupported()) return emptyList()

        // takeLast(11) will extract 10 lines (exclusing current line) of left context as the query parameter
        val query = targetContext.caretContext.leftFileContext.split("\n").takeLast(11).joinToString("\n")

        // step 1: prepare data
        val first60Chunks: List<Chunk> = try {
            runReadAction { codewhispererCodeChunksIndex.getFileData(psiFile) }
        } catch (e: TimeoutCancellationException) {
            throw e
        }

        yield()

        if (first60Chunks.isEmpty()) {
            LOG.warn {
                "0 chunks was found for supplemental context, fileName=${targetContext.filename}, " +
                    "programmingLanaugage: ${targetContext.programmingLanguage}"
            }
            return emptyList()
        }

        // we need to keep the reference to Chunk object because we will need to get "nextChunk" later after calculation
        val contentToChunk = first60Chunks.associateBy { it.content }

        // BM250 only take list of string as argument
        // step 2: bm25 calculation
        val timeBeforeBm25 = System.currentTimeMillis()
        val top3Chunks: List<BM25Result> = BM250kapi(first60Chunks.map { it.content }).topN(query)
        LOG.info { "Time ellapsed for BM25 algorithm: ${System.currentTimeMillis() - timeBeforeBm25} ms" }

        yield()

        // we use nextChunk as supplemental context
        return top3Chunks.mapNotNull { bm25Result ->
            contentToChunk[bm25Result.docString]?.let {
                Chunk(content = it.nextChunk, path = it.path, score = bm25Result.score)
            }
        }
    }

    @VisibleForTesting
    fun extractSupplementalFileContextForTst(psiFile: PsiFile, targetContext: FileContextInfo): List<Chunk> {
        if (!targetContext.programmingLanguage.isUTGSupported()) return emptyList()

        val focalFile = getFileCrawlerForLanguage(targetContext.programmingLanguage).findFocalFileForTest(psiFile)

        return focalFile?.let { file ->
            val relativePath = contentRootPathProvider.getPathToElement(project, file, null) ?: file.path
            listOf(
                Chunk(
                    content = UTG_PREFIX + file.content().let { it.substring(0, minOf(it.length, UTG_SEGMENT_SIZE)) },
                    path = relativePath
                )
            )
        }.orEmpty()
    }

    companion object {
        private val LOG = getLogger<DefaultCodeWhispererFileContextProvider>()
        private const val CHUNK_SIZE = 60
        private const val UTG_SEGMENT_SIZE = 10200
        private const val UTG_PREFIX = "UTG\n"
    }
}


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

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

import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.project.guessProjectDir
import com.intellij.openapi.vfs.VfsUtil
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiManager

/**
 * An interface define how do we parse and fetch files provided a psi file or project
 * since different language has its own way importing other files or its own naming style for test file
 */
interface FileCrawler {
    /**
     * parse the import statements provided a file
     * @param psiFile of the file we are search with
     * @return list of file reference from the import statements
     */
    suspend fun listFilesImported(psiFile: PsiFile): List<VirtualFile>

    fun listFilesUnderProjectRoot(project: Project): List<VirtualFile>

    /**
     * @param psiFile the file we are searching with, aka target file
     * @return Files under the same package as the given file and exclude the given file
     */
    fun listFilesWithinSamePackage(psiFile: PsiFile): List<VirtualFile>

    /**
     * should be invoked at test files e.g. MainTest.java, or test_main.py
     * @param psiFile psi of the test file we are searching with, e.g. MainTest.java
     * @return its source file e.g. Main.java, main.py or most relevant file if any
     */
    fun findFocalFileForTest(psiFile: PsiFile): VirtualFile?

    fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile>
}

class NoOpFileCrawler : FileCrawler {
    override suspend fun listFilesImported(psiFile: PsiFile): List<VirtualFile> = emptyList()

    override fun listFilesUnderProjectRoot(project: Project): List<VirtualFile> = emptyList()
    override fun findFocalFileForTest(psiFile: PsiFile): VirtualFile? = null

    override fun listFilesWithinSamePackage(psiFile: PsiFile): List<VirtualFile> = emptyList()

    override fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile> = emptyList()
}

abstract class CodeWhispererFileCrawler : FileCrawler {
    abstract val fileExtension: String
    abstract val testFilenamePattern: Regex

    override fun listFilesUnderProjectRoot(project: Project): List<VirtualFile> = project.guessProjectDir()?.let { rootDir ->
        VfsUtil.collectChildrenRecursively(rootDir).filter {
            it.path.endsWith(fileExtension)
        }
    }.orEmpty()

    abstract fun guessSourceFileName(tstFileName: String): String

    companion object {
        fun searchRelevantFileInEditors(target: PsiFile, keywordProducer: (psiFile: PsiFile) -> List<String>): VirtualFile? {
            val project = target.project
            val targetElements = keywordProducer(target)

            return FileEditorManager.getInstance(project).openFiles
                .filter { openedFile ->
                    openedFile.name != target.virtualFile.name && openedFile.extension == target.virtualFile.extension
                }
                .mapNotNull { openedFile -> PsiManager.getInstance(project).findFile(openedFile) }
                .maxByOrNull {
                    val elementsToCheck = keywordProducer(it)
                    countSubstringMatches(targetElements, elementsToCheck)
                }?.virtualFile
        }

        /**
         * how many elements in elementsToCheck is contained (as substring) in targetElements
         */
        fun countSubstringMatches(targetElements: List<String>, elementsToCheck: List<String>): Int = elementsToCheck.fold(0) { acc, elementToCheck ->
            val hasTarget = targetElements.any { it.contains(elementToCheck, ignoreCase = true) }
            if (hasTarget) {
                acc + 1
            } else {
                acc
            }
        }
    }
}


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

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

data class CodeWhispererMetadata(
    var insertEnd: Int = -1
)