// 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.testFramework.DisposableRule
import com.intellij.testFramework.RuleChain
import com.intellij.testFramework.replaceService
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.After
import org.junit.Before
import org.junit.Rule
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doAnswer
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import org.mockito.kotlin.timeout
import org.mockito.kotlin.verify
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.paginators.GenerateCompletionsIterable
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.codeWhispererRecommendationActionId
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.CodeWhispererTestUtil.pythonTestLeftContext
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testValidAccessToken
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
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.CodeWhispererExploreActionState
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreStateType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererRecommendationManager
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.settings.CodeWhispererConfiguration
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurationType
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.resources.message

// TODO: restructure testbase, too bulky and hard to debug
open class CodeWhispererTestBase {
    var projectRule = PythonCodeInsightTestFixtureRule()
    val mockClientManagerRule = MockClientManagerRule()
    val mockCredentialRule = MockCredentialManagerRule()
    val disposableRule = DisposableRule()

    @Rule
    @JvmField
    val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, disposableRule)

    protected lateinit var mockClient: CodeWhispererRuntimeClient

    protected lateinit var popupManagerSpy: CodeWhispererPopupManager
    protected lateinit var clientAdaptorSpy: CodeWhispererClientAdaptor
    protected lateinit var invocationStatusSpy: CodeWhispererInvocationStatus
    internal lateinit var stateManager: CodeWhispererExplorerActionManager
    protected lateinit var recommendationManager: CodeWhispererRecommendationManager
    protected lateinit var codewhispererService: CodeWhispererService
    protected lateinit var editorManager: CodeWhispererEditorManager
    protected lateinit var settingsManager: CodeWhispererSettings
    private lateinit var originalExplorerActionState: CodeWhispererExploreActionState
    private lateinit var originalSettings: CodeWhispererConfiguration
    private lateinit var userGroupSettings: CodeWhispererUserGroupSettings

    @Before
    open fun setUp() {
        mockClient = mockClientManagerRule.create()
        val requestCaptor = argumentCaptor<GenerateCompletionsRequest>()
        mockClient.stub {
            on {
                mockClient.generateCompletionsPaginator(requestCaptor.capture())
            } doAnswer {
                GenerateCompletionsIterable(mockClient, requestCaptor.lastValue)
            }
        }
        mockClient.stub {
            on {
                mockClient.generateCompletions(any<GenerateCompletionsRequest>())
            } doAnswer {
                pythonResponse
            }
        }

        popupManagerSpy = spy(CodeWhispererPopupManager.getInstance())
        popupManagerSpy.reset()
        doNothing().`when`(popupManagerSpy).showPopup(any(), any(), any(), any(), any())
        ApplicationManager.getApplication().replaceService(CodeWhispererPopupManager::class.java, popupManagerSpy, disposableRule.disposable)

        invocationStatusSpy = spy(CodeWhispererInvocationStatus.getInstance())
        invocationStatusSpy.stub {
            on {
                hasEnoughDelayToShowCodeWhisperer()
            } doAnswer {
                true
            }
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererInvocationStatus::class.java, invocationStatusSpy, disposableRule.disposable)

        stateManager = spy(CodeWhispererExplorerActionManager.getInstance())
        recommendationManager = CodeWhispererRecommendationManager.getInstance()
        codewhispererService = CodeWhispererService.getInstance()
        editorManager = CodeWhispererEditorManager.getInstance()
        settingsManager = CodeWhispererSettings.getInstance()

        setFileContext(pythonFileName, pythonTestLeftContext, "")

        originalExplorerActionState = stateManager.state
        originalSettings = settingsManager.state
        stateManager.loadState(
            CodeWhispererExploreActionState().apply {
                CodeWhispererExploreStateType.values().forEach {
                    value[it] = true
                }
                token = testValidAccessToken
            }
        )
        settingsManager.loadState(
            CodeWhispererConfiguration().apply {
                value[CodeWhispererConfigurationType.IsIncludeCodeWithReference] = true
            }
        )

        clientAdaptorSpy = spy(CodeWhispererClientAdaptor.getInstance(projectRule.project))
        projectRule.project.replaceService(CodeWhispererClientAdaptor::class.java, clientAdaptorSpy, disposableRule.disposable)
        ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, stateManager, disposableRule.disposable)
        stateManager.setAutoEnabled(false)

        userGroupSettings = mock { on { getUserGroup() } doReturn CodeWhispererUserGroup.Control }
        ApplicationManager.getApplication().replaceService(CodeWhispererUserGroupSettings::class.java, userGroupSettings, disposableRule.disposable)
    }

    @After
    open fun tearDown() {
        stateManager.loadState(originalExplorerActionState)
        settingsManager.loadState(originalSettings)
        popupManagerSpy.reset()
    }

    fun withCodeWhispererServiceInvokedAndWait(runnable: (InvocationContext) -> Unit) {
        val statesCaptor = argumentCaptor<InvocationContext>()
        invokeCodeWhispererService()
        verify(popupManagerSpy, timeout(5000).atLeastOnce())
            .showPopup(statesCaptor.capture(), any(), any(), any(), any())
        val states = statesCaptor.lastValue

        runInEdtAndWait {
            try {
                runnable(states)
            } finally {
                CodeWhispererPopupManager.getInstance().closePopup(states.popup)
            }
        }

        // To make sure the previous runnable on EDT thread is complete before the test proceeds
        runInEdtAndWait {}
    }

    fun invokeCodeWhispererService() {
        runInEdtAndWait {
            projectRule.fixture.performEditorAction(codeWhispererRecommendationActionId)
        }
        while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
            Thread.sleep(10)
        }
    }

    fun checkRecommendationInfoLabelText(selected: Int, total: Int) {
        // should show "[selected] of [total]"
        val actual = popupManagerSpy.popupComponents.recommendationInfoLabel.text
        val expected = message("codewhisperer.popup.recommendation_info", selected, total, POPUP_DIM_HEX)
        assertThat(actual).isEqualTo(expected)
    }

    fun setFileContext(filename: String, leftContext: String, rightContext: String) {
        projectRule.fixture.configureByText(filename, leftContext + rightContext)
        runInEdtAndWait {
            projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(leftContext.length)
        }
    }

    fun mockCodeWhispererEnabledStatus(enabled: Boolean) {
        stateManager.stub {
            onGeneric {
                checkActiveCodeWhispererConnectionType(any())
            } doAnswer {
                if (enabled) CodeWhispererLoginType.Sono else CodeWhispererLoginType.Logout
            }
        }
    }
}


// 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.gradle.internal.impldep.com.amazonaws.ResponseMetadata.AWS_REQUEST_ID
import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata
import software.amazon.awssdk.awscore.exception.AwsErrorDetails
import software.amazon.awssdk.http.SdkHttpResponse
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
import software.amazon.awssdk.services.codewhispererruntime.model.FileContext
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.ProgrammingLanguage
import software.amazon.awssdk.services.codewhispererruntime.model.RecommendationsWithReferencesPreference
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
import software.amazon.awssdk.services.codewhispererruntime.model.Span
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import kotlin.random.Random

object CodeWhispererTestUtil {
    const val testSessionId = "test_codewhisperer_session_id"
    const val testRequestId = "test_aws_request_id"
    const val testRequestIdForCodeWhispererException = "test_request_id_for_codewhispererException"
    const val codeWhispererRecommendationActionId = "CodeWhispererRecommendationAction"
    const val codeWhispererCodeScanActionId = "codewhisperer.toolbar.security.scan"
    const val testValidAccessToken = "test_valid_access_token"
    const val testNextToken = "test_next_token"
    private val testReferenceInfoPair = listOf(
        Pair("MIT", "testRepo1"),
        Pair("Apache-2.0", "testRepo2"),
        Pair("BSD-4-Clause", "testRepo3")
    )
    val metadata: DefaultAwsResponseMetadata = DefaultAwsResponseMetadata.create(
        mapOf(AWS_REQUEST_ID to testRequestId)
    )
    val sdkHttpResponse = SdkHttpResponse.builder().headers(
        mapOf(CodeWhispererService.KET_SESSION_ID to listOf(testSessionId))
    ).build()
    private val errorDetail = AwsErrorDetails.builder()
        .errorCode("123")
        .errorMessage("something went wrong")
        .sdkHttpResponse(sdkHttpResponse)
        .build()
    val testCodeWhispererException = CodeWhispererRuntimeException.builder()
        .requestId(testRequestIdForCodeWhispererException)
        .awsErrorDetails(errorDetail)
        .build() as CodeWhispererRuntimeException

    val pythonRequest: GenerateCompletionsRequest = GenerateCompletionsRequest.builder()
        .fileContext(
            FileContext.builder()
                .filename("test.py")
                .programmingLanguage(
                    ProgrammingLanguage.builder()
                        .languageName("python")
                        .build()
                )
                .build()
        )
        .nextToken("")
        .referenceTrackerConfiguration { it.recommendationsWithReferences(RecommendationsWithReferencesPreference.ALLOW) }
        .maxResults(5)
        .build()

