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

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

import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import software.amazon.awssdk.services.codewhisperer.model.ArtifactType
import software.amazon.awssdk.services.codewhisperer.model.CodeScanFindingsSchema
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.CodeAnalysisFindingsSchema
import software.amazon.awssdk.services.codewhispererruntime.model.CodeAnalysisStatus
import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisResponse
import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisRequest
import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.util.transform

class CodeWhispererSdkMapperUtilTest {
    // TODO: revisit these test cases since if we miss one field, they don't capture

    @Test
    fun `CreateCodeScanRequest`() {
        val expected = CreateCodeScanRequest.builder()
            .programmingLanguage { it.languageName("python") }
            .artifacts(mapOf(ArtifactType.SOURCE_CODE to "1", ArtifactType.BUILT_JARS to "2"))
            .clientToken("token")
            .build()

        val actual = expected.transform()
        assertThat(actual)
            .isNotNull
            .isInstanceOf(StartCodeAnalysisRequest::class.java)
            .usingRecursiveComparison()
            .isEqualTo(expected)
    }

    @Test
    fun `StartCodeAnalysisResponse`() {
        val expected = StartCodeAnalysisResponse.builder()
            .jobId("jobid")
            .status(CodeAnalysisStatus.PENDING)
            .errorMessage("msg")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as StartCodeAnalysisResponse

        val actual = expected.transform()
        assertThat(actual)
            .isNotNull
            .isInstanceOf(CreateCodeScanResponse::class.java)
            .usingRecursiveComparison()
            .isEqualTo(expected)
    }

    @Test
    fun `GetCodeScanRequest`() {
        val expected = GetCodeScanRequest.builder()
            .jobId("123")
            .build()

        val actual = expected.transform()
        assertThat(actual)
            .isNotNull
            .isInstanceOf(GetCodeAnalysisRequest::class.java)
            .usingRecursiveComparison()
            .isEqualTo(expected)
    }

    @Test
    fun `GetCodeAnalysisResponse`() {
        val expected = GetCodeAnalysisResponse.builder()
            .status(CodeAnalysisStatus.COMPLETED)
            .errorMessage("msg")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as GetCodeAnalysisResponse

        val actual = expected.transform()
        assertThat(actual)
            .isNotNull
            .isInstanceOf(GetCodeScanResponse::class.java)
            .usingRecursiveComparison()
            .isEqualTo(expected)
    }

    @Test
    fun `ListCodeScanFindingsRequest`() {
        val expected = ListCodeScanFindingsRequest.builder()
            .codeScanFindingsSchema(CodeScanFindingsSchema.CODESCAN_FINDINGS_1_0)
            .jobId("id")
            .nextToken("123")
            .build()

        val actual = expected.transform()
        assertThat(actual.jobId()).isEqualTo(expected.jobId())
        assertThat(actual.nextToken()).isEqualTo(expected.nextToken())
        assertThat(actual.codeAnalysisFindingsSchema()).isEqualTo(CodeAnalysisFindingsSchema.CODEANALYSIS_FINDINGS_1_0)
    }

    @Test
    fun `ListCodeAnalysisFindingsResponse`() {
        val expected = ListCodeAnalysisFindingsResponse.builder()
            .codeAnalysisFindings(CodeScanFindingsSchema.UNKNOWN_TO_SDK_VERSION.toString()) // the model provided only has builder.codeScanFindings(String)
            .nextToken("token")
            .responseMetadata(metadata)
            .sdkHttpResponse(sdkHttpResponse)
            .build() as ListCodeAnalysisFindingsResponse

        val actual = expected.transform()
        assertThat(actual.nextToken())
            .isNotNull
            .isInstanceOf(String::class.java)
            .isEqualTo(expected.nextToken())
        assertThat(actual.codeScanFindings())
            .isNotNull
            .isInstanceOf(String::class.java)
            .isEqualTo(expected.codeAnalysisFindings())
    }
}


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

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

import com.intellij.openapi.application.ApplicationManager
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndWait
import kotlinx.coroutines.TimeoutCancellationException
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.doThrow
import org.mockito.kotlin.mock
import org.mockito.kotlin.whenever
import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
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.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererTriggerType

class CodeWhispererServiceTest {
    @Rule
    @JvmField
    val projectRule = JavaCodeInsightTestFixtureRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    private lateinit var sut: CodeWhispererService
    private lateinit var userGroupSetting: CodeWhispererUserGroupSettings

    @Before
    fun setUp() {
        sut = CodeWhispererService.getInstance()
        userGroupSetting = mock()

        ApplicationManager.getApplication().replaceService(CodeWhispererUserGroupSettings::class.java, userGroupSetting, disposableRule.disposable)
    }

    @Test
    fun `getRequestContext - cross file context should be non-null for cross-file user group`() {
        whenever(userGroupSetting.getUserGroup()).thenReturn(CodeWhispererUserGroup.CrossFile)
        val mockFileContextProvider = mock<FileContextProvider> {
            on { this.extractFileContext(any(), any()) } doReturn aFileContextInfo()
            onBlocking { this.extractSupplementalFileContext(any(), any()) } doThrow TimeoutCancellationException::class
        }
        projectRule.project.replaceService(FileContextProvider::class.java, mockFileContextProvider, disposableRule.disposable)
        val file = projectRule.fixture.addFileToProject("main.java", "public class Main {}")
        runInEdtAndWait {
            projectRule.fixture.openFileInEditor(file.virtualFile)
        }

        val actual = sut.getRequestContext(
            TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown()),
            projectRule.fixture.editor,
            projectRule.project,
            file,
            LatencyContext()
        )

        actual.supplementalContext.let {
            assertThat(it).isNotNull
            assertThat(it?.isProcessTimeout)
                .isNotNull
                .isEqualTo(
                    it?.latency?.let { latency ->
                        latency > CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
                    }
                )

            assertThat(it?.contents).isNotNull.isEmpty()
        }

        assertThat(actual.supplementalContext).isNotNull
    }

    @Test
    fun `getRequestContext - cross file context should be null for non-cross-file user group`() {
        whenever(userGroupSetting.getUserGroup()).thenReturn(CodeWhispererUserGroup.Control)
        val file = projectRule.fixture.addFileToProject("main.java", "public class Main {}")

        runInEdtAndWait {
            projectRule.fixture.openFileInEditor(file.virtualFile)
        }

        val actual = sut.getRequestContext(
            TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown()),
            projectRule.fixture.editor,
            projectRule.project,
            file,
            LatencyContext()
        )

        assertThat(actual.supplementalContext).isNull()
    }
}


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

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

import com.intellij.analysis.problemsView.toolWindow.ProblemsView
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.openapi.wm.RegisterToolWindowTask
import com.intellij.openapi.wm.ToolWindow
import com.intellij.openapi.wm.ToolWindowManager
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetsManager
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Ignore
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.never
import org.mockito.kotlin.spy
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.aws.toolkits.jetbrains.core.ToolWindowHeadlessManagerImpl
import software.aws.toolkits.jetbrains.core.explorer.AwsToolkitExplorerFactory
import software.aws.toolkits.jetbrains.core.explorer.AwsToolkitExplorerToolWindow
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfiguration
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidgetFactory
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceToolWindowFactory
import kotlin.test.fail

class CodeWhispererSettingsTest : CodeWhispererTestBase() {

    private lateinit var codewhispererServiceSpy: CodeWhispererService
    private lateinit var toolWindowHeadlessManager: ToolWindowHeadlessManagerImpl

    @Before
    override fun setUp() {
        super.setUp()
        codewhispererServiceSpy = spy(codewhispererService)
        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)