    val pythonResponse: GenerateCompletionsResponse = GenerateCompletionsResponse.builder()
        .completions(
            generateMockCompletionDetail("(x, y):\n    return x + y"),
            generateMockCompletionDetail("(a, b):\n    return a + b"),
            generateMockCompletionDetail("test recommendation 3"),
            generateMockCompletionDetail("test recommendation 4"),
            generateMockCompletionDetail("test recommendation 5")
        )
        .nextToken("")
        .responseMetadata(metadata)
        .sdkHttpResponse(sdkHttpResponse)
        .build() as GenerateCompletionsResponse
    val pythonResponseWithNonEmptyToken = pythonResponseWithToken(testNextToken)
    val javaResponse: GenerateCompletionsResponse = GenerateCompletionsResponse.builder()
        .completions(
            generateMockCompletionDetail("(x, y) {\n        return x + y\n    }"),
            generateMockCompletionDetail("(a, b) {\n        return a + b\n    }"),
            generateMockCompletionDetail("test recommendation 3"),
            generateMockCompletionDetail("test recommendation 4"),
            generateMockCompletionDetail("test recommendation 5")
        )
        .nextToken("")
        .responseMetadata(metadata)
        .sdkHttpResponse(sdkHttpResponse)
        .build() as GenerateCompletionsResponse
    val emptyListResponse: GenerateCompletionsResponse = GenerateCompletionsResponse.builder()
        .completions(listOf())
        .nextToken("")
        .responseMetadata(metadata)
        .sdkHttpResponse(sdkHttpResponse)
        .build() as GenerateCompletionsResponse
    val listOfEmptyRecommendationResponse: GenerateCompletionsResponse = GenerateCompletionsResponse.builder()
        .completions(
            generateMockCompletionDetail(""),
            generateMockCompletionDetail(""),
            generateMockCompletionDetail(""),
        )
        .nextToken("")
        .responseMetadata(metadata)
        .sdkHttpResponse(sdkHttpResponse)
        .build() as GenerateCompletionsResponse
    val listOfMixedEmptyAndNonEmptyRecommendationResponse: GenerateCompletionsResponse = GenerateCompletionsResponse.builder()
        .completions(
            generateMockCompletionDetail(""),
            generateMockCompletionDetail("test recommendation 3"),
            generateMockCompletionDetail(""),
            generateMockCompletionDetail("test recommendation 4"),
            generateMockCompletionDetail("test recommendation 5")
        )
        .nextToken("")
        .responseMetadata(metadata)
        .sdkHttpResponse(sdkHttpResponse)
        .build() as GenerateCompletionsResponse

    const val pythonFileName = "test.py"
    const val javaFileName = "test.java"
    const val cppFileName = "test.cpp"
    const val jsFileName = "test.js"
    const val pythonTestLeftContext = "def addTwoNumbers"
    const val cppTestLeftContext = "int addTwoNumbers"
    const val javaTestContext = "public class Test {\n    public static void main\n}"

    internal fun pythonResponseWithToken(token: String): GenerateCompletionsResponse =
        pythonResponse.toBuilder().nextToken(token).build()

    internal fun generateMockCompletionDetail(content: String): Completion {
        val referenceInfo = getReferenceInfo()
        return Completion.builder().content(content)
            .references(
                generateMockReferences(referenceInfo.first, referenceInfo.second, 0, content.length)
            )
            .build()
    }

    internal fun getReferenceInfo() = testReferenceInfoPair[Random.nextInt(testReferenceInfoPair.size)]

    internal fun generateMockCompletionDetail(
        content: String,
        licenseName: String,
        repository: String,
        start: Int,
        end: Int
    ): Completion =
        Completion.builder()
            .content(content)
            .references(generateMockReferences(licenseName, repository, start, end))
            .build()

    private fun generateMockReferences(licenseName: String, repository: String, start: Int, end: Int) =
        Reference.builder()
            .licenseName(licenseName)
            .repository(repository)
            .recommendationContentSpan(
                Span.builder()
                    .start(start)
                    .end(end)
                    .build()
            )
            .build()
}


// 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.testFramework.runInEdtAndWait
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.pythonTestLeftContext

class CodeWhispererTypeaheadTest : CodeWhispererTestBase() {

    @Test
    fun `test typing typeahead should update typeahead state when no right context`() {
        testTypingTypeaheadMatchingRecommendationShouldMatchRightContext("")
    }

    @Test
    fun `test typing typeahead should update typeahead state when there is right context with extra opening bracket`() {
        val testRightContext = "() {\n"
        testTypingTypeaheadMatchingRecommendationShouldMatchRightContext(testRightContext)
    }

    @Test
    fun `test typing typeahead should update typeahead state when there is right context with extra closing bracket`() {
        val testRightContext = "() }"
        testTypingTypeaheadMatchingRecommendationShouldMatchRightContext(testRightContext)
    }

    @Test
    fun `test typing typeahead should update typeahead state when there is right context with trailing spaces`() {
        val testRightContext = "() {\n    }      "
        testTypingTypeaheadMatchingRecommendationShouldMatchRightContext(testRightContext)
    }

    @Test
    fun `test typing blank typeahead should correctly update typeahead state`() {
        val testTypeaheadOriginal = "     "
        testTypingTypeaheadWithLeadingSpaceShouldMatchTypeaheadStateCorrectly(testTypeaheadOriginal, 5, 1)
    }

    @Test
    fun `test typing typeahead with leading spaces and matching suffix should correctly update typeahead state`() {
        val testTypeaheadOriginal = "   test"
        testTypingTypeaheadWithLeadingSpaceShouldMatchTypeaheadStateCorrectly(testTypeaheadOriginal, 3, 3)
    }

    private fun testTypingTypeaheadWithLeadingSpaceShouldMatchTypeaheadStateCorrectly(
        expectedTypeaheadOriginal: String,
        expectedNumOfValidRecommendation: Int,
        expectedSelectedAfterBackspace: Int
    ) {
        withCodeWhispererServiceInvokedAndWait { states ->
            val editor = projectRule.fixture.editor
            val startOffset = editor.caretModel.offset
            expectedTypeaheadOriginal.forEach { char ->
                projectRule.fixture.type(char)
                val caretOffset = editor.caretModel.offset
                val actualTypeaheadOriginal = editor.document.charsSequence.subSequence(startOffset, caretOffset).toString()
                val actualTypeahead = actualTypeaheadOriginal.trimStart()
                assertThat(popupManagerSpy.sessionContext.typeaheadOriginal).isEqualTo(actualTypeaheadOriginal)
                assertThat(popupManagerSpy.sessionContext.typeahead).isEqualTo(actualTypeahead)
                assertThat(states.popup.isDisposed).isFalse
            }
            checkRecommendationInfoLabelText(1, expectedNumOfValidRecommendation)

            // Backspacing for the same amount of times
            expectedTypeaheadOriginal.forEach { _ ->
                projectRule.fixture.type('\b')
                val caretOffset = editor.caretModel.offset
                val actualTypeaheadOriginal = editor.document.charsSequence.subSequence(startOffset, caretOffset).toString()
                val actualTypeahead = actualTypeaheadOriginal.trimStart()
                assertThat(popupManagerSpy.sessionContext.typeaheadOriginal).isEqualTo(actualTypeaheadOriginal)
                assertThat(popupManagerSpy.sessionContext.typeahead).isEqualTo(actualTypeahead)
                assertThat(states.popup.isDisposed).isFalse
            }
            checkRecommendationInfoLabelText(expectedSelectedAfterBackspace, 5)
        }
    }

    private fun testTypingTypeaheadMatchingRecommendationShouldMatchRightContext(rightContext: String) {
        projectRule.fixture.configureByText(pythonFileName, pythonTestLeftContext + rightContext)
        runInEdtAndWait {
            projectRule.fixture.editor.caretModel.moveToOffset(pythonTestLeftContext.length)
        }
        withCodeWhispererServiceInvokedAndWait { states ->
            val recommendation = states.recommendationContext.details[0].reformatted.content()
            val editor = projectRule.fixture.editor
            val startOffset = editor.caretModel.offset
            recommendation.forEachIndexed { index, char ->
                if (index < editor.caretModel.offset - startOffset) return@forEachIndexed
                projectRule.fixture.type(char)
                val caretOffset = editor.caretModel.offset
                val typeahead = editor.document.charsSequence.subSequence(startOffset, caretOffset).toString()
                assertThat(popupManagerSpy.sessionContext.typeahead).isEqualTo(typeahead)
            }
        }
    }
}


// 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.codeInsight.codeVision.ui.visibleAreaChanged
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_DELETE_LINE
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_DELETE_TO_WORD_START
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_LEFT_WITH_SELECTION
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_RIGHT_WITH_SELECTION
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_LINE_END_WITH_SELECTION
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_LINE_START_WITH_SELECTION
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_SELECT_WORD_AT_CARET
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TEXT_END_WITH_SELECTION
import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TEXT_START_WITH_SELECTION
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.editor.event.VisibleAreaEvent
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.testFramework.runInEdtAndWait
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.mockito.Mockito.mock
import org.mockito.Mockito.times
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import org.mockito.kotlin.timeout
import org.mockito.kotlin.verify
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
import java.awt.Rectangle

class CodeWhispererUserActionsTest : CodeWhispererTestBase() {

    @Before
    override fun setUp() {
        super.setUp()
        WriteCommandAction.runWriteCommandAction(projectRule.project) {
            projectRule.fixture.editor.document.insertString(projectRule.fixture.editor.caretModel.offset, "test")
        }
    }