        // Create a mock ToolWindowManager with working implementation of setAvailable() and isAvailable()
        toolWindowHeadlessManager = object : ToolWindowHeadlessManagerImpl(projectRule.project) {
            private val myToolWindows: MutableMap<String, ToolWindow> = HashMap()
            override fun doRegisterToolWindow(id: String): ToolWindow {
                val toolWindow = object : MockToolWindow(projectRule.project) {
                    private var isAvailable = false
                    override fun setAvailable(value: Boolean) {
                        isAvailable = value
                    }

                    override fun isAvailable() = isAvailable
                }
                myToolWindows[id] = toolWindow
                return toolWindow
            }

            override fun getToolWindow(id: String?): ToolWindow? {
                return myToolWindows[id]
            }
        }
        projectRule.project.replaceService(ToolWindowManager::class.java, toolWindowHeadlessManager, disposableRule.disposable)

        ToolWindowManager.getInstance(projectRule.project).registerToolWindow(
            RegisterToolWindowTask(
                id = ProblemsView.ID
            )
        )
        ToolWindowManager.getInstance(projectRule.project).registerToolWindow(
            RegisterToolWindowTask(
                id = AwsToolkitExplorerFactory.TOOLWINDOW_ID
            )
        )
        ToolWindowManager.getInstance(projectRule.project).registerToolWindow(
            RegisterToolWindowTask(
                id = CodeWhispererCodeReferenceToolWindowFactory.id
            )
        )
        projectRule.project.service<StatusBarWidgetsManager>().updateWidget(CodeWhispererStatusBarWidgetFactory::class.java)
    }

    @Test
    fun `when isCodeWhispererEnabled is false, user not able to trigger CodeWhisperer manually`() {
        whenever(stateManager.checkActiveCodeWhispererConnectionType(projectRule.project)).thenReturn(CodeWhispererLoginType.Logout)
        assertThat(isCodeWhispererEnabled(projectRule.project)).isFalse
        invokeCodeWhispererService()
        verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
    }

    @Test
    fun `test disable auto trigger should make user not able to trigger CodeWhisperer automatically`() {
        stateManager.setAutoEnabled(false)
        assertThat(stateManager.isAutoEnabled()).isFalse
        runInEdtAndWait {
            projectRule.fixture.type(':')
            verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
        }
    }

    @Test
    fun `test CodeWhisperer components should have correct states on initialization with no persistent states`() {
        stateManager.loadState(CodeWhispererExploreActionState())
        CodeWhispererSettings.getInstance().loadState(CodeWhispererConfiguration())

        val problemsWindow = ProblemsView.getToolWindow(projectRule.project) ?: fail("Problems window not found")
        val codeReferenceWindow = ToolWindowManager.getInstance(projectRule.project).getToolWindow(
            CodeWhispererCodeReferenceToolWindowFactory.id
        ) ?: fail("Code Reference Log window not found")
        val statusBarWidgetFactory = projectRule.project.service<StatusBarWidgetsManager>().getWidgetFactories().firstOrNull {
            it.id == CodeWhispererStatusBarWidgetFactory.ID
        } ?: fail("CodeWhisperer status bar widget not found")

        runInEdtAndWait {
            assertThat(
                AwsToolkitExplorerToolWindow.toolWindow(projectRule.project)
            ).isNotNull
            assertThat(problemsWindow.contentManager.contentCount).isEqualTo(0)
            assertThat(codeReferenceWindow.isAvailable).isFalse
            assertThat(statusBarWidgetFactory.isAvailable(projectRule.project)).isFalse
            assertThat(settingsManager.isIncludeCodeWithReference()).isFalse
        }
    }

    // TODO: update this to be enable on enabling CodeWhisperer
    @Ignore
    @Test
    fun `test accept CodeWhisperer TOS will show CodeWhisperer UI components, and vice-versa`() {
        val problemsWindow = ProblemsView.getToolWindow(projectRule.project) ?: fail("Problems window not found")
        val codeReferenceWindow = ToolWindowManager.getInstance(projectRule.project).getToolWindow(
            CodeWhispererCodeReferenceToolWindowFactory.id
        ) ?: fail("Code Reference Log window not found")
        val statusBarWidgetFactory = projectRule.project.service<StatusBarWidgetsManager>().getWidgetFactories().firstOrNull {
            it.id == CodeWhispererStatusBarWidgetFactory.ID
        } ?: fail("CodeWhisperer status bar widget not found")
        val originalIsIncludeCodeWithReference = settingsManager.isIncludeCodeWithReference()

        runInEdtAndWait {
            assertThat(problemsWindow.contentManager.contentCount).isEqualTo(0)
            assertThat(codeReferenceWindow.isAvailable).isFalse
            assertThat(statusBarWidgetFactory.isAvailable(projectRule.project)).isFalse
            assertThat(settingsManager.isIncludeCodeWithReference()).isEqualTo(originalIsIncludeCodeWithReference)
        }

        runInEdtAndWait {
            assertThat(problemsWindow.contentManager.contentCount).isEqualTo(1)
            assertThat(codeReferenceWindow.isAvailable).isTrue
            assertThat(statusBarWidgetFactory.isAvailable(projectRule.project)).isTrue
            assertThat(settingsManager.isIncludeCodeWithReference()).isEqualTo(true)
        }
    }
}


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

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

import com.intellij.openapi.util.TextRange
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import software.aws.toolkits.telemetry.CodewhispererCompletionType
import software.aws.toolkits.telemetry.CodewhispererLanguage
import software.aws.toolkits.telemetry.CodewhispererTriggerType

class CodeWhispererStateTest : CodeWhispererTestBase() {

    @Test
    fun `test CodeWhisperer invocation sets request metadata correctly`() {
        withCodeWhispererServiceInvokedAndWait { states ->
            val actualRequestContext = states.requestContext
            val editor = projectRule.fixture.editor
            val (actualProject, actualEditor, actualTriggerTypeInfo, actualCaretPosition, actualFileContextInfo) = actualRequestContext
            val (actualCaretContext, actualFilename, actualProgrammingLanguage) = actualFileContextInfo

            assertThat(actualProject).isEqualTo(projectRule.project)
            assertThat(actualEditor).isEqualTo(editor)
            assertThat(actualTriggerTypeInfo.triggerType).isEqualTo(CodewhispererTriggerType.OnDemand)
            assertThat(actualTriggerTypeInfo.automatedTriggerType is CodeWhispererAutomatedTriggerType.Unknown).isTrue()
            assertThat(actualFilename).isEqualTo(pythonFileName)
            assertThat(actualProgrammingLanguage.languageId).isEqualTo(CodewhispererLanguage.Python.toString())

            val expectedCurrOffset = editor.caretModel.offset
            val document = editor.document
            val expectedCaretLeftFileContext = document.getText(TextRange(0, expectedCurrOffset))
            val expectedCaretRightFileContext = document.getText(TextRange(expectedCurrOffset, document.textLength))
            val (actualCaretLeftFileContext, actualCaretRightFileContext) = actualCaretContext
            assertThat(actualCaretLeftFileContext).isEqualTo(expectedCaretLeftFileContext)
            assertThat(actualCaretRightFileContext).isEqualTo(expectedCaretRightFileContext)

            val (actualOffset, actualLine) = actualCaretPosition
            assertThat(actualOffset).isEqualTo(expectedCurrOffset)
            assertThat(actualLine).isEqualTo(document.getLineNumber(expectedCurrOffset))
        }
    }

    @Test
    fun `test CodeWhisperer invocation sets response metadata correctly`() {
        withCodeWhispererServiceInvokedAndWait { states ->
            val actualResponseContext = states.responseContext
            assertThat(listOf(actualResponseContext.sessionId)).isEqualTo(
                pythonResponse.sdkHttpResponse().headers()[CodeWhispererService.KET_SESSION_ID]
            )
            assertThat(actualResponseContext.completionType).isEqualTo(CodewhispererCompletionType.Block)
        }
    }

    @Test
    fun `test CodeWhisperer invocation sets recommendation metadata correctly`() {
        withCodeWhispererServiceInvokedAndWait { states ->
            val actualRecommendationContext = states.recommendationContext
            val (actualDetailContexts, actualUserInput) = actualRecommendationContext

            assertThat(actualUserInput).isEqualTo("")
            val expectedCount = pythonResponse.completions().size
            assertThat(actualDetailContexts.size).isEqualTo(expectedCount)
            actualDetailContexts.forEachIndexed { i, actualDetailContext ->
                val (actualRequestId, actualRecommendationDetail, _, actualIsDiscarded) = actualDetailContext
                assertThat(actualRecommendationDetail.content()).isEqualTo(pythonResponse.completions()[i].content())
                assertThat(actualRequestId).isEqualTo(pythonResponse.responseMetadata().requestId())
                assertThat(actualIsDiscarded).isEqualTo(false)
            }
        }
    }

    @Test
    fun `test CodeWhisperer invocation sets initial typeahead and selected index correctly`() {
        withCodeWhispererServiceInvokedAndWait {
            val sessionContext = popupManagerSpy.sessionContext
            val actualSelectedIndex = sessionContext.selectedIndex
            val actualTypeahead = sessionContext.typeahead
            val actualTypeaheadOriginal = sessionContext.typeaheadOriginal

            assertThat(actualSelectedIndex).isEqualTo(0)
            assertThat(actualTypeahead).isEqualTo("")
            assertThat(actualTypeaheadOriginal).isEqualTo("")
        }
    }
}


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

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

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.openapi.wm.StatusBarWidgetFactory
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetSettings
import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetsManager
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.replaceService
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.eq
import org.mockito.kotlin.mock
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
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.status.CodeWhispererStatusBarManager
import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidgetFactory
import kotlin.test.fail

class CodeWhispererStatusBarManagerTest {
    @JvmField
    @Rule
    val applicationRule = ApplicationRule()

    @JvmField
    @Rule
    val projectRule = ProjectRule()

    @JvmField
    @Rule
    val disposableRule = DisposableRule()

    private lateinit var sut: CodeWhispererStatusBarManager
    private lateinit var widgetManager: StatusBarWidgetsManager
    private lateinit var settings: StatusBarWidgetSettings
    private lateinit var explorerActionManager: CodeWhispererExplorerActionManager
    private var factory: StatusBarWidgetFactory? = null

    @Before
    fun setup() {
        widgetManager = mock()
        projectRule.project.replaceService(StatusBarWidgetsManager::class.java, widgetManager, disposableRule.disposable)

        settings = mock()
        ApplicationManager.getApplication().replaceService(StatusBarWidgetSettings::class.java, settings, disposableRule.disposable)

        explorerActionManager = mock()
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, explorerActionManager, disposableRule.disposable)

        factory = ExtensionPointName<StatusBarWidgetFactory>("com.intellij.statusBarWidgetFactory").extensionList.find {
            it.id == CodeWhispererStatusBarWidgetFactory.ID
        }
    }

    @Test
    fun `updateWidget - logout`() {
        assertUpdateWidgetArgumentCorrect(CodeWhispererLoginType.Logout)
    }

    @Test
    fun `udpateWidget - Sono`() {
        assertUpdateWidgetArgumentCorrect(CodeWhispererLoginType.Sono)
    }

    @Test
    fun `updateWidget - SSO`() {
        assertUpdateWidgetArgumentCorrect(CodeWhispererLoginType.SSO)
    }

    @Test
    fun `updateWidget - Accountless`() {
        assertUpdateWidgetArgumentCorrect(CodeWhispererLoginType.Accountless)
    }

    private fun assertUpdateWidgetArgumentCorrect(loginType: CodeWhispererLoginType) {
        val isWidgetVisible = loginType != CodeWhispererLoginType.Logout
        sut = CodeWhispererStatusBarManager(projectRule.project)

        configureConnection(loginType)

        sut.updateWidget()

        factory?.let {
            verify(settings).setEnabled(eq(it), eq(isWidgetVisible))
            verify(widgetManager).updateWidget(eq(it))
        } ?: fail("factory should not be null")
    }

    private fun configureConnection(loginType: CodeWhispererLoginType) {
        whenever(explorerActionManager.checkActiveCodeWhispererConnectionType(any())).thenReturn(loginType)
    }
}


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

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

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.editor.VisualPosition
import com.intellij.openapi.project.Project
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.replaceService
import org.assertj.core.api.Assertions.assertThat
import org.junit.After
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.atLeastOnce
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.Import
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.aws.toolkits.core.telemetry.MetricEvent
import software.aws.toolkits.core.telemetry.TelemetryBatcher
import software.aws.toolkits.core.telemetry.TelemetryPublisher
import software.aws.toolkits.core.utils.test.aString
import software.aws.toolkits.jetbrains.AwsToolkit
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.model.CaretContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.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.CodeWhispererUserGroupStates
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.telemetry.CodeWhispererTelemetryService
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
import software.aws.toolkits.jetbrains.settings.AwsSettings
import software.aws.toolkits.telemetry.CodewhispererCompletionType
import software.aws.toolkits.telemetry.CodewhispererPreviousSuggestionState
import software.aws.toolkits.telemetry.CodewhispererSuggestionState
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import java.time.Duration
import kotlin.random.Random

// TODO: add more tests
class CodeWhispererTelemetryServiceTest {
    private class NoOpToolkitTelemetryService(
        publisher: TelemetryPublisher = NoOpPublisher(),
        batcher: TelemetryBatcher
    ) : TelemetryService(publisher, batcher)

    @Rule
    @JvmField
    val applicationRule = ApplicationRule()

    @Rule
    @JvmField
    val projectRule = ProjectRule()

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    private lateinit var sut: CodeWhispererTelemetryService
    private lateinit var telemetryServiceSpy: TelemetryService
    private lateinit var batcher: TelemetryBatcher

    @Before
    fun setup() {
        AwsSettings.getInstance().isTelemetryEnabled = true

        sut = spy(CodeWhispererTelemetryService.getInstance())

        batcher = mock()

        telemetryServiceSpy = NoOpToolkitTelemetryService(batcher = batcher)
        ApplicationManager.getApplication().replaceService(TelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)

        val groupSettings = CodeWhispererUserGroupSettings().apply {
            loadState(
                CodeWhispererUserGroupStates(
                    version = AwsToolkit.PLUGIN_VERSION,
                    settings = mapOf(CodeWhispererUserGroupSettings.USER_GROUP_KEY to CodeWhispererUserGroup.Control.name)
                )
            )
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererUserGroupSettings::class.java, groupSettings, disposableRule.disposable)
    }

    @After
    fun cleanup() {
        sut.previousDecisions().clear()
    }

    @Test
    fun `test recordSuggestionState`() {
        fun assertSuggestionStates(expectedStates: List<CodewhispererSuggestionState>) {
            val (recommendationContext, sessionContext) = aRecommendationContextAndSessionContext(expectedStates)
            val hasUserAccepted = expectedStates.any { it == CodewhispererSuggestionState.Accept }

            val details = recommendationContext.details
            val actualStates = mutableListOf<CodewhispererSuggestionState>()
            details.forEachIndexed { index, detail ->
                val suggestionState = sut.recordSuggestionState(
                    index,
                    sessionContext.selectedIndex,
                    sessionContext.seen.contains(index),
                    hasUserAccepted,
                    detail.isDiscarded,
                    detail.recommendation.content().isEmpty()
                )
                actualStates.add(suggestionState)
            }

            assertThat(actualStates).hasSize(expectedStates.size)
            actualStates.forEachIndexed { i, actual ->
                assertThat(actual).isEqualTo(expectedStates[i])
            }
        }

        assertSuggestionStates(listOf(CodewhispererSuggestionState.Ignore, CodewhispererSuggestionState.Ignore, CodewhispererSuggestionState.Accept))
        assertSuggestionStates(
            listOf(
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Unseen
            )
        )
        assertSuggestionStates(listOf(CodewhispererSuggestionState.Empty, CodewhispererSuggestionState.Empty, CodewhispererSuggestionState.Empty))
        assertSuggestionStates(
            listOf(
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Discard
            )
        )
        assertSuggestionStates(
            listOf(
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Accept,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Unseen
            )
        )
    }