    @Test
    fun `test delete to word start should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_DELETE_TO_WORD_START)
    }

    @Test
    fun `test delete line should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_DELETE_LINE)
    }

    @Test
    fun `test text start with selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_TEXT_START_WITH_SELECTION)
    }

    @Test
    fun `test text end with selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_TEXT_END_WITH_SELECTION)
    }

    @Test
    fun `test move line start with selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_MOVE_LINE_START_WITH_SELECTION)
    }

    @Test
    fun `test move line end selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_MOVE_LINE_END_WITH_SELECTION)
    }

    @Test
    fun `test select word at caret should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_SELECT_WORD_AT_CARET)
    }

    @Test
    fun `test move caret left with selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_MOVE_CARET_LEFT_WITH_SELECTION)
    }

    @Test
    fun `test move caret right with selection should cancel popup`() {
        testUserActionsShouldCancelPopup(ACTION_EDITOR_MOVE_CARET_RIGHT_WITH_SELECTION)
    }

    private fun testUserActionsShouldCancelPopup(actionId: String) {
        runInEdtAndWait {
            projectRule.fixture.editor.caretModel.moveToOffset(pythonTestLeftContext.length)
        }
        withCodeWhispererServiceInvokedAndWait {
            projectRule.fixture.performEditorAction(actionId)
            verify(popupManagerSpy, timeout(5000)).cancelPopup(any())
        }
    }

    @Test
    fun `test hitting enter after non-whitespace characters should trigger CodeWhisperer`() {
        testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer(pythonTestLeftContext, 1)
    }

    @Test
    fun `test hitting enter after whitespace characters should trigger CodeWhisperer`() {
        testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer("$pythonTestLeftContext ", 1)
        testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer("$pythonTestLeftContext\t", 2)
        testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer("$pythonTestLeftContext\n", 3)
    }

    @Test
    fun `test hitting enter inside braces in Java file should auto-trigger CodeWhisperer and keep the formatting correct`() {
        val testLeftContext = "public class Test {\n    public static void main() {"
        val testRightContext = "}\n}"
        setFileContext(javaFileName, testLeftContext, testRightContext)
        CodeWhispererExplorerActionManager.getInstance().setAutoEnabled(true)
        projectRule.fixture.type('\n')
        val expectedFileContext = "$testLeftContext\n        \n    $testRightContext"
        assertThat(projectRule.fixture.editor.document.text).isEqualTo(expectedFileContext)
        val popupCaptor = argumentCaptor<JBPopup>()
        verify(popupManagerSpy, timeout(5000))
            .showPopup(any(), any(), popupCaptor.capture(), any(), any())
        runInEdtAndWait {
            popupManagerSpy.closePopup(popupCaptor.lastValue)
        }
    }

    @Test
    fun `test CodeWhispererScrollListener re-render popup when visibleArea gets changed when popup is active`() {
        // Arrange
        val oldRect = Rectangle(0, 0, 20, 10)
        val newRect = Rectangle(0, 10, 20, 10)
        val event = mock<VisibleAreaEvent> {
            on { this.oldRectangle } doReturn oldRect
            on { this.newRectangle } doReturn newRect
        }
        withCodeWhispererServiceInvokedAndWait { states ->
            CodeWhispererInvocationStatus.getInstance().setPopupActive(true)
            val listener = CodeWhispererScrollListener(states)
            listener.visibleAreaChanged(event)
            verify(popupManagerSpy, times(2)).showPopup(any(), any(), any(), any(), any())
        }
    }

    private fun testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer(prompt: String, times: Int) {
        CodeWhispererExplorerActionManager.getInstance().setAutoEnabled(true)
        setFileContext(pythonFileName, prompt, "")
        projectRule.fixture.type('\n')
        val popupCaptor = argumentCaptor<JBPopup>()
        verify(popupManagerSpy, timeout(5000).atLeast(times))
            .showPopup(any(), any(), popupCaptor.capture(), any(), any())
        runInEdtAndWait {
            popupManagerSpy.closePopup(popupCaptor.lastValue)
        }
    }
}


// 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.testFramework.ApplicationRule
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import org.mockito.kotlin.spy
import org.mockito.kotlin.verify
import software.aws.toolkits.jetbrains.AwsToolkit
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererExpThresholdGroup
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.CodeWhispererUserGroupSettings.Companion.EXP_THRESHOLD_KEY
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings.Companion.USER_GROUP_KEY
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupStates
import kotlin.reflect.KMutableProperty
import kotlin.reflect.full.memberProperties
import kotlin.reflect.jvm.isAccessible

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

    private lateinit var sut: CodeWhispererUserGroupSettings

    @Before
    fun setup() {
        sut = CodeWhispererUserGroupSettings()
    }

    @Test
    fun `getGroup should return the current states if the group does exist`() {
        val settingsField = sut::class.memberProperties.find { it.name == "settings" }
        assertThat(settingsField).isNotNull

        settingsField?.let {
            it.isAccessible = true

            @Suppress("UNCHECKED_CAST")
            val settings = it.call(sut) as MutableMap<String, String>

            settings[USER_GROUP_KEY] = CodeWhispererUserGroup.Control.name
            settings[EXP_THRESHOLD_KEY] = CodeWhispererExpThresholdGroup.Exp.name

            assertThat(sut.getGroup<CodeWhispererUserGroup>()).isEqualTo(CodeWhispererUserGroup.Control)
            assertThat(sut.getGroup<CodeWhispererExpThresholdGroup>()).isEqualTo(CodeWhispererExpThresholdGroup.Exp)
        }
    }

    @Test
    fun `getGroup will return null if the group does not exist - mix case`() {
        val settingsField = sut::class.memberProperties.find { it.name == "settings" }
        assertThat(settingsField).isNotNull

        settingsField?.let {
            it.isAccessible = true

            @Suppress("UNCHECKED_CAST")
            val settings = it.call(sut) as MutableMap<String, String>

            settings[USER_GROUP_KEY] = CodeWhispererUserGroup.CrossFile.name
            settings[EXP_THRESHOLD_KEY] = "group that doesn't exist"

            assertThat(sut.getGroup<CodeWhispererUserGroup>()).isEqualTo(CodeWhispererUserGroup.CrossFile)
            assertThat(sut.getGroup<CodeWhispererExpThresholdGroup>()).isNull()
        }
    }

    @Test
    fun `getGroup will return null if the group does not exist - both not exist`() {
        val settingsField = sut::class.memberProperties.find { it.name == "settings" }
        assertThat(settingsField).isNotNull

        settingsField?.let {
            it.isAccessible = true

            @Suppress("UNCHECKED_CAST")
            val settings = it.call(sut) as MutableMap<String, String>

            settings[USER_GROUP_KEY] = "group that doesn't exist"
            settings[EXP_THRESHOLD_KEY] = "another group that doesn't exist"

            assertThat(sut.getGroup<CodeWhispererUserGroup>()).isNull()
            assertThat(sut.getGroup<CodeWhispererExpThresholdGroup>()).isNull()
        }
    }

    @Test
    fun `getUserGroup will assign a group if group is not set yet`() {
        val settingsField = sut::class.memberProperties.find { it.name == "settings" }
        assertThat(settingsField).isNotNull
        settingsField?.let {
            it.isAccessible = true

            @Suppress("UNCHECKED_CAST")
            val settings = it.call(sut) as Map<String, String>
            assertThat(settings).hasSize(0)

            val userGroupAssigned = sut.getUserGroup()
            assertThat(userGroupAssigned).isEqualTo(sut.getUserGroup())

            // should be the same no matter how many times we invoke
            assertThat(userGroupAssigned).isEqualTo(sut.getUserGroup())
            assertThat(userGroupAssigned).isEqualTo(sut.getUserGroup())
        }
    }

    @Test
    fun `getUserGroup will return the previously set group if version doesn't change`() {
        val settingsField = sut::class.memberProperties.find { it.name == "settings" }
        assertThat(settingsField).isNotNull

        // if version differs, will re-assign the group
        val oldVersionStoreInAwsXml = AwsToolkit.PLUGIN_VERSION

        settingsField?.let {
            // set up CodeWhispererUserGroupSettings.settings field
            it.isAccessible = true
            @Suppress("UNCHECKED_CAST")
            val settings = it.getter.call(sut) as MutableMap<String, String>
            settings[USER_GROUP_KEY] = CodeWhispererUserGroup.Control.name

            // set up CodeWhispererUserGroupSettings.version field
            val versionField = sut::class.memberProperties.find { field -> field.name == "version" }
            assertThat(versionField).isNotNull
            versionField?.let { myVersion ->
                myVersion as KMutableProperty<*>
                myVersion.isAccessible = true
                myVersion.setter.call(sut, oldVersionStoreInAwsXml)
            }

            assertThat(sut.getUserGroup()).isEqualTo(CodeWhispererUserGroup.Control)

            // will not change the value no matter how many times we invoke getUserGroup
            assertThat(sut.getUserGroup()).isEqualTo(CodeWhispererUserGroup.Control)
        }
    }

    @Test
    fun `getUserGroup will re-assign the group if version changes`() {
        sut = spy(sut)
        sut.loadState(
            CodeWhispererUserGroupStates(
                version = "old version",
                settings = mapOf(USER_GROUP_KEY to CodeWhispererUserGroup.Control.name)
            )
        )

        sut.getUserGroup()

        verify(sut).determineUserGroup()
        assertThat(sut.getVersion()).isEqualTo(AwsToolkit.PLUGIN_VERSION)
    }

    @Test
    fun loadState() {
        // version must use AwsToolkit.PLUGIN_VERSION, otherwise group will be assigned
        val state = CodeWhispererUserGroupStates(
            version = AwsToolkit.PLUGIN_VERSION,
            settings = mapOf(USER_GROUP_KEY to CodeWhispererUserGroup.CrossFile.name)
        )

        sut.loadState(state)

        assertThat(sut.getUserGroup()).isEqualTo(CodeWhispererUserGroup.CrossFile)
        assertThat(sut.getVersion()).isEqualTo(AwsToolkit.PLUGIN_VERSION)
    }
}