    @Test
    fun `test aggregateUserDecision`() {
        fun assertAggregateUserDecision(decisions: List<CodewhispererSuggestionState>, expected: CodewhispererPreviousSuggestionState) {
            val actual = sut.aggregateUserDecision(decisions)
            assertThat(actual).isEqualTo(expected)
        }

        assertAggregateUserDecision(
            listOf(
                CodewhispererSuggestionState.Accept,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Unseen
            ),
            CodewhispererPreviousSuggestionState.Accept
        )

        assertAggregateUserDecision(
            listOf(
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Ignore
            ),
            CodewhispererPreviousSuggestionState.Reject
        )

        assertAggregateUserDecision(
            listOf(
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Discard,
                CodewhispererSuggestionState.Empty
            ),
            CodewhispererPreviousSuggestionState.Discard
        )

        assertAggregateUserDecision(
            listOf(
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty
            ),
            CodewhispererPreviousSuggestionState.Empty
        )
    }

    @Test
    fun `sendUserTriggerDecisionEvent`() {
        val timeSinceDocumentChanged = Random.nextDouble(0.0, 1000.0)
        val mockCodeWhispererInvocationStatus: CodeWhispererInvocationStatus = mock {
            on { getTimeSinceDocumentChanged() }.thenReturn(timeSinceDocumentChanged)
        }

        ApplicationManager.getApplication().replaceService(
            CodeWhispererInvocationStatus::class.java,
            mockCodeWhispererInvocationStatus,
            disposableRule.disposable
        )

        val supplementalContextInfo = aSupplementalContextInfo()
        val requestContext = aRequestContext(projectRule.project, mySupplementalContextInfo = supplementalContextInfo)
        val responseContext = aResponseContext()
        val recommendationContext = aRecommendationContext()
        val popupShownDuration = Duration.ofSeconds(Random.nextLong(0, 30))
        val suggestionState = CodewhispererSuggestionState.Reject

        val expectedTotalImportCount = recommendationContext.details.fold(0) { grandTotal, detail ->
            grandTotal + detail.recommendation.mostRelevantMissingImports().size
        }

        sut.sendUserTriggerDecisionEvent(requestContext, responseContext, recommendationContext, suggestionState, popupShownDuration)

        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeastOnce()).enqueue(capture())
            CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
                allValues,
                "codewhisperer_userTriggerDecision",
                1,
                "codewhispererSessionId" to responseContext.sessionId,
                "codewhispererFirstRequestId" to requestContext.latencyContext.firstRequestId,
                "codewhispererCompletionType" to responseContext.completionType,
                "codewhispererLanguage" to requestContext.fileContextInfo.programmingLanguage.toTelemetryType(),
                "codewhispererTriggerType" to requestContext.triggerTypeInfo.triggerType,
                "codewhispererAutomatedTriggerType" to requestContext.triggerTypeInfo.automatedTriggerType.telemetryType,
                "codewhispererLineNumber" to requestContext.caretPosition.line,
                "codewhispererCursorOffset" to requestContext.caretPosition.offset,
                "codewhispererSuggestionCount" to recommendationContext.details.size,
                "codewhispererSuggestionImportCount" to expectedTotalImportCount,
                "codewhispererTotalShownTime" to popupShownDuration?.toMillis()?.toDouble(),
                "codewhispererTriggerCharacter" to requestContext.triggerTypeInfo.automatedTriggerType.let {
                    if (it is CodeWhispererAutomatedTriggerType.SpecialChar) {
                        it.specialChar.toString()
                    } else {
                        null
                    }
                },
                "codewhispererTypeaheadLength" to recommendationContext.userInputSinceInvocation.length,
                "codewhispererTimeSinceLastDocumentChange" to timeSinceDocumentChanged,
                "codewhispererUserGroup" to "Control",
                "codewhispererSupplementalContextTimeout" to supplementalContextInfo.isProcessTimeout,
                "codewhispererSupplementalContextIsUtg" to supplementalContextInfo.isUtg,
                "codewhispererSupplementalContextLength" to supplementalContextInfo.contentLength
            )
        }
    }

    @Test
    fun `sendUserDecisionEventForAll will send userDecision event for all suggestions`() {
        doNothing().whenever(sut).sendUserTriggerDecisionEvent(any(), any(), any(), any(), any())
        val eventCount = mutableMapOf<CodewhispererSuggestionState, Int>()
        var totalEventCount = 0
        val requestContext = aRequestContext(projectRule.project)
        val responseContext = aResponseContext()

        fun assertUserDecision(decisions: List<CodewhispererSuggestionState>) {
            decisions.forEach { eventCount[it] = 1 + (eventCount[it] ?: 0) }
            totalEventCount += decisions.size

            val (recommendationContext, sessionContext) = aRecommendationContextAndSessionContext(decisions)
            val hasUserAccept = decisions.any { it == CodewhispererSuggestionState.Accept }
            val popupShownDuration = Duration.ofSeconds(Random.nextLong(0, 30))

            sut.sendUserDecisionEventForAll(requestContext, responseContext, recommendationContext, sessionContext, hasUserAccept, popupShownDuration)
            argumentCaptor<MetricEvent>().apply {
                verify(batcher, atLeastOnce()).enqueue(capture())

                CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
                    allValues,
                    "codewhisperer_userDecision",
                    count = totalEventCount,
                    "codewhispererUserGroup" to "Control",
                )

                eventCount.forEach { (k, v) ->
                    CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
                        allValues,
                        "codewhisperer_userDecision",
                        count = v,
                        "codewhispererSuggestionState" to k
                    )
                }
            }
        }

        assertUserDecision(
            listOf(
                CodewhispererSuggestionState.Accept,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Unseen
            )
        )

        assertUserDecision(
            listOf(
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Unseen
            ),
        )

        assertUserDecision(
            listOf(
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty
            ),
        )
    }

    @Test
    fun `sendUserDecisionEventForAll will aggregate suggestion level decisions and send out userTriggerDecision`() {
        fun helper(
            decisions: List<CodewhispererSuggestionState>,
            expectedState: CodewhispererSuggestionState,
            expectedPreviousSuggestionState: CodewhispererPreviousSuggestionState?
        ) {
            val supplementalContextInfo = aSupplementalContextInfo()
            val requestContext = aRequestContext(projectRule.project, mySupplementalContextInfo = supplementalContextInfo)
            val responseContext = aResponseContext()
            val (recommendationContext, sessionContext) = aRecommendationContextAndSessionContext(decisions)
            val hasUserAccept = decisions.any { it == CodewhispererSuggestionState.Accept }
            val popupShownDuration = Duration.ofSeconds(Random.nextLong(0, 30))

            sut.sendUserDecisionEventForAll(requestContext, responseContext, recommendationContext, sessionContext, hasUserAccept, popupShownDuration)
            argumentCaptor<MetricEvent>().apply {
                verify(batcher, atLeastOnce()).enqueue(capture())

                CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
                    allValues,
                    "codewhisperer_userTriggerDecision",
                    count = 1,
                    "codewhispererSuggestionState" to expectedState,
                    "codewhispererUserGroup" to "Control",
                    "codewhispererSupplementalContextTimeout" to supplementalContextInfo.isProcessTimeout,
                    "codewhispererSupplementalContextIsUtg" to supplementalContextInfo.isUtg,
                    "codewhispererSupplementalContextLength" to supplementalContextInfo.contentLength,
                    atLeast = true
                )

                CodeWhispererTelemetryTest.assertEventsContainsFieldsAndCount(
                    allValues,
                    "codewhisperer_userTriggerDecision",
                    count = 1,
                    "codewhispererPreviousSuggestionState" to expectedPreviousSuggestionState,
                    "codewhispererSupplementalContextTimeout" to supplementalContextInfo.isProcessTimeout,
                    "codewhispererSupplementalContextIsUtg" to supplementalContextInfo.isUtg,
                    "codewhispererSupplementalContextLength" to supplementalContextInfo.contentLength,
                    "codewhispererUserGroup" to "Control",
                    atLeast = true
                )
            }
        }

        val decisionQueue = sut.previousDecisions()
        assertThat(decisionQueue).isEmpty()

        helper(
            listOf(
                CodewhispererSuggestionState.Accept,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Ignore,
                CodewhispererSuggestionState.Unseen
            ),
            CodewhispererSuggestionState.Accept,
            null
        )
        assertThat(decisionQueue).hasSize(1)

        helper(
            listOf(
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Reject,
                CodewhispererSuggestionState.Unseen,
                CodewhispererSuggestionState.Unseen
            ),
            CodewhispererSuggestionState.Reject,
            CodewhispererPreviousSuggestionState.Accept
        )
        assertThat(decisionQueue).hasSize(2)

        helper(
            listOf(
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty,
                CodewhispererSuggestionState.Empty
            ),
            CodewhispererSuggestionState.Empty,
            CodewhispererPreviousSuggestionState.Reject
        )
        assertThat(decisionQueue).hasSize(3)
    }
}