// 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.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import com.intellij.testFramework.replaceService
import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doAnswer
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
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.CodeWhispererService

class CodeWhispererUserInputTest : CodeWhispererTestBase() {

    @Test
    fun `test no user input should show all recommendations`() {
        addUserInputAfterInvocation("")

        withCodeWhispererServiceInvokedAndWait { states ->
            val actualRecommendations = states.recommendationContext.details.map {
                it.recommendation.content()
            }
            assertThat(actualRecommendations).isEqualTo(pythonResponse.completions().map { it.content() })
        }
    }

    @Test
    fun `test have user input should show that recommendations prefix-matching user input are valid`() {
        val userInput = "test"
        addUserInputAfterInvocation(userInput)

        val expectedRecommendations = pythonResponse.completions().map { it.content() }

        withCodeWhispererServiceInvokedAndWait { states ->
            val actualRecommendations = states.recommendationContext.details.map { it.recommendation.content() }
            assertThat(actualRecommendations).isEqualTo(expectedRecommendations)
            states.recommendationContext.details.forEachIndexed { index, context ->
                val expectedDiscarded = !pythonResponse.completions()[index].content().startsWith(userInput)
                val actualDiscarded = context.isDiscarded
                assertThat(actualDiscarded).isEqualTo(expectedDiscarded)
            }
        }
    }

    @Test
    fun `test have user input and typeahead should show that recommendations prefix-matching user input + typeahead are valid`() {
        val userInput = "test"
        addUserInputAfterInvocation(userInput)

        val typeahead = " recommendation"

        withCodeWhispererServiceInvokedAndWait { states ->
            projectRule.fixture.type(typeahead)
            assertThat(popupManagerSpy.sessionContext.typeahead).isEqualTo(typeahead)
            states.recommendationContext.details.forEachIndexed { index, actualContext ->
                val actualDiscarded = actualContext.isDiscarded
                val expectedDiscarded = !pythonResponse.completions()[index].content().startsWith(userInput + typeahead)
                assertThat(actualDiscarded).isEqualTo(expectedDiscarded)
            }
        }
    }

    @Test
    fun `test have blank user input should show that all recommendations are valid`() {
        val blankUserInput = "    "
        addUserInputAfterInvocation(blankUserInput)
        val userInput = blankUserInput.trimStart()

        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.userInputSinceInvocation).isEqualTo(userInput)
            states.recommendationContext.details.forEachIndexed { _, actualContext ->
                assertThat(actualContext.isDiscarded).isEqualTo(false)
            }
        }
    }

    @Test
    fun `test have user input with leading spaces and matching suffix should show recommendations prefix-matching suffix are valid`() {
        val userInputWithLeadingSpaces = "   test"
        addUserInputAfterInvocation(userInputWithLeadingSpaces)
        val userInput = userInputWithLeadingSpaces.trimStart()

        withCodeWhispererServiceInvokedAndWait { states ->
            assertThat(states.recommendationContext.userInputSinceInvocation).isEqualTo(userInput)
            states.recommendationContext.details.forEachIndexed { index, actualContext ->
                val actualDiscarded = actualContext.isDiscarded
                val expectedDiscarded = !pythonResponse.completions()[index].content().startsWith(userInput)
                assertThat(actualDiscarded).isEqualTo(expectedDiscarded)
            }
        }
    }

    private fun addUserInputAfterInvocation(userInput: String) {
        val codewhispererServiceSpy = spy(codewhispererService)
        val triggerTypeCaptor = argumentCaptor<TriggerTypeInfo>()
        val editorCaptor = argumentCaptor<Editor>()
        val projectCaptor = argumentCaptor<Project>()
        val psiFileCaptor = argumentCaptor<PsiFile>()
        val latencyContextCaptor = argumentCaptor<LatencyContext>()
        codewhispererServiceSpy.stub {
            onGeneric {
                getRequestContext(
                    triggerTypeCaptor.capture(),
                    editorCaptor.capture(),
                    projectCaptor.capture(),
                    psiFileCaptor.capture(),
                    latencyContextCaptor.capture()
                )
            }.doAnswer {
                val requestContext = codewhispererServiceSpy.getRequestContext(
                    triggerTypeCaptor.firstValue,
                    editorCaptor.firstValue,
                    projectCaptor.firstValue,
                    psiFileCaptor.firstValue,
                    latencyContextCaptor.firstValue
                )
                projectRule.fixture.type(userInput)
                requestContext
            }.thenCallRealMethod()
        }
        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
    }
}


// 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.testFramework.fixtures.CodeInsightTestFixture
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Rule
import org.junit.Test
import software.aws.toolkits.jetbrains.services.codewhisperer.util.toCodeChunk
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule

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

    lateinit var fixture: CodeInsightTestFixture

    @Before
    fun setup() {
        fixture = projectRule.fixture
    }

    @Test
    fun `toCodeChunk case_1`() {
        val psiFile = fixture.configureByText(
            "Sample.java",
            """public class Main {
            |    public static void main() {
            |    }
            |}
            """.trimMargin()
        )

        val result = runBlocking {
            psiFile.virtualFile.toCodeChunk("fake/path")
        }.toList()

        assertThat(result).hasSize(2)

        assertThat(result[0].content).isEqualTo(
            """public class Main {
                |    public static void main() {
                |    }
            """.trimMargin()
        )
        assertThat(result[1].content).isEqualTo(
            """public class Main {
            |    public static void main() {
            |    }
            |}
            """.trimMargin()
        )
    }

    @Test
    fun `toCodeChunk case_2`() {
        val psiFile = fixture.configureByText("Sample.java", codeSample33Lines)

        val result = runBlocking {
            psiFile.virtualFile.toCodeChunk("fake/path")
        }.toList()

        assertThat(result).hasSize(5)

        // 0th
        assertThat(result[0].content).isEqualTo(
            """public int runBinarySearchRecursively(int[] sortedArray, int key, int low, int high) {
                |    int middle = low  + ((high - low) / 2);
            """.trimMargin()
        )
        assertThat(result[0].path).isEqualTo("fake/path")
        assertThat(result[0].nextChunk).isEqualTo(result[1].content)

        // 1st
        assertThat(result[1].content).isEqualTo(
            """|public int runBinarySearchRecursively(int[] sortedArray, int key, int low, int high) {
                    |    int middle = low  + ((high - low) / 2);
                    |
                    |    if (high < low) {
                    |        return -1;
                    |    }
                    |
                    |    if (key == sortedArray[middle]) {
                    |        return middle;
                    |    } else if (key < sortedArray[middle]) {
            """.trimMargin()
        )
        assertThat(result[1].path).isEqualTo("fake/path")
        assertThat(result[1].nextChunk).isEqualTo(result[2].content)

        // 2nd
        assertThat(result[2].content).isEqualTo(
            """|        return runBinarySearchRecursively(sortedArray, key, low, middle - 1);
               |    } else {
               |        return runBinarySearchRecursively(sortedArray, key, middle + 1, high);
               |    }
               |}
               |
               |public int runBinarySearchIteratively(int[] sortedArray, int key, int low, int high) {
               |    int index = Integer.MAX_VALUE;
               |
               |    while (low <= high) {
            """.trimMargin()
        )
        assertThat(result[2].path).isEqualTo("fake/path")
        assertThat(result[2].nextChunk).isEqualTo(result[3].content)

        // 3rd
        assertThat(result[3].content).isEqualTo(
            """|        int mid = low  + ((high - low) / 2);
       |        if (sortedArray[mid] < key) {
       |            low = mid + 1;
       |        } else if (sortedArray[mid] > key) {
       |            high = mid - 1;
       |        } else if (sortedArray[mid] == key) {
       |            index = mid;
       |            break;
       |        }
       |     }
            """.trimMargin()
        )
        assertThat(result[3].path).isEqualTo("fake/path")
        assertThat(result[3].nextChunk).isEqualTo(result[4].content)

        // 4th
        assertThat(result[4].content).isEqualTo(
            """|
               |    return index;
               |}
            """.trimMargin()
        )
        assertThat(result[4].path).isEqualTo("fake/path")
        assertThat(result[4].nextChunk).isEqualTo(result[4].content)
    }
}

private val codeSample33Lines =
    """public int runBinarySearchRecursively(int[] sortedArray, int key, int low, int high) {
       |    int middle = low  + ((high - low) / 2);
       |
       |    if (high < low) {
       |        return -1;
       |    }
       |
       |    if (key == sortedArray[middle]) {
       |        return middle;
       |    } else if (key < sortedArray[middle]) {
       |        return runBinarySearchRecursively(sortedArray, key, low, middle - 1);
       |    } else {
       |        return runBinarySearchRecursively(sortedArray, key, middle + 1, high);
       |    }
       |}
       |
       |public int runBinarySearchIteratively(int[] sortedArray, int key, int low, int high) {
       |    int index = Integer.MAX_VALUE;
       |
       |    while (low <= high) {
       |        int mid = low  + ((high - low) / 2);
       |        if (sortedArray[mid] < key) {
       |            low = mid + 1;
       |        } else if (sortedArray[mid] > key) {
       |            high = mid - 1;
       |        } else if (sortedArray[mid] == key) {
       |            index = mid;
       |            break;
       |        }
       |     }
       |
       |    return index;
       |}
       |
    """.trimMargin()


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

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