// TODO: move these to testUtil
fun aRequestContext(project: Project, myFileContextInfo: FileContextInfo? = null, mySupplementalContextInfo: SupplementalContextInfo? = null): RequestContext {
    val triggerType = listOf(CodewhispererTriggerType.AutoTrigger, CodewhispererTriggerType.OnDemand).random()
    val automatedTriggerType = if (triggerType == CodewhispererTriggerType.AutoTrigger) {
        listOf(
            CodeWhispererAutomatedTriggerType.IdleTime(),
            CodeWhispererAutomatedTriggerType.Enter(),
            CodeWhispererAutomatedTriggerType.SpecialChar('a'),
            CodeWhispererAutomatedTriggerType.IntelliSense()
        ).random()
    } else {
        CodeWhispererAutomatedTriggerType.Unknown()
    }

    return RequestContext(
        project,
        mock(),
        TriggerTypeInfo(triggerType, automatedTriggerType),
        CaretPosition(Random.nextInt(), Random.nextInt()),
        fileContextInfo = myFileContextInfo ?: aFileContextInfo(),
        supplementalContext = mySupplementalContextInfo ?: aSupplementalContextInfo(),
        null,
        LatencyContext(
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextDouble(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            Random.nextLong(),
            aString()
        )
    )
}

fun aSupplementalContextInfo(myContents: List<Chunk>? = null, myIsUtg: Boolean? = null, myLatency: Long? = null): SupplementalContextInfo {
    val contents = mutableListOf<Chunk>()
    val numberOfContent = Random.nextInt(1, 4)
    repeat(numberOfContent) {
        contents.add(
            Chunk(
                content = aString(),
                path = aString(),
            )
        )
    }

    val isUtg = Random.nextBoolean()
    val latency = Random.nextLong(from = 0L, until = 100L)

    return SupplementalContextInfo(
        isUtg = myIsUtg ?: isUtg,
        latency = myLatency ?: latency,
        contents = myContents ?: contents,
        targetFileName = aString()
    )
}

fun aRecommendationContext(): RecommendationContext {
    val details = mutableListOf<DetailContext>()
    val size = Random.nextInt(1, 5)
    for (i in 1..size) {
        details.add(DetailContext(aString(), aCompletion(), aCompletion(), listOf(true, false).random(), listOf(true, false).random()))
    }

    return RecommendationContext(
        details,
        aString(),
        aString(),
        VisualPosition(Random.nextInt(1, 100), Random.nextInt(1, 100))
    )
}

/**
 * util to generate a RecommendationContext and a SessionContext given expected decisions
 */
fun aRecommendationContextAndSessionContext(decisions: List<CodewhispererSuggestionState>): Pair<RecommendationContext, SessionContext> {
    val table = CodewhispererSuggestionState.values().associateWith { 0 }.toMutableMap()
    decisions.forEach {
        table[it]?.let { curCount -> table[it] = 1 + curCount }
    }

    val details = mutableListOf<DetailContext>()
    decisions.forEach { decision ->
        val toAdd = if (decision == CodewhispererSuggestionState.Empty) {
            val completion = aCompletion(true, 0, 0)
            DetailContext(aString(), completion, completion, Random.nextBoolean(), Random.nextBoolean())
        } else if (decision == CodewhispererSuggestionState.Discard) {
            val completion = aCompletion()
            DetailContext(aString(), completion, completion, true, Random.nextBoolean())
        } else {
            val completion = aCompletion()
            DetailContext(aString(), completion, completion, false, Random.nextBoolean())
        }

        details.add(toAdd)
    }

    val recommendationContext = RecommendationContext(
        details,
        aString(),
        aString(),
        VisualPosition(Random.nextInt(1, 100), Random.nextInt(1, 100))
    )

    val selectedIndex = decisions.indexOfFirst { it == CodewhispererSuggestionState.Accept }.let {
        if (it != -1) {
            it
        } else {
            0
        }
    }

    val seen = mutableSetOf<Int>()
    decisions.forEachIndexed { index, decision ->
        if (decision != CodewhispererSuggestionState.Unseen) {
            seen.add(index)
        }
    }

    val sessionContext = SessionContext(
        selectedIndex = selectedIndex,
        seen = seen
    )
    return recommendationContext to sessionContext
}

fun aCompletion(isEmpty: Boolean = false, referenceCount: Int? = null, importCount: Int? = null): Completion {
    val myReferenceCount = referenceCount ?: Random.nextInt(0, 4)
    val myImportCount = importCount ?: Random.nextInt(0, 4)

    val references = List(myReferenceCount) {
        Reference.builder()
            .licenseName(aString())
            .build()
    }

    val imports = List(myImportCount) {
        Import.builder()
            .statement(aString())
            .build()
    }

    return Completion.builder()
        .content(if (!isEmpty) aString() else "")
        .references(references)
        .mostRelevantMissingImports(imports)
        .build()
}

fun aResponseContext(): ResponseContext = ResponseContext(aString(), listOf(CodewhispererCompletionType.Block, CodewhispererCompletionType.Line).random())

fun aFileContextInfo(language: CodeWhispererProgrammingLanguage? = null): FileContextInfo {
    val caretContextInfo = CaretContext(aString(), aString(), aString())
    val fileName = aString()

    val programmingLanguage = language ?: listOf(
        CodeWhispererPython.INSTANCE,
        CodeWhispererJava.INSTANCE
    ).random()

    return FileContextInfo(caretContextInfo, fileName, programmingLanguage)
}


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

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

import com.google.gson.Gson
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.psi.PsiDocumentManager
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.gradle.internal.impldep.com.amazonaws.ResponseMetadata
import org.junit.After
import org.junit.Before
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.atLeast
import org.mockito.kotlin.atLeastOnce
import org.mockito.kotlin.doAnswer
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.doReturnConsecutively
import org.mockito.kotlin.mock
import org.mockito.kotlin.never
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import org.mockito.kotlin.timeout
import org.mockito.kotlin.verify
import org.mockito.kotlin.whenever
import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata
import software.amazon.awssdk.http.SdkHttpResponse
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.aws.toolkits.core.telemetry.MetricEvent
import software.aws.toolkits.core.telemetry.TelemetryBatcher
import software.aws.toolkits.core.telemetry.TelemetryPublisher
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.emptyListResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.listOfEmptyRecommendationResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.listOfMixedEmptyAndNonEmptyRecommendationResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithNonEmptyToken
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testCodeWhispererException
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testRequestId
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testRequestIdForCodeWhispererException
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testSessionId
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.PauseCodeWhispererNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.ResumeCodeWhispererNode
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
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.telemetry.AcceptedSuggestionEntry
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererCodeCoverageTracker
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererUserModificationTracker
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
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.CodewhispererSuggestionState
import software.aws.toolkits.telemetry.CodewhispererTriggerType
import software.aws.toolkits.telemetry.Result
import java.time.Instant

class CodeWhispererTelemetryTest : CodeWhispererTestBase() {
    private val userDecision = "codewhisperer_userDecision"
    private val userModification = "codewhisperer_userModification"
    private val serviceInvocation = "codewhisperer_serviceInvocation"
    private val codePercentage = "codewhisperer_codePercentage"
    private val awsModifySetting = "aws_modifySetting"
    private val codewhispererSuggestionState = "codewhispererSuggestionState"

    private class TestTelemetryService(
        publisher: TelemetryPublisher = NoOpPublisher(),
        batcher: TelemetryBatcher
    ) : TelemetryService(publisher, batcher)

    private lateinit var telemetryService: TelemetryService
    private lateinit var batcher: TelemetryBatcher
    private lateinit var telemetryServiceSpy: TelemetryService
    private var isTelemetryEnabledDefault: Boolean = false

    @Before
    override fun setUp() {
        super.setUp()
        batcher = mock<TelemetryBatcher>()
        telemetryService = TestTelemetryService(batcher = batcher)
        telemetryServiceSpy = spy(telemetryService)
        ApplicationManager.getApplication().replaceService(TelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)
        isTelemetryEnabledDefault = AwsSettings.getInstance().isTelemetryEnabled
        AwsSettings.getInstance().isTelemetryEnabled = true
    }

    @Test
    fun `test pre-setup failure will send service invocation event with failed status`() {
        val codewhispererServiceSpy = spy(codewhispererService) {
            onGeneric { getRequestContext(any(), any(), any(), any(), any()) }
                .doAnswer { throw Exception() }
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)

        invokeCodeWhispererService()

        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeastOnce()).enqueue(capture())
            assertEventsContainsFieldsAndCount(
                allValues,
                serviceInvocation,
                1,
                "result" to Result.Failed.toString()
            )
        }
    }

    @Test
    fun `test accepting recommendation will send user modification events with 1 accepted other unseen`() {
        val trackerSpy = spy(CodeWhispererUserModificationTracker(projectRule.project))
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        projectRule.project.replaceService(CodeWhispererUserModificationTracker::class.java, trackerSpy, disposableRule.disposable)

        runInEdtAndWait {
            val fakeSuggestionEntry = AcceptedSuggestionEntry(
                Instant.now().minusSeconds(350),
                projectRule.fixture.file.virtualFile,
                projectRule.fixture.editor.document.createRangeMarker(0, 0, true),
                "",
                testSessionId,
                testRequestId,
                0,
                CodewhispererTriggerType.OnDemand,
                CodewhispererCompletionType.Line,
                CodewhispererLanguage.Java,
                CodewhispererRuntime.Java11,
                "",
                null
            )
            trackerSpy.stub {
                onGeneric {
                    trackerSpy.enqueue(any())
                }.doAnswer {
                    trackerSpy.enqueue(fakeSuggestionEntry)
                }.thenCallRealMethod()
            }
        }

        withCodeWhispererServiceInvokedAndWait {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }

        trackerSpy.dispose()
        val count = pythonResponse.completions().size
        argumentCaptor<MetricEvent>().apply {
            // 1 serviceInvocation + 1 userModification + 1 userDecision for accepted +
            // (count - 1) userDecisions for ignored
            verify(batcher, atLeast(2 + count)).enqueue(capture())
            assertEventsContainsFieldsAndCount(allValues, serviceInvocation, 1)
            assertEventsContainsFieldsAndCount(
                allValues,
                userModification,
                1,
                "codewhispererSessionId" to testSessionId,
                "codewhispererUserGroup" to userGroup.name,
            )
            assertEventsContainsFieldsAndCount(
                allValues,
                userDecision,
                1,
                codewhispererSuggestionState to CodewhispererSuggestionState.Accept.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
            assertEventsContainsFieldsAndCount(
                allValues,
                userDecision,
                count - 1,
                codewhispererSuggestionState to CodewhispererSuggestionState.Unseen.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test cancelling popup will send user decision event for all unseen but one rejected`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        withCodeWhispererServiceInvokedAndWait { states ->
            popupManagerSpy.cancelPopup(states.popup)

            val count = pythonResponse.completions().size
            argumentCaptor<MetricEvent>().apply {
                verify(batcher, atLeast(1 + count)).enqueue(capture())
                assertEventsContainsFieldsAndCount(
                    allValues,
                    serviceInvocation,
                    1,
                    "result" to Result.Succeeded.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
                assertEventsContainsFieldsAndCount(
                    allValues,
                    userDecision,
                    1,
                    codewhispererSuggestionState to CodewhispererSuggestionState.Reject.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
                assertEventsContainsFieldsAndCount(
                    allValues,
                    userDecision,
                    count - 1,
                    codewhispererSuggestionState to CodewhispererSuggestionState.Unseen.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
            }
        }
    }

    @Test
    fun `test invoking CodeWhisperer will send service invocation event with succeeded status`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        withCodeWhispererServiceInvokedAndWait {
            argumentCaptor<MetricEvent>().apply {
                verify(batcher, atLeastOnce()).enqueue(capture())
                assertEventsContainsFieldsAndCount(
                    allValues,
                    serviceInvocation,
                    1,
                    "result" to Result.Succeeded.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
            }
        }
    }

    @Test
    fun `test moving caret backwards before getting back response will send user decision events for all discarded`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val editorManagerSpy = spy(editorManager)
        editorManagerSpy.stub {
            onGeneric {
                getCaretMovement(any(), any())
            } doAnswer {
                CaretMovement.MOVE_BACKWARD
            }
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererEditorManager::class.java, editorManagerSpy, disposableRule.disposable)

        invokeCodeWhispererService()

        val count = pythonResponse.completions().size
        runInEdtAndWait {
            argumentCaptor<MetricEvent>().apply {
                verify(batcher, atLeast(1 + count)).enqueue(capture())
                assertEventsContainsFieldsAndCount(
                    allValues,
                    serviceInvocation,
                    1,
                    "result" to Result.Succeeded.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
                assertEventsContainsFieldsAndCount(
                    allValues,
                    userDecision,
                    count,
                    codewhispererSuggestionState to CodewhispererSuggestionState.Discard.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                )
            }
        }
    }

    @Test
    fun `test user's typeahead before getting response will discard recommendations whose prefix not matching`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val editorManagerSpy = spy(editorManager)
        val typeahead = "(x, y)"
        editorManagerSpy.stub {
            onGeneric {
                getUserInputSinceInvocation(any(), any())
            } doAnswer {
                typeahead
            }
            onGeneric {
                getCaretMovement(any(), any())
            } doAnswer {
                CaretMovement.MOVE_FORWARD
            }
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererEditorManager::class.java, editorManagerSpy, disposableRule.disposable)

        withCodeWhispererServiceInvokedAndWait { }
        val prefixNotMatchCount = pythonResponse.completions().filter {
            !it.content().startsWith(typeahead)
        }.size
        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeast(1 + prefixNotMatchCount)).enqueue(capture())
            assertEventsContainsFieldsAndCount(
                allValues,
                serviceInvocation,
                1,
                "result" to Result.Succeeded.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
            assertEventsContainsFieldsAndCount(
                allValues,
                userDecision,
                prefixNotMatchCount,
                codewhispererSuggestionState to CodewhispererSuggestionState.Discard.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test getting non-CodeWhisperer Exception will return empty request id`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        mockClient.stub {
            on {
                this.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                throw Exception("wrong path")
            }
        }

        invokeCodeWhispererService()

        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeastOnce()).enqueue(capture())
            assertEventsContainsFieldsAndCount(
                allValues,
                serviceInvocation,
                1,
                "codewhispererRequestId" to "",
                "result" to Result.Failed.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test getting CodeWhispererException will capture request id`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        mockClient.stub {
            on {
                this.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                throw testCodeWhispererException
            }
        }

        invokeCodeWhispererService()

        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeastOnce()).enqueue(capture())
            assertEventsContainsFieldsAndCount(
                allValues,
                serviceInvocation,
                1,
                "codewhispererRequestId" to testRequestIdForCodeWhispererException,
                "result" to Result.Failed.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test user Decision events record CodeWhisperer reference info`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        withCodeWhispererServiceInvokedAndWait {}
        argumentCaptor<MetricEvent>().apply {
            verify(batcher, atLeastOnce()).enqueue(capture())
            pythonResponse.completions().forEach {
                assertEventsContainsFieldsAndCount(
                    allValues,
                    userDecision,
                    1,
                    "codewhispererSuggestionReferences" to Gson().toJson(it.references().map { ref -> ref.licenseName() }.toSet()),
                    "codewhispererSuggestionReferenceCount" to it.references().size.toString(),
                    "codewhispererUserGroup" to userGroup.name,
                    atLeast = true
                )
            }
        }
    }

    @Test
    fun `test invoking CodeWhisperer will send service invocation event with sessionId and requestId from response`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        withCodeWhispererServiceInvokedAndWait { states ->
            val metricCaptor = argumentCaptor<MetricEvent>()
            verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                serviceInvocation,
                1,
                "codewhispererSessionId" to states.responseContext.sessionId,
                "codewhispererRequestId" to states.recommendationContext.details[0].requestId,
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test userDecision events will record sessionId and requestId from response`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val statesCaptor = argumentCaptor<InvocationContext>()
        withCodeWhispererServiceInvokedAndWait {}
        verify(popupManagerSpy, timeout(5000).atLeastOnce()).render(statesCaptor.capture(), any(), any(), any(), any())
        val states = statesCaptor.lastValue
        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            userDecision,
            states.recommendationContext.details.size,
            "codewhispererSessionId" to states.responseContext.sessionId,
            "codewhispererRequestId" to states.recommendationContext.details[0].requestId,
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    @Test
    fun `test showing IntelliSense after triggering CodeWhisperer will send userDecision events of state Discard`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val codewhispererServiceSpy = spy(codewhispererService)
        codewhispererServiceSpy.stub {
            onGeneric {
                canDoInvocation(any(), any())
            } doAnswer {
                true
            }
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
        popupManagerSpy.stub {
            onGeneric {
                hasConflictingPopups(any())
            } doAnswer {
                true
            }
        }
        invokeCodeWhispererService()

        runInEdtAndWait {
            val metricCaptor = argumentCaptor<MetricEvent>()
            verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                pythonResponse.completions().size,
                codewhispererSuggestionState to CodewhispererSuggestionState.Discard.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test codePercentage tracker will not be activated if CWSPR terms of service is not accepted`() {
        val exploreManagerMock = mock<CodeWhispererExplorerActionManager> {
            on { checkActiveCodeWhispererConnectionType(projectRule.project) } doReturn CodeWhispererLoginType.Logout
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreManagerMock, disposableRule.disposable)
        val project = projectRule.project
        val fixture = projectRule.fixture
        val tracker = CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE)
        assertThat(tracker.isTrackerActive()).isFalse
        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString("arbitrary string to trigger documentChanged")
            }
        }
        assertThat(tracker.isTrackerActive()).isFalse
    }

    @Test
    fun `test codePercentage tracker will not be initialized with unsupportedLanguage`() {
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(0)
        val project = projectRule.project
        val fixture = projectRule.fixture
        val plainTxtFile = fixture.addFileToProject("/notSupported.txt", "")

        runInEdtAndWait {
            fixture.openFileInEditor(plainTxtFile.virtualFile)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString("random txt string")
            }
        }

        // document changes in unsupported files will not trigger initialization of tracker
        assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(0)
    }

    @Test
    fun `test codePercentage metric is correct - 1`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val project = projectRule.project
        val fixture = projectRule.fixture
        val emptyFile = fixture.addFileToProject("/anotherFile.py", "")
        // simulate users typing behavior of the following
        // def addTwoNumbers(
        runInEdtAndWait {
            fixture.openFileInEditor(emptyFile.virtualFile)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString(pythonTestLeftContext)
            }
        }
        // simulate users accepting the recommendation and delete part of the recommendation
        // (x, y):\n    return x + y
        val deletedTokenByUser = 4
        withCodeWhispererServiceInvokedAndWait {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }

        runInEdtAndWait {
            val offset = fixture.caretOffset
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.document.deleteString(offset - deletedTokenByUser, offset)
                fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
            }
        }

        CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE).dispose()

        val acceptedTokensSize = pythonResponse.completions()[0].content().length - deletedTokenByUser
        val totalTokensSize = pythonTestLeftContext.length + pythonResponse.completions()[0].content().length - deletedTokenByUser

        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            codePercentage,
            1,
            "codewhispererAcceptedTokens" to acceptedTokensSize.toString(),
            "codewhispererTotalTokens" to totalTokensSize.toString(),
            "codewhispererPercentage" to CodeWhispererCodeCoverageTracker.calculatePercentage(acceptedTokensSize, totalTokensSize).toString(),
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    @Test
    fun `test codePercentage metric is correct - 2`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val project = projectRule.project
        val fixture = projectRule.fixture
        val emptyFile = fixture.addFileToProject("/anotherFile.py", "")
        // simulate users typing behavior of the following
        // def addTwoNumbers
        runInEdtAndWait {
            fixture.openFileInEditor(emptyFile.virtualFile)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString(pythonTestLeftContext)
            }
        }
        // simulate users accepting the recommendation
        // (x, y):\n    return x + y
        val anotherCodeSnippet = "\ndef functionWritenByMyself():\n\tpass()"
        withCodeWhispererServiceInvokedAndWait {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }

        runInEdtAndWait {
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString(anotherCodeSnippet)
                val currentOffset = fixture.editor.caretModel.offset
                // delete 1 char
                fixture.editor.document.deleteString(currentOffset - 1, currentOffset)
            }
            // use dispose() to force tracker to emit telemetry
            CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE).dispose()
        }

        val acceptedTokensSize = pythonResponse.completions()[0].content().length
        val totalTokensSize = pythonTestLeftContext.length + pythonResponse.completions()[0].content().length + anotherCodeSnippet.length - 1

        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            codePercentage,
            1,
            "codewhispererAcceptedTokens" to acceptedTokensSize.toString(),
            "codewhispererTotalTokens" to totalTokensSize.toString(),
            "codewhispererPercentage" to CodeWhispererCodeCoverageTracker.calculatePercentage(acceptedTokensSize, totalTokensSize).toString(),
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    /**
     * When user accept recommendation in file1, then switch and delete code in file2, if the deletion code in file2 is on pre-existing code,
     * we should not decrement totalTokens by the length of the code deleted
     */
    @Test
    fun `test codePercentage metric - switching files and delete tokens`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val project = projectRule.project
        val fixture = projectRule.fixture
        fixture.configureByText("/file1.py", pythonTestLeftContext)
        runInEdt {
            fixture.editor.caretModel.moveToOffset(fixture.editor.document.textLength)
        }
        val file2 = fixture.addFileToProject("./file2.py", "Pre-existing string")

        // accept recommendation in file1.py
        withCodeWhispererServiceInvokedAndWait {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }

        // switch to file2.py and delete code there
        runInEdtAndWait {
            fixture.openFileInEditor(file2.virtualFile)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.document.deleteString(0, file2.textLength)
            }
        }

        // use dispose() to force tracker to emit telemetry
        CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE).dispose()

        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            codePercentage,
            1,
            "codewhispererAcceptedTokens" to pythonResponse.completions()[0].content().length.toString(),
            "codewhispererTotalTokens" to pythonResponse.completions()[0].content().length.toString(),
            "codewhispererPercentage" to "100",
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    @Test
    fun `test codePercentage metric is correct - simulate IDE adding right closing paren`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        val project = projectRule.project
        val fixture = projectRule.fixture
        val emptyFile = fixture.addFileToProject("/anotherFile.py", "")
        // simulate users typing behavior of the following
        // def addTwoNumbers(
        whenever(mockClient.generateCompletions(any<GenerateCompletionsRequest>())).thenReturn(
            GenerateCompletionsResponse.builder()
                .completions(
                    CodeWhispererTestUtil.generateMockCompletionDetail("x, y):\n    return x + y"),
                    CodeWhispererTestUtil.generateMockCompletionDetail("a, b):\n    return a + b"),
                )
                .nextToken("")
                .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(ResponseMetadata.AWS_REQUEST_ID to testRequestId)))
                .sdkHttpResponse(SdkHttpResponse.builder().headers(mapOf(CodeWhispererService.KET_SESSION_ID to listOf(testSessionId))).build())
                .build() as GenerateCompletionsResponse
        )

        runInEdtAndWait {
            fixture.openFileInEditor(emptyFile.virtualFile)
            WriteCommandAction.runWriteCommandAction(project) {
                fixture.editor.appendString("$pythonTestLeftContext(")
                // add closing paren but not move the caret position, simulating IDE's behavior
                fixture.editor.document.insertString(fixture.editor.caretModel.offset, ")")
            }
        }

        withCodeWhispererServiceInvokedAndWait {
            popupManagerSpy.popupComponents.acceptButton.doClick()
        }
        CodeWhispererCodeCoverageTracker.getInstance(CodeWhispererPython.INSTANCE).dispose()

        val acceptedTokensSize = "x, y):\n    return x + y".length
        val totalTokensSize = "$pythonTestLeftContext(".length + acceptedTokensSize
        val metricCaptor = argumentCaptor<MetricEvent>()
        verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            codePercentage,
            1,
            "codewhispererAcceptedTokens" to acceptedTokensSize.toString(),
            "codewhispererTotalTokens" to totalTokensSize.toString(),
            "codewhispererPercentage" to CodeWhispererCodeCoverageTracker.calculatePercentage(acceptedTokensSize, totalTokensSize).toString(),
            "codewhispererUserGroup" to userGroup.name,
        )
    }

    @Test
    fun `test empty list of recommendations should sent 1 empty userDecision event and no popup shown`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                emptyListResponse
            }
        }
        invokeCodeWhispererService()

        verify(popupManagerSpy, never()).showPopup(any(), any(), any(), any(), any())
        runInEdtAndWait {
            val metricCaptor = argumentCaptor<MetricEvent>()
            verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                1,
                "codewhispererSuggestionState" to CodewhispererSuggestionState.Empty.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test getting some recommendations and a final empty list of recommendations at session end should not sent empty userDecision events`() {
        val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doReturnConsecutively(listOf(pythonResponseWithNonEmptyToken, emptyListResponse))
        }

        withCodeWhispererServiceInvokedAndWait { }

        runInEdtAndWait {
            val metricCaptor = argumentCaptor<MetricEvent>()
            // 2 serviceInvocation + 5 userDecision
            verify(batcher, atLeast(pythonResponseWithNonEmptyToken.completions().size + 2)).enqueue(metricCaptor.capture())
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                0,
                "codewhispererSuggestionState" to CodewhispererSuggestionState.Empty.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                1,
                "codewhispererSuggestionState" to CodewhispererSuggestionState.Accept.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                4,
                "codewhispererSuggestionState" to CodewhispererSuggestionState.Unseen.toString(),
                "codewhispererUserGroup" to userGroup.name,
            )
        }
    }

    @Test
    fun `test empty recommendations should send empty user decision events`() {
        testSendEmptyUserDecisionEventForEmptyRecommendations(listOfEmptyRecommendationResponse)
    }

    @Test
    fun `test a mix of empty and non-empty recommendations should send empty user decision events accordingly`() {
        testSendEmptyUserDecisionEventForEmptyRecommendations(listOfMixedEmptyAndNonEmptyRecommendationResponse)
    }

    @Test
    fun `test toggle autoSugestion will emit autoSuggestionActivation telemetry`() {
        val metricCaptor = argumentCaptor<MetricEvent>()
        doNothing().`when`(batcher).enqueue(metricCaptor.capture())

        PauseCodeWhispererNode(projectRule.project).onDoubleClick(mock())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            awsModifySetting,
            1,
            "settingId" to CodeWhispererConstants.AutoSuggestion.SETTING_ID,
            "settingState" to CodeWhispererConstants.AutoSuggestion.DEACTIVATED
        )

        ResumeCodeWhispererNode(projectRule.project).onDoubleClick(mock())
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            awsModifySetting,
            1,
            "settingId" to CodeWhispererConstants.AutoSuggestion.SETTING_ID,
            "settingState" to CodeWhispererConstants.AutoSuggestion.ACTIVATED
        )
        assertEventsContainsFieldsAndCount(
            metricCaptor.allValues,
            awsModifySetting,
            2,
        )
    }

    private fun testSendEmptyUserDecisionEventForEmptyRecommendations(response: GenerateCompletionsResponse) {
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                response
            }
        }

        invokeCodeWhispererService()

        val numOfEmptyRecommendations = response.completions().filter { it.content().isEmpty() }.size
        if (numOfEmptyRecommendations == response.completions().size) {
            verify(popupManagerSpy, never()).showPopup(any(), any(), any(), any(), any())
        } else {
            val popupCaptor = argumentCaptor<JBPopup>()
            verify(popupManagerSpy, timeout(5000))
                .showPopup(any(), any(), popupCaptor.capture(), any(), any())
            runInEdtAndWait {
                popupManagerSpy.closePopup(popupCaptor.lastValue)
            }
        }
        runInEdtAndWait {
            val metricCaptor = argumentCaptor<MetricEvent>()
            verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture())
            assertEventsContainsFieldsAndCount(
                metricCaptor.allValues,
                userDecision,
                numOfEmptyRecommendations,
                "codewhispererSuggestionState" to CodewhispererSuggestionState.Empty.toString()
            )
        }
    }

    private fun Editor.appendString(string: String) {
        val currentOffset = caretModel.primaryCaret.offset
        document.insertString(currentOffset, string)
        caretModel.moveToOffset(currentOffset + string.length)
        PsiDocumentManager.getInstance(projectRule.project).commitDocument(document)
    }

    @After
    override fun tearDown() {
        super.tearDown()
        telemetryService.dispose()
        AwsSettings.getInstance().isTelemetryEnabled = isTelemetryEnabledDefault
        CodeWhispererCodeCoverageTracker.getInstancesMap().clear()
    }

    companion object {
        // TODO: move this to util and tweak it to show what telemetry field not matching assertions
        fun assertEventsContainsFieldsAndCount(
            events: Collection<MetricEvent>,
            name: String,
            count: Int,
            vararg keyValues: Pair<String, Any?>,
            atLeast: Boolean = false
        ) {
            assertThat(events).filteredOn { event ->
                event.data.any {
                    it.name == name && isThisMapContains(it.metadata, *keyValues)
                }
            }.apply {
                if (atLeast) {
                    hasSizeGreaterThanOrEqualTo(count)
                } else {
                    hasSize(count)
                }
            }
        }

        private fun isThisMapContains(map: Map<String, String>, vararg keyValues: Pair<String, Any?>): Boolean {
            keyValues.forEach { pair ->
                val fieldName = pair.first
                val expectedValue = pair.second
                expectedValue?.let {
                    if (it.toString() != map[fieldName]) {
                        return false
                    }
                }
            }
            return true
        }
    }
}