import com.intellij.psi.PsiFile
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import org.apache.commons.codec.digest.DigestUtils
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.mockito.internal.verification.Times
import org.mockito.kotlin.any
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.eq
import org.mockito.kotlin.inOrder
import org.mockito.kotlin.isNull
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import org.mockito.kotlin.verify
import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
import software.aws.toolkits.core.utils.WaiterTimeoutException
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.Payload
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PythonCodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_MILLIS_IN_SECOND
import software.aws.toolkits.jetbrains.utils.isInstanceOf
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.File
import java.io.FileInputStream
import java.util.Base64
import kotlin.io.path.relativeTo
import kotlin.test.assertNotNull

class CodeWhispererCodeScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsightTestFixtureRule()) {
    private lateinit var psifile: PsiFile
    private lateinit var file: File
    private lateinit var sessionConfigSpy: PythonCodeScanSessionConfig
    private val payloadContext = PayloadContext(CodewhispererLanguage.Python, 1, 1, 10, listOf(), 600, 200)
    private lateinit var codeScanSessionContext: CodeScanSessionContext
    private lateinit var codeScanSessionSpy: CodeWhispererCodeScanSession

    @Before
    override fun setup() {
        super.setup()
        psifile = projectRule.fixture.addFileToProject(
            "/test.py",
            """import numpy as np
               import from module1 import helper

               def add(a, b):
                  return a + b

            """.trimMargin()
        )
        file = psifile.virtualFile.toNioPath().toFile()

        sessionConfigSpy = spy(CodeScanSessionConfig.create(psifile.virtualFile, project) as PythonCodeScanSessionConfig)
        setupResponse(psifile.virtualFile.toNioPath().relativeTo(sessionConfigSpy.projectRoot.toNioPath()))

        sessionConfigSpy.stub {
            onGeneric { sessionConfigSpy.createPayload() }.thenReturn(Payload(payloadContext, file))
        }

        // Mock CodeWhispererClient needs to be setup before initializing CodeWhispererCodeScanSession
        codeScanSessionContext = CodeScanSessionContext(project, sessionConfigSpy)
        codeScanSessionSpy = spy(CodeWhispererCodeScanSession(codeScanSessionContext))
        doNothing().`when`(codeScanSessionSpy).uploadArtifactToS3(any(), any(), any(), any(), isNull())

        mockClient.stub {
            onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
        }
    }

    @Test
    fun `test createUploadUrlAndUpload()`() {
        val fileMd5: String = Base64.getEncoder().encodeToString(DigestUtils.md5(FileInputStream(file)))
        codeScanSessionSpy.stub {
            onGeneric { codeScanSessionSpy.createUploadUrl(any(), any()) }
                .thenReturn(fakeCreateUploadUrlResponse)
        }

        codeScanSessionSpy.createUploadUrlAndUpload(file, "artifactType")

        val inOrder = inOrder(codeScanSessionSpy)
        inOrder.verify(codeScanSessionSpy).createUploadUrl(eq(fileMd5), eq("artifactType"))
        inOrder.verify(codeScanSessionSpy).uploadArtifactToS3(
            eq(fakeCreateUploadUrlResponse.uploadUrl()),
            eq(fakeCreateUploadUrlResponse.uploadId()),
            eq(file),
            eq(fileMd5),
            eq(null)
        )
    }

    @Test
    fun `test createUploadUrl()`() {
        val response = codeScanSessionSpy.createUploadUrl("md5", "type")

        argumentCaptor<CreateUploadUrlRequest>().apply {
            verify(mockClient).createUploadUrl(capture())
            assertThat(response.uploadUrl()).isEqualTo(s3endpoint)
            assertThat(response.uploadId()).isEqualTo(UPLOAD_ID)
            assertThat(firstValue.contentMd5()).isEqualTo("md5")
            assertThat(firstValue.artifactTypeAsString()).isEqualTo("type")
        }
    }

    @Test
    fun `test mapToCodeScanIssues`() {
        val recommendations = listOf(
            fakeListCodeScanFindingsResponse.codeScanFindings(),
            getFakeRecommendationsOnNonExistentFile()
        )
        val res = codeScanSessionSpy.mapToCodeScanIssues(recommendations)
        assertThat(res).hasSize(2)
    }

    @Test
    fun `test run() - happypath`() {
        assertNotNull(sessionConfigSpy)
        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Success>()
            assertThat(codeScanResponse.issues).hasSize(2)
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat(codeScanResponse.responseContext.codeScanJobId).isEqualTo("jobId")
        }

        val inOrder = inOrder(codeScanSessionSpy)
        inOrder.verify(codeScanSessionSpy, Times(1)).createUploadUrlAndUpload(eq(file), eq("SourceCode"))
        inOrder.verify(codeScanSessionSpy, Times(1)).createCodeScan(eq(CodewhispererLanguage.Python.toString()))
        inOrder.verify(codeScanSessionSpy, Times(1)).getCodeScan(any())
        inOrder.verify(codeScanSessionSpy, Times(1)).listCodeScanFindings(eq("jobId"))
    }

    @Test
    fun `test run() - createCodeScan failed`() {
        mockClient.stub {
            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponseFailed)
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<CodeWhispererCodeScanException>()
        }
    }

    @Test
    fun `test run() - createCodeScan error`() {
        mockClient.stub {
            onGeneric { createCodeScan(any(), any()) }.thenThrow(CodeWhispererException::class.java)
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<CodeWhispererException>()
        }
    }

    @Test
    fun `test run() - getCodeScan failed`() {
        mockClient.stub {
            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponseFailed)
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<CodeWhispererCodeScanException>()
        }
    }

    @Test
    fun `test run() - getCodeScan pending timeout`() {
        sessionConfigSpy.stub {
            onGeneric { overallJobTimeoutInSeconds() }.thenReturn(5)
        }
        mockClient.stub {
            onGeneric { getCodeScan(any(), any()) }.thenAnswer {
                runBlocking {
                    delay(TIMEOUT)
                }
                fakeGetCodeScanResponsePending
            }
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<WaiterTimeoutException>()
        }
    }

    @Test
    fun `test run() - getCodeScan error`() {
        mockClient.stub {
            onGeneric { getCodeScan(any(), any()) }.thenThrow(CodeWhispererException::class.java)
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<CodeWhispererException>()
        }
    }

    @Test
    fun `test run() - listCodeScanFindings error`() {
        mockClient.stub {
            onGeneric { listCodeScanFindings(any(), any()) }.thenThrow(CodeWhispererException::class.java)
        }

        runBlocking {
            val codeScanResponse = codeScanSessionSpy.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Failure>()
            assertThat(codeScanResponse.responseContext.payloadContext).isEqualTo(payloadContext)
            assertThat((codeScanResponse as CodeScanResponse.Failure).failureReason).isInstanceOf<CodeWhispererException>()
        }
    }

    companion object {
        const val TIMEOUT = 10L * TOTAL_MILLIS_IN_SECOND
    }
}


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

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

import com.github.tomakehurst.wiremock.core.WireMockConfiguration
import com.github.tomakehurst.wiremock.junit.WireMockRule
import com.intellij.analysis.problemsView.toolWindow.ProblemsView
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.wm.RegisterToolWindowTask
import com.intellij.openapi.wm.ToolWindowManager
import com.intellij.testFramework.ApplicationRule
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.replaceService
import com.intellij.util.io.systemIndependentPath
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.gradle.internal.impldep.com.amazonaws.ResponseMetadata
import org.junit.Before
import org.junit.Rule
import org.mockito.kotlin.any
import org.mockito.kotlin.doNothing
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.isNull
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.whenever
import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata
import software.amazon.awssdk.services.codewhisperer.model.CodeScanStatus
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.utils.isInstanceOf
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
import java.nio.file.Path
import kotlin.test.assertNotNull

open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule) {
    @Rule
    @JvmField
    val applicationRule = ApplicationRule()

    @Rule
    @JvmField
    val projectRule: CodeInsightTestFixtureRule = projectRule

    @Rule
    @JvmField
    val disposableRule = DisposableRule()

    @Rule
    @JvmField
    val mockClientManagerRule = MockClientManagerRule()

    @Rule
    @JvmField
    val wireMock = WireMockRule(WireMockConfiguration.wireMockConfig().dynamicPort())

    protected lateinit var mockClient: CodeWhispererClientAdaptor

    internal lateinit var s3endpoint: String

    internal lateinit var fakeCreateUploadUrlResponse: CreateUploadUrlResponse
    internal lateinit var fakeCreateCodeScanResponse: CreateCodeScanResponse
    internal lateinit var fakeCreateCodeScanResponseFailed: CreateCodeScanResponse
    internal lateinit var fakeCreateCodeScanResponsePending: CreateCodeScanResponse
    internal lateinit var fakeListCodeScanFindingsResponse: ListCodeScanFindingsResponse
    internal lateinit var fakeGetCodeScanResponse: GetCodeScanResponse
    internal lateinit var fakeGetCodeScanResponsePending: GetCodeScanResponse
    internal lateinit var fakeGetCodeScanResponseFailed: GetCodeScanResponse

    internal val metadata: DefaultAwsResponseMetadata = DefaultAwsResponseMetadata.create(
        mapOf(ResponseMetadata.AWS_REQUEST_ID to CodeWhispererTestUtil.testRequestId)
    )

    internal lateinit var scanManagerSpy: CodeWhispererCodeScanManager
    internal lateinit var project: Project

    @Before
    open fun setup() {
        project = projectRule.project
        s3endpoint = "http://127.0.0.1:${wireMock.port()}"

        scanManagerSpy = spy(CodeWhispererCodeScanManager.getInstance(project))
        doNothing().whenever(scanManagerSpy).addCodeScanUI(any())

        mockClient = mock<CodeWhispererClientAdaptor>().also {
            project.replaceService(CodeWhispererClientAdaptor::class.java, it, disposableRule.disposable)
        }

        ApplicationManager.getApplication().replaceService(
            CodeWhispererExplorerActionManager::class.java,
            mock {
                on { checkActiveCodeWhispererConnectionType(any()) } doReturn CodeWhispererLoginType.Accountless
            },
            disposableRule.disposable
        )
    }

    private fun setupCodeScanFindings(filePath: Path) = """
        [
            {
                "filePath": "${filePath.systemIndependentPath}",
                "startLine": 1,
                "endLine": 2,
                "title": "test",
                "description": {
                    "text": "global variable",
                    "markdown": "### global variable"
                }
            },
            {
                "filePath": "${filePath.systemIndependentPath}",
                "startLine": 1,
                "endLine": 2,
                "title": "test",
                "description": {
                    "text": "global variable",
                    "markdown": "### global variable"
                }
            }
        ]
    """

    protected fun setupResponse(filePath: Path) {
        fakeCreateUploadUrlResponse = CreateUploadUrlResponse.builder()
            .uploadId(UPLOAD_ID)
            .uploadUrl(s3endpoint)
            .responseMetadata(metadata)
            .build() as CreateUploadUrlResponse

        fakeCreateCodeScanResponse = CreateCodeScanResponse.builder()
            .status(CodeScanStatus.COMPLETED)
            .jobId(JOB_ID)
            .responseMetadata(metadata)
            .build() as CreateCodeScanResponse

        fakeCreateCodeScanResponseFailed = CreateCodeScanResponse.builder()
            .status(CodeScanStatus.FAILED)
            .jobId(JOB_ID)
            .responseMetadata(metadata)
            .build() as CreateCodeScanResponse

        fakeCreateCodeScanResponsePending = CreateCodeScanResponse.builder()
            .status(CodeScanStatus.PENDING)
            .jobId(JOB_ID)
            .responseMetadata(metadata)
            .build() as CreateCodeScanResponse

        fakeListCodeScanFindingsResponse = ListCodeScanFindingsResponse.builder()
            .codeScanFindings(setupCodeScanFindings(filePath))
            .responseMetadata(metadata)
            .build() as ListCodeScanFindingsResponse

        fakeGetCodeScanResponse = GetCodeScanResponse.builder()
            .status(CodeScanStatus.COMPLETED)
            .responseMetadata(metadata)
            .build() as GetCodeScanResponse

        fakeGetCodeScanResponsePending = GetCodeScanResponse.builder()
            .status(CodeScanStatus.PENDING)
            .responseMetadata(metadata)
            .build() as GetCodeScanResponse

        fakeGetCodeScanResponseFailed = GetCodeScanResponse.builder()
            .status(CodeScanStatus.FAILED)
            .responseMetadata(metadata)
            .build() as GetCodeScanResponse
    }

    protected fun getFakeRecommendationsOnNonExistentFile() = """
        [
            {
                "filePath": "non-exist.py",
                "startLine": 1,
                "endLine": 2,
                "title": "test",
                "description": {
                    "text": "global variable",
                    "markdown": "### global variable"
                }
            }
        ]
    """

    internal fun assertE2ERunsSuccessfully(
        sessionConfigSpy: CodeScanSessionConfig,
        project: Project,
        expectedTotalLines: Long,
        expectedTotalFiles: Int,
        expectedTotalSize: Long,
        expectedTotalIssues: Int
    ) {
        val codeScanContext = CodeScanSessionContext(project, sessionConfigSpy)
        val sessionMock = spy(CodeWhispererCodeScanSession(codeScanContext))
        doNothing().`when`(sessionMock).uploadArtifactToS3(any(), any(), any(), any(), isNull())
        doNothing().`when`(sessionMock).sleepThread()

        ToolWindowManager.getInstance(project).registerToolWindow(
            RegisterToolWindowTask(
                id = ProblemsView.ID
            )
        )

        runBlocking {
            val codeScanResponse = sessionMock.run()
            assertThat(codeScanResponse).isInstanceOf<CodeScanResponse.Success>()
            assertThat(codeScanResponse.issues).hasSize(expectedTotalIssues)
            assertThat(codeScanResponse.responseContext.codeScanJobId).isEqualTo("jobId")
            val payloadContext = codeScanResponse.responseContext.payloadContext
            assertThat(payloadContext.totalLines).isEqualTo(expectedTotalLines)
            assertThat(payloadContext.totalFiles).isEqualTo(expectedTotalFiles)
            assertThat(payloadContext.srcPayloadSize).isEqualTo(expectedTotalSize)
            scanManagerSpy.testRenderResponseOnUIThread(
                codeScanResponse.issues,
                codeScanResponse.responseContext.payloadContext.scannedFiles,
                sessionConfigSpy.isProjectTruncated()
            )
            assertNotNull(scanManagerSpy.getScanTree().model)
            val treeModel = scanManagerSpy.getScanTree().model as? CodeWhispererCodeScanTreeModel
            assertNotNull(treeModel)
            assertThat(treeModel.getTotalIssuesCount()).isEqualTo(expectedTotalIssues)
        }
    }

    companion object {
        const val UPLOAD_ID = "uploadId"
        const val JOB_ID = "jobId"
    }
}


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

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

import com.intellij.compiler.CompilerTestUtil
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.module.ModuleManager
import com.intellij.openapi.projectRoots.JavaSdk
import com.intellij.openapi.projectRoots.ProjectJdkTable
import com.intellij.openapi.projectRoots.impl.SdkConfigurationUtil
import com.intellij.openapi.roots.CompilerProjectExtension
import com.intellij.openapi.roots.ModuleRootModificationUtil
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.openapi.vfs.newvfs.impl.VfsRootAccess
import com.intellij.testFramework.IdeaTestUtil
import com.intellij.testFramework.PlatformTestUtil
import com.intellij.testFramework.runInEdtAndWait
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.junit.jupiter.api.assertThrows
import org.mockito.kotlin.any
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import software.aws.toolkits.jetbrains.core.compileProjectAndWait
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.JavaCodeScanSessionConfig
import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
import software.aws.toolkits.jetbrains.utils.rules.addClass
import software.aws.toolkits.jetbrains.utils.rules.addModule
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.BufferedInputStream
import java.util.zip.ZipInputStream
import kotlin.io.path.relativeTo
import kotlin.test.assertNotNull

class CodeWhispererJavaCodeScanTest : CodeWhispererCodeScanTestBase(HeavyJavaCodeInsightTestFixtureRule()) {
    private lateinit var utilsJava: VirtualFile
    private lateinit var test1Java: VirtualFile
    private lateinit var test2Java: VirtualFile
    private lateinit var sessionConfigSpy: JavaCodeScanSessionConfig

    private var totalSize: Long = 0
    private var totalLines: Long = 0

    @Before
    override fun setup() {
        super.setup()
        setupJavaProject()
        sessionConfigSpy = spy(CodeScanSessionConfig.create(utilsJava, project) as JavaCodeScanSessionConfig)
        setupResponse(utilsJava.toNioPath().relativeTo(sessionConfigSpy.projectRoot.toNioPath()))

        mockClient.stub {
            onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
        }
    }

    @Test
    fun `test createPayload`() {
        val payload = sessionConfigSpy.createPayload()
        assertThat(payload.context.totalFiles).isEqualTo(3)

        assertThat(payload.context.scannedFiles.size).isEqualTo(3)
        assertThat(payload.context.scannedFiles).containsExactly(utilsJava, test1Java, test2Java)

        assertThat(payload.context.srcPayloadSize).isEqualTo(totalSize)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Java)
        assertThat(payload.context.totalLines).isEqualTo(totalLines)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }
        assertThat(filesInZip).isEqualTo(6)
    }

    @Test
    fun `test getSourceFilesUnderProjectRoot`() {
        assertThat(sessionConfigSpy.getSourceFilesUnderProjectRoot(utilsJava).size).isEqualTo(3)
    }

    @Test
    fun `test parseImports()`() {
        val utilsJavaImports = sessionConfigSpy.parseImports(utilsJava)
        assertThat(utilsJavaImports.imports.size).isEqualTo(4)

        val test1JavaImports = sessionConfigSpy.parseImports(test1Java)
        assertThat(test1JavaImports.imports.size).isEqualTo(1)
    }

    @Test
    fun `test includeDependencies()`() {
        val payloadMetadata = sessionConfigSpy.includeDependencies()
        assertNotNull(payloadMetadata)
        assertThat(sessionConfigSpy.isProjectTruncated()).isFalse
        val (includedSourceFiles, srcPayloadSize, totalLines, buildPaths) = payloadMetadata
        assertThat(includedSourceFiles).hasSize(3)
        assertThat(srcPayloadSize).isEqualTo(totalSize)
        assertThat(totalLines).isEqualTo(this.totalLines)
        assertThat(buildPaths).hasSize(3)
    }

    @Test
    fun `test getTotalProjectSizeInBytes()`() {
        runBlocking {
            assertThat(sessionConfigSpy.getTotalProjectSizeInBytes()).isEqualTo(totalSize)
        }
    }

    @Test
    fun `selected file larger than payload limit throws exception`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(100)
        }
        assertThrows<CodeWhispererCodeScanException> {
            sessionConfigSpy.createPayload()
        }
    }

    @Test
    fun `test createPayload with custom payload limit`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(900)
        }
        val payload = sessionConfigSpy.createPayload()
        assertNotNull(payload)
        assertThat(sessionConfigSpy.isProjectTruncated()).isTrue
        assertThat(payload.context.totalFiles).isEqualTo(2)

        assertThat(payload.context.scannedFiles.size).isEqualTo(2)
        assertThat(payload.context.scannedFiles).containsExactly(utilsJava, test2Java)

        assertThat(payload.context.srcPayloadSize).isEqualTo(612)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Java)
        assertThat(payload.context.totalLines).isEqualTo(30)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }
        assertThat(filesInZip).isEqualTo(4)
    }

    @Test
    fun `e2e happy path integration test`() {
        assertE2ERunsSuccessfully(sessionConfigSpy, project, totalLines, 3, totalSize, 2)
    }

    private fun compileProject() {
        setUpCompiler()
        compileProjectAndWait(project)
    }

    private fun setUpCompiler() {
        val modules = ModuleManager.getInstance(project).modules

        WriteCommandAction.writeCommandAction(project).run<Nothing> {
            val compilerExtension = CompilerProjectExtension.getInstance(project)
            assertNotNull(compilerExtension)
            compilerExtension.compilerOutputUrl = projectRule.fixture.tempDirFixture.findOrCreateDir("out").url
            val jdkHome = IdeaTestUtil.requireRealJdkHome()
            VfsRootAccess.allowRootAccess(projectRule.fixture.testRootDisposable, jdkHome)
            val jdkHomeDir = LocalFileSystem.getInstance().refreshAndFindFileByPath(jdkHome)
            assertNotNull(jdkHomeDir)
            val jdkName = "Real JDK"
            val jdk = SdkConfigurationUtil.setupSdk(emptyArray(), jdkHomeDir, JavaSdk.getInstance(), false, null, jdkName)
            assertNotNull(jdk)
            ProjectJdkTable.getInstance().addJdk(jdk, projectRule.fixture.testRootDisposable)

            for (module in modules) {
                ModuleRootModificationUtil.setModuleSdk(module, jdk)
            }
        }

        runInEdtAndWait {
            PlatformTestUtil.saveProject(project)
            CompilerTestUtil.saveApplicationSettings()
        }
    }

    private fun setupJavaProject() {
        projectRule as HeavyJavaCodeInsightTestFixtureRule
        val module = projectRule.fixture.addModule("main")

        val utilsClass = projectRule.fixture.addClass(
            module,
            """
            package com.example;

            import java.io.BufferedInputStream;
            import java.util.concurrent.CompletableFuture;
            import java.util.concurrent.TimeUnit;
            import java.util.zip.ZipInputStream;

            public class Utils {
                public int add(int a, int b) {
                    return a + b;
                }

                public int sub(int a, int b) {
                    return a - b;
                }
            }
            """.trimIndent()
        )
        utilsJava = utilsClass.containingFile.virtualFile
        totalSize += utilsJava.length
        totalLines += utilsJava.toNioPath().toFile().readLines().size

        val test1Class = projectRule.fixture.addClass(
            module,
            """
            package com.example2;

            import com.example.Utils;

            public class Test1 {
                Utils utils = new Utils();

                public int fib(int n) {
                   if (n == 0 || n == 1) {
                      return n;
                   }
                  return utils.add(fib(utils.sub(n,1)), fib(utils.sub(n,2)));
                }

                /**
                * Bubble sort algorithm to sort integer array.
                */
                public void bubbleSort(int[] arr) {
                    int n = arr.length;
                    int temp;
                    for(int i=0; i < n; i++) {
                         for(int j=1; j < (n-i); j++) {
                             if(arr[j-1] > arr[j]) {
                                 //swap elements
                                 temp = arr[j-1];
                                 arr[j-1] = arr[j];
                                 arr[j] = temp;
                             }
                         }
                    }
                }
            }

            """.trimIndent()
        )
        test1Java = test1Class.containingFile.virtualFile
        totalSize += test1Java.length
        totalLines += test1Java.toNioPath().toFile().readLines().size

        val test2Class = projectRule.fixture.addClass(
            module,
            """
            package com.example2;

            import com.example.*;

            public class Test2 {
                private Utils utils = new Utils();

                int fib(int n) {
                   if (n == 0 || n == 1) {
                      return n;
                   }
                  return utils.add(fib(utils.sub(n,1)), fib(utils.sub(n,2)));
                }
            }
            """.trimIndent()
        )
        test2Java = test2Class.containingFile.virtualFile
        totalSize += test2Java.length
        totalLines += test2Java.toNioPath().toFile().readLines().size

        projectRule.fixture.addFileToProject("/notIncluded.md", "### should NOT be included")

        compileProject()
    }
}


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

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

import com.intellij.openapi.vfs.VirtualFile
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.junit.jupiter.api.assertThrows
import org.mockito.kotlin.any
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.JavaScriptCodeScanSessionConfig
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.BufferedInputStream
import java.util.zip.ZipInputStream
import kotlin.io.path.relativeTo
import kotlin.test.assertNotNull

class CodeWhispererJavaScriptCodeScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsightTestFixtureRule()) {

    internal lateinit var testJs: VirtualFile
    internal lateinit var utilsJs: VirtualFile
    internal lateinit var helperJs: VirtualFile
    internal lateinit var sessionConfigSpy: JavaScriptCodeScanSessionConfig

    private var totalSize: Long = 0
    private var totalLines: Long = 0

    @Before
    override fun setup() {
        super.setup()
        setupJavaScriptProject()
        sessionConfigSpy = spy(CodeScanSessionConfig.create(testJs, project) as JavaScriptCodeScanSessionConfig)
        setupResponse(testJs.toNioPath().relativeTo(sessionConfigSpy.projectRoot.toNioPath()))

        mockClient.stub {
            onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
        }
    }

    @Test
    fun `test createPayload`() {
        val payload = sessionConfigSpy.createPayload()
        assertNotNull(payload)
        assertThat(payload.context.totalFiles).isEqualTo(3)

        assertThat(payload.context.scannedFiles.size).isEqualTo(3)
        assertThat(payload.context.scannedFiles).containsExactly(testJs, utilsJs, helperJs)

        assertThat(payload.context.srcPayloadSize).isEqualTo(totalSize)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Javascript)
        assertThat(payload.context.totalLines).isEqualTo(totalLines)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }

        assertThat(filesInZip).isEqualTo(3)
    }

    @Test
    fun `test getSourceFilesUnderProjectRoot`() {
        assertThat(sessionConfigSpy.getSourceFilesUnderProjectRoot(testJs).size).isEqualTo(3)
    }

    @Test
    fun `test parseImport()`() {
        val testJsImports = sessionConfigSpy.parseImports(testJs)
        assertThat(testJsImports.size).isEqualTo(3)

        val helperJsImports = sessionConfigSpy.parseImports(helperJs)
        assertThat(helperJsImports.size).isEqualTo(1)
    }

    @Test
    fun `test getImportedFiles()`() {
        val files = sessionConfigSpy.getImportedFiles(testJs, setOf())
        assertNotNull(files)
        assertThat(files).hasSize(2)
        assertThat(files).contains(utilsJs.path)
        assertThat(files).contains(helperJs.path)
    }

    @Test
    fun `test includeDependencies()`() {
        val payloadMetadata = sessionConfigSpy.includeDependencies()
        assertNotNull(payloadMetadata)
        assertThat(sessionConfigSpy.isProjectTruncated()).isFalse
        assertThat(payloadMetadata.sourceFiles.size).isEqualTo(3)
        assertThat(payloadMetadata.payloadSize).isEqualTo(totalSize)
        assertThat(payloadMetadata.linesScanned).isEqualTo(this.totalLines)
        assertThat(payloadMetadata.buildPaths).hasSize(0)
    }

    @Test
    fun `selected file larger than payload limit throws exception`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(100)
        }
        assertThrows<CodeWhispererCodeScanException> {
            sessionConfigSpy.createPayload()
        }
    }

    @Test
    fun `test createPayload with custom payload limit`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(900)
        }
        val payload = sessionConfigSpy.createPayload()
        assertNotNull(payload)
        assertThat(sessionConfigSpy.isProjectTruncated()).isTrue
        assertThat(payload.context.totalFiles).isEqualTo(2)

        assertThat(payload.context.scannedFiles.size).isEqualTo(2)
        assertThat(payload.context.scannedFiles).containsExactly(testJs, utilsJs)

        assertThat(payload.context.srcPayloadSize).isEqualTo(816)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Javascript)
        assertThat(payload.context.totalLines).isEqualTo(31)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }

        assertThat(filesInZip).isEqualTo(2)
    }

    @Test
    fun `e2e happy path integration test`() {
        assertE2ERunsSuccessfully(sessionConfigSpy, projectRule.project, totalLines, 3, totalSize, 2)
    }

    private fun setupJavaScriptProject() {
        testJs = projectRule.fixture.addFileToProject(
            "/test.js",
            """
            import * as utils from "./utils";
            import * from 'nodeJs';
            var myVar = require('./helpers/helper.js');

            var a = 1;
            var b = 2;

            var c = utils.add(a, b);
            var d = myVar.subtract(a, b);
            var e = utils.fib(a, b);
            """.trimIndent()
        ).virtualFile
        totalSize += testJs.length
        totalLines += testJs.toNioPath().toFile().readLines().size

        utilsJs = projectRule.fixture.addFileToProject(
            "/utils.js",
            """
            function add(num1, num2) {
              return num1 + num2;
            }

            function bblSort(arr) {
                for(var i = 0; i < arr.length; i++) {
                    // Last i elements are already in place
                    for(var j = 0; j < ( arr.length - i -1 ); j++) {
                        // Checking if the item at present iteration
                        // is greater than the next iteration
                        if(arr[j] > arr[j+1]) {
                            // If the condition is true then swap them
                            var temp = arr[j]
                            arr[j] = arr[j + 1]
                            arr[j+1] = temp
                        }
                    }
                }
                // Print the sorted array
                console.log(arr);
            }
            """.trimIndent()
        ).virtualFile
        totalSize += utilsJs.length
        totalLines += utilsJs.toNioPath().toFile().readLines().size

        helperJs = projectRule.fixture.addFileToProject(
            "/helpers/helper.js",
            """
            import * as h from './helpers/helper.js'
            function subtract(num1, num2) {
              return num1 - num2;
            }

            function fibonacci(num) {
                var num1=0;
                var num2=1;
                var sum;
                var i=0;
                for (i = 0; i < num; i++) {
                    sum=h.add(num1, num2);
                    num1=num2;
                    num2=sum;
                }
                return num2;
            }

            """.trimIndent()
        ).virtualFile
        totalSize += helperJs.length
        totalLines += helperJs.toNioPath().toFile().readLines().size

        projectRule.fixture.addFileToProject("/notIncluded.md", "### should NOT be included")
    }
}


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

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

import com.intellij.openapi.vfs.VirtualFile
import kotlinx.coroutines.runBlocking
import org.assertj.core.api.Assertions.assertThat
import org.junit.Before
import org.junit.Test
import org.junit.jupiter.api.assertThrows
import org.mockito.kotlin.any
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PythonCodeScanSessionConfig
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.BufferedInputStream
import java.util.zip.ZipInputStream
import kotlin.io.path.relativeTo
import kotlin.test.assertNotNull

class CodeWhispererPythonCodeScanTest : CodeWhispererCodeScanTestBase(PythonCodeInsightTestFixtureRule()) {
    private lateinit var testPy: VirtualFile
    private lateinit var utilsPy: VirtualFile
    private lateinit var helperPy: VirtualFile
    private lateinit var sessionConfigSpy: PythonCodeScanSessionConfig

    private var totalSize: Long = 0
    private var totalLines: Long = 0

    @Before
    override fun setup() {
        super.setup()
        setupPythonProject()
        sessionConfigSpy = spy(CodeScanSessionConfig.create(testPy, project) as PythonCodeScanSessionConfig)
        setupResponse(testPy.toNioPath().relativeTo(sessionConfigSpy.projectRoot.toNioPath()))

        mockClient.stub {
            onGeneric { createUploadUrl(any()) }.thenReturn(fakeCreateUploadUrlResponse)
            onGeneric { createCodeScan(any(), any()) }.thenReturn(fakeCreateCodeScanResponse)
            onGeneric { getCodeScan(any(), any()) }.thenReturn(fakeGetCodeScanResponse)
            onGeneric { listCodeScanFindings(any(), any()) }.thenReturn(fakeListCodeScanFindingsResponse)
        }
    }

    @Test
    fun `test createPayload`() {
        val payload = sessionConfigSpy.createPayload()
        assertNotNull(payload)
        assertThat(payload.context.totalFiles).isEqualTo(3)

        assertThat(payload.context.scannedFiles.size).isEqualTo(3)
        assertThat(payload.context.scannedFiles).containsExactly(testPy, utilsPy, helperPy)

        assertThat(payload.context.srcPayloadSize).isEqualTo(totalSize)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Python)
        assertThat(payload.context.totalLines).isEqualTo(totalLines)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }

        assertThat(filesInZip).isEqualTo(3)
    }

    @Test
    fun `test getSourceFilesUnderProjectRoot`() {
        assertThat(sessionConfigSpy.getSourceFilesUnderProjectRoot(testPy).size).isEqualTo(3)
    }

    @Test
    fun `test parseImport()`() {
        val testPyImports = sessionConfigSpy.parseImports(testPy)
        assertThat(testPyImports.size).isEqualTo(4)

        val helperPyImports = sessionConfigSpy.parseImports(helperPy)
        assertThat(helperPyImports.size).isEqualTo(1)
    }

    @Test
    fun `test getImportedFiles()`() {
        val files = sessionConfigSpy.getImportedFiles(testPy, setOf())
        assertNotNull(files)
        assertThat(files).hasSize(2)
        assertThat(files).contains(utilsPy.path)
        assertThat(files).contains(helperPy.path)
    }

    @Test
    fun `test includeDependencies()`() {
        val payloadMetadata = sessionConfigSpy.includeDependencies()
        assertNotNull(payloadMetadata)
        assertThat(sessionConfigSpy.isProjectTruncated()).isFalse
        val (includedSourceFiles, srcPayloadSize, totalLines, buildPaths) = payloadMetadata
        assertThat(includedSourceFiles.size).isEqualTo(3)
        assertThat(srcPayloadSize).isEqualTo(totalSize)
        assertThat(totalLines).isEqualTo(this.totalLines)
        assertThat(buildPaths).hasSize(0)
    }

    @Test
    fun `test getTotalProjectSizeInBytes()`() {
        runBlocking {
            assertThat(sessionConfigSpy.getTotalProjectSizeInBytes()).isEqualTo(totalSize)
        }
    }

    @Test
    fun `selected file larger than payload limit throws exception`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(100)
        }
        assertThrows<CodeWhispererCodeScanException> {
            sessionConfigSpy.createPayload()
        }
    }

    @Test
    fun `test createPayload with custom payload limit`() {
        sessionConfigSpy.stub {
            onGeneric { getPayloadLimitInBytes() }.thenReturn(900)
        }
        val payload = sessionConfigSpy.createPayload()
        assertNotNull(payload)
        assertThat(sessionConfigSpy.isProjectTruncated()).isTrue
        assertThat(payload.context.totalFiles).isEqualTo(2)

        assertThat(payload.context.scannedFiles.size).isEqualTo(2)
        assertThat(payload.context.scannedFiles).containsExactly(testPy, helperPy)

        assertThat(payload.context.srcPayloadSize).isEqualTo(363)
        assertThat(payload.context.language).isEqualTo(CodewhispererLanguage.Python)
        assertThat(payload.context.totalLines).isEqualTo(18)
        assertNotNull(payload.srcZip)

        val bufferedInputStream = BufferedInputStream(payload.srcZip.inputStream())
        val zis = ZipInputStream(bufferedInputStream)
        var filesInZip = 0
        while (zis.nextEntry != null) {
            filesInZip += 1
        }

        assertThat(filesInZip).isEqualTo(2)
    }

    @Test
    fun `e2e happy path integration test`() {
        assertE2ERunsSuccessfully(sessionConfigSpy, project, totalLines, 3, totalSize, 2)
    }

    private fun setupPythonProject() {
        testPy = projectRule.fixture.addFileToProject(
            "/test.py",
            """
            import numpy as np
            import utils
            import helpers.helper
            import test2

            a = 1
            b = 2
            print(utils.add(a, b))
            println(helper.subtract(a, b))
            println(utils.fib(5))
            """.trimIndent()
        ).virtualFile
        totalSize += testPy.length
        totalLines += testPy.toNioPath().toFile().readLines().size

        utilsPy = projectRule.fixture.addFileToProject(
            "/utils.py",
            """
            def add(num1, num2
                return num1 + num2

            def multiply(num1, num2)
                return num1 * num2

            def divide(num1, num2)
                return num1 / num2

            def bubbleSort(arr):
            	n = len(arr)
            	# optimize code, so if the array is already sorted, it doesn't need
            	# to go through the entire process
            	swapped = False
            	# Traverse through all array elements
            	for i in range(n-1):
            		# range(n) also work but outer loop will
            		# repeat one time more than needed.
            		# Last i elements are already in place
            		for j in range(0, n-i-1):

            			# traverse the array from 0 to n-i-1
            			# Swap if the element found is greater
            			# than the next element
            			if arr[j] > arr[j + 1]:
            				swapped = True
            				arr[j], arr[j + 1] = arr[j + 1], arr[j]

            		if not swapped:
            			# if we haven't needed to make a single swap, we
            			# can just exit the main loop.
            			return

            """.trimIndent()
        ).virtualFile
        totalSize += utilsPy.length
        totalLines += utilsPy.toNioPath().toFile().readLines().size

        helperPy = projectRule.fixture.addFileToProject(
            "/helpers/helper.py",
            """
            from helpers import helper as h
            def subtract(num1, num2)
                return num1 - num2

            def fib(num):
                if num == 0: return 0
                if num in [1,2]: return 1
                return h.add(fib(num-1), fib(num-2))

            """.trimIndent()
        ).virtualFile
        totalSize += helperPy.length
        totalLines += helperPy.toNioPath().toFile().readLines().size

        projectRule.fixture.addFileToProject("/notIncluded.md", "### should NOT be included")
    }
}