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

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

import com.intellij.openapi.actionSystem.ActionPromoter
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.editor.actionSystem.EditorAction
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.handlers.CodeWhispererPopupLeftArrowHandler
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.handlers.CodeWhispererPopupRightArrowHandler
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.handlers.CodeWhispererPopupTabHandler

class CodeWhispererActionPromoter : ActionPromoter {
    override fun promote(actions: MutableList<out AnAction>, context: DataContext): MutableList<AnAction> {
        val results = actions.toMutableList()
        results.sortWith { a, b ->
            if (isCodeWhispererPopupAction(a)) {
                return@sortWith -1
            } else if (isCodeWhispererPopupAction(b)) {
                return@sortWith 1
            } else {
                0
            }
        }
        return results
    }

    private fun isCodeWhispererAcceptAction(action: AnAction): Boolean =
        action is EditorAction && action.handler is CodeWhispererPopupTabHandler

    private fun isCodeWhispererNavigateAction(action: AnAction): Boolean =
        action is EditorAction && (
            action.handler is CodeWhispererPopupRightArrowHandler ||
                action.handler is CodeWhispererPopupLeftArrowHandler
            )

    private fun isCodeWhispererPopupAction(action: AnAction): Boolean =
        isCodeWhispererAcceptAction(action) || isCodeWhispererNavigateAction(action)
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.ide.BrowserUtil
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAware
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODEWHISPERER_LEARN_MORE_URI
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODEWHISPERER_LOGIN_LEARN_MORE_URI
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODEWHISPERER_SSO_LEARN_MORE_URI
import software.aws.toolkits.resources.message
import java.net.URI

class CodeWhispererLearnMoreAction :
    AnAction(
        message("codewhisperer.explorer.learn_more"),
        null,
        AllIcons.Actions.Help
    ),
    DumbAware {

    override fun actionPerformed(e: AnActionEvent) {
        BrowserUtil.browse(URI(CODEWHISPERER_LEARN_MORE_URI))
    }
}

class CodeWhispererSsoLearnMoreAction : AnAction(message("aws.settings.learn_more")), DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        BrowserUtil.browse(URI(CODEWHISPERER_SSO_LEARN_MORE_URI))
    }
}

class CodeWhispererLoginLearnMoreAction : AnAction(message("aws.settings.learn_more")), DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        BrowserUtil.browse(URI(CODEWHISPERER_LOGIN_LEARN_MORE_URI))
    }
}


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

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

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.LangDataKeys
import com.intellij.openapi.project.DumbAware
import icons.AwsIcons
import software.aws.toolkits.jetbrains.ui.feedback.FeedbackDialog
import software.aws.toolkits.resources.message

class CodeWhispererProvideFeedbackAction :
    AnAction(
        message("codewhisperer.feedback"),
        null,
        AwsIcons.Misc.SMILE_GREY
    ),
    DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        FeedbackDialog(e.getRequiredData(LangDataKeys.PROJECT)).showAndGet()
    }
}


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

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

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.CommonDataKeys
import com.intellij.openapi.project.DumbAware
import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.CodewhispererTriggerType

class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigger.service")), DumbAware {
    override fun update(e: AnActionEvent) {
        e.presentation.isEnabled = e.project != null && e.getData(CommonDataKeys.EDITOR) != null
    }

    override fun actionPerformed(e: AnActionEvent) {
        val latencyContext = LatencyContext()
        latencyContext.codewhispererPreprocessingStart = System.nanoTime()
        latencyContext.codewhispererEndToEndStart = System.nanoTime()
        val editor = e.getRequiredData(CommonDataKeys.EDITOR)
        if (!CodeWhispererService.getInstance().canDoInvocation(editor, CodewhispererTriggerType.OnDemand)) {
            return
        }

        val triggerType = TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown())
        CodeWhispererService.getInstance().showRecommendationsInPopup(editor, triggerType, latencyContext)
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.LangDataKeys
import com.intellij.openapi.options.ShowSettingsUtil
import com.intellij.openapi.project.DumbAware
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
import software.aws.toolkits.resources.message

class CodeWhispererShowSettingsAction :
    AnAction(
        message("codewhisperer.settings.show.label"),
        null,
        AllIcons.General.GearPlain
    ),
    DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        ShowSettingsUtil.getInstance().showSettingsDialog(e.getRequiredData(LangDataKeys.PROJECT), CodeWhispererConfigurable::class.java)
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.ide.BrowserUtil
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAware
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.resources.message
import java.net.URI

class CodeWhispererWhatIsAction :
    AnAction(
        message("codewhisperer.explorer.what_is"),
        null,
        AllIcons.Actions.Help
    ),
    DumbAware {
    override fun update(e: AnActionEvent) {
        e.project?.let {
            e.presentation.isEnabledAndVisible = isCodeWhispererEnabled(it)
        }
    }

    override fun actionPerformed(e: AnActionEvent) {
        BrowserUtil.browse(URI(CodeWhispererConstants.CODEWHISPERER_LEARN_MORE_URI))
    }
}


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

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

import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginDialog
import software.aws.toolkits.resources.message

/**
 * Action prompting users to switch to SSO based credential, will nullify accountless credential (delete)
 */
class ConnectWithAwsToContinueActionWarn : DumbAwareAction(message("codewhisperer.notification.accountless.warn.action.connect")) {
    override fun actionPerformed(e: AnActionEvent) {
        e.project?.let {
            runInEdt {
                CodeWhispererLoginDialog(it).showAndGet()
            }
        }
    }
}
class ConnectWithAwsToContinueActionError : DumbAwareAction(message("codewhisperer.notification.accountless.error.action.connect")) {
    override fun actionPerformed(e: AnActionEvent) {
        e.project?.let {
            runInEdt {
                CodeWhispererLoginDialog(it).showAndGet()
            }
        }
    }
}


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

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

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAware
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyInfoAccountless
import software.aws.toolkits.resources.message

class DoNotShowAgainActionWarn : AnAction(message("codewhisperer.notification.accountless.warn.dont.show.again")), DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        notifyInfoAccountless()
        CodeWhispererExplorerActionManager.getInstance().setDoNotShowAgainWarn(true)
    }
}

class DoNotShowAgainActionError : AnAction(message("codewhisperer.notification.accountless.warn.dont.show.again")), DumbAware {
    override fun actionPerformed(e: AnActionEvent) {
        notifyInfoAccountless()
        CodeWhispererExplorerActionManager.getInstance().setDoNotShowAgainError(true)
    }
}


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

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

import software.aws.toolkits.resources.message

open class CodeWhispererCodeScanException(override val message: String?) : RuntimeException()

internal fun noFileOpenError(): Nothing =
    throw CodeWhispererCodeScanException(message("codewhisperer.codescan.no_file_open"))

internal fun codeScanFailed(): Nothing =
    throw CodeWhispererCodeScanException(message("codewhisperer.codescan.service_error"))

internal fun cannotFindFile(file: String?): Nothing =
    error(message("codewhisperer.codescan.file_not_found", file ?: ""))

internal fun cannotFindBuildArtifacts(): Nothing =
    throw CodeWhispererCodeScanException(message("codewhisperer.codescan.build_artifacts_not_found"))

internal fun fileFormatNotSupported(format: String): Nothing =
    throw CodeWhispererCodeScanException(message("codewhisperer.codescan.file_ext_not_supported", format))

internal fun fileTooLarge(presentableSize: String): Nothing =
    throw CodeWhispererCodeScanException(message("codewhisperer.codescan.file_too_large", presentableSize))


// Copyright 2023 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.application.runInEdt
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.fileEditor.OpenFileDescriptor
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.ui.ColoredTreeCellRenderer
import com.intellij.ui.DoubleClickListener
import com.intellij.ui.ScrollPaneFactory
import com.intellij.ui.SimpleTextAttributes
import com.intellij.ui.TreeSpeedSearch
import com.intellij.ui.treeStructure.Tree
import com.intellij.util.ui.UIUtil
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.resources.message
import java.awt.BorderLayout
import java.awt.event.MouseEvent
import javax.swing.BorderFactory
import javax.swing.JLabel
import javax.swing.JPanel
import javax.swing.JTree
import javax.swing.tree.DefaultMutableTreeNode
import javax.swing.tree.DefaultTreeModel

internal class CodeWhispererCodeScanHighlightingFilesPanel(private val project: Project, files: List<VirtualFile>) : JPanel(BorderLayout()) {

    init {
        removeAll()
        val scannedFilesTreeNodeRoot = DefaultMutableTreeNode("CodeWhisperer scanned files for security scan")
        files.forEach {
            scannedFilesTreeNodeRoot.add(DefaultMutableTreeNode(it))
        }
        val scannedFilesTreeModel = DefaultTreeModel(scannedFilesTreeNodeRoot)
        val scannedFilesTree: Tree = Tree().apply {
            isRootVisible = false
            object : DoubleClickListener() {
                var status = true
                override fun onDoubleClick(event: MouseEvent): Boolean {
                    val fileNode = (event.source as Tree).selectionPath?.lastPathComponent as? DefaultMutableTreeNode
                    val file = fileNode?.userObject as? VirtualFile ?: return false
                    runInEdt {
                        val editor = FileEditorManager.getInstance(project).openTextEditor(
                            OpenFileDescriptor(project, file),
                            true
                        )
                        if (editor == null) {
                            LOG.error { "Cannot fetch editor for the file ${file.path}" }
                            status = false
                            return@runInEdt
                        }
                    }
                    return status
                }
            }.installOn(this)
            model = scannedFilesTreeModel
            cellRenderer = ScannedFilesTreeCellRenderer()
            repaint()
        }
        TreeSpeedSearch(scannedFilesTree, TreeSpeedSearch.NODE_DESCRIPTOR_TOSTRING, true)
        val scrollPane = ScrollPaneFactory.createScrollPane(scannedFilesTree, true)
        val htmlText = message("codewhisperer.codescan.scanned_files_heading", files.size)
        val label = JLabel(htmlText, JLabel.LEFT).apply {
            border = BorderFactory.createEmptyBorder(7, 26, 8, 11)
        }

        add(BorderLayout.NORTH, label)
        add(BorderLayout.CENTER, scrollPane)
        isVisible = true
        revalidate()
    }

    private class ScannedFilesTreeCellRenderer : ColoredTreeCellRenderer() {
        override fun customizeCellRenderer(
            tree: JTree,
            value: Any,
            selected: Boolean,
            expanded: Boolean,
            leaf: Boolean,
            row: Int,
            hasFocus: Boolean
        ) {
            value as DefaultMutableTreeNode
            val file = value.userObject as? VirtualFile
            val attributes = SimpleTextAttributes.LINK_ATTRIBUTES
            icon = file?.fileType?.icon
            font = UIUtil.getTreeFont()
            file?.path?.let { append(it, attributes) }
        }
    }

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanHighlightingFilesPanel>()
    }
}


// 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.analysis.problemsView.toolWindow.ProblemsView
import com.intellij.codeHighlighting.HighlightDisplayLevel
import com.intellij.codeInspection.util.InspectionMessage
import com.intellij.icons.AllIcons
import com.intellij.openapi.application.ApplicationInfo
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.EditorFactory
import com.intellij.openapi.editor.ex.RangeHighlighterEx
import com.intellij.openapi.editor.impl.DocumentMarkupModel
import com.intellij.openapi.editor.markup.HighlighterLayer
import com.intellij.openapi.editor.markup.HighlighterTargetArea
import com.intellij.openapi.editor.markup.MarkupModel
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.fileEditor.impl.FileDocumentManagerImpl
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.MessageDialogBuilder
import com.intellij.openapi.util.Disposer
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.refactoring.suggested.range
import com.intellij.ui.content.ContentManagerEvent
import com.intellij.ui.content.ContentManagerListener
import com.intellij.ui.treeStructure.Tree
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.Job
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.async
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.time.withTimeout
import kotlinx.coroutines.withContext
import org.jetbrains.annotations.TestOnly
import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
import software.aws.toolkits.core.utils.WaiterTimeoutException
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineUiContext
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.core.explorer.refreshDevToolTree
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.listeners.CodeWhispererCodeScanDocumentListener
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.listeners.CodeWhispererCodeScanEditorMouseMotionListener
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.overlaps
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanTelemetryEvent
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.INACTIVE_TEXT_COLOR
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.ISSUE_HIGHLIGHT_TEXT_ATTRIBUTES
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.Result
import java.time.Duration
import java.time.Instant
import java.util.concurrent.atomic.AtomicBoolean
import javax.swing.Icon
import javax.swing.tree.DefaultMutableTreeNode
import javax.swing.tree.TreePath
import kotlin.coroutines.CoroutineContext

class CodeWhispererCodeScanManager(val project: Project) {
    private val codeScanResultsPanel by lazy {
        CodeWhispererCodeScanResultsView(project)
    }
    private val codeScanIssuesContent by lazy {
        val contentManager = getProblemsWindow().contentManager
        contentManager.factory.createContent(
            codeScanResultsPanel,
            message("codewhisperer.codescan.scan_display"),
            false
        ).also {
            Disposer.register(contentManager, it)
            contentManager.addContentManagerListener(object : ContentManagerListener {
                override fun contentRemoved(event: ContentManagerEvent) {
                    if (event.content == it) reset()
                }
            })
        }
    }

    private val fileNodeLookup = mutableMapOf<VirtualFile, DefaultMutableTreeNode>()
    private val scanNodesLookup = mutableMapOf<VirtualFile, MutableList<DefaultMutableTreeNode>>()

    private val documentListener = CodeWhispererCodeScanDocumentListener(project)
    private val editorMouseListener = CodeWhispererCodeScanEditorMouseMotionListener(project)

    private val isCodeScanInProgress = AtomicBoolean(false)

    private lateinit var codeScanJob: Job

    /**
     * Returns true if the code scan is in progress.
     * This function will return true for a cancelled code scan job which is in cancellation state.
     */
    fun isCodeScanInProgress(): Boolean = isCodeScanInProgress.get()

    /**
     * Code scan job is active when the [Job] is started and is in active state.
     */
    fun isCodeScanJobActive(): Boolean = this::codeScanJob.isInitialized && codeScanJob.isActive

    fun getRunActionButtonIcon(): Icon = if (isCodeScanInProgress()) AllIcons.Process.Step_1 else AllIcons.Actions.Execute

    fun getActionButtonIconForExplorerNode(): Icon = if (isCodeScanInProgress()) AllIcons.Actions.Suspend else AllIcons.Actions.Execute

    fun getActionButtonText(): String = if (!isCodeScanInProgress()) message("codewhisperer.codescan.run_scan") else message("codewhisperer.codescan.stop_scan")

    /**
     * Triggers a code scan and displays results in the new tab in problems view panel.
     */
    fun runCodeScan() {
        if (!isCodeWhispererEnabled(project)) return

        // Return if a scan is already in progress.
        if (isCodeScanInProgress.getAndSet(true)) return
        if (promptReAuth(project)) {
            isCodeScanInProgress.set(false)
            return
        }

        // Prepare for a code scan
        beforeCodeScan()
        // launch code scan coroutine
        codeScanJob = launchCodeScanCoroutine()
    }

    fun stopCodeScan() {
        // Return if code scan job is not active.
        if (!codeScanJob.isActive) return
        if (isCodeScanInProgress() && confirmCancelCodeScan()) {
            LOG.info { "Security scan stopped by user..." }
            // Checking `codeScanJob.isActive` to ensure that the job is not already completed by the time user confirms.
            if (codeScanJob.isActive) {
                codeScanResultsPanel.setStoppingCodeScan()
                codeScanJob.cancel(CancellationException("User requested cancellation"))
            }
        }
    }

    private fun confirmCancelCodeScan(): Boolean = MessageDialogBuilder
        .okCancel(message("codewhisperer.codescan.stop_scan"), message("codewhisperer.codescan.stop_scan_confirm_message"))
        .yesText(message("codewhisperer.codescan.stop_scan_confirm_button"))
        .ask(project)

    private fun launchCodeScanCoroutine() = projectCoroutineScope(project).launch {
        var codeScanStatus: Result = Result.Failed
        val startTime = Instant.now().toEpochMilli()
        var codeScanResponseContext = defaultCodeScanResponseContext()
        var getProjectSize: Deferred<Long?> = async { null }
        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
        try {
            val file = FileEditorManager.getInstance(project).selectedEditor?.file
                ?: noFileOpenError()
            val codeScanSessionConfig = CodeScanSessionConfig.create(file, project)
            withTimeout(Duration.ofSeconds(codeScanSessionConfig.overallJobTimeoutInSeconds())) {
                // 1. Generate truncation (zip files) based on the current editor.
                LOG.debug { "Creating context truncation for file ${file.path}" }
                val sessionContext = CodeScanSessionContext(project, codeScanSessionConfig)
                val session = CodeWhispererCodeScanSession(sessionContext)
                val codeScanResponse = session.run()
                codeScanResponseContext = codeScanResponse.responseContext
                when (codeScanResponse) {
                    is CodeScanResponse.Success -> {
                        val issues = codeScanResponse.issues
                        coroutineContext.ensureActive()
                        renderResponseOnUIThread(
                            issues,
                            codeScanResponse.responseContext.payloadContext.scannedFiles,
                            codeScanSessionConfig.isProjectTruncated()
                        )
                        codeScanStatus = Result.Succeeded
                    }

                    is CodeScanResponse.Failure -> {
                        if (codeScanResponse.failureReason !is TimeoutCancellationException && codeScanResponse.failureReason is CancellationException) {
                            codeScanStatus = Result.Cancelled
                        }
                        throw codeScanResponse.failureReason
                    }
                }
                LOG.info { "Security scan completed." }
            }
            getProjectSize = async {
                codeScanSessionConfig.getTotalProjectSizeInBytes()
            }
        } catch (e: Exception) {
            isCodeScanInProgress.set(false)
            val errorMessage = handleException(coroutineContext, e)
            codeScanResponseContext = codeScanResponseContext.copy(reason = errorMessage)
        } finally {
            // After code scan
            afterCodeScan()
            launch {
                val duration = (Instant.now().toEpochMilli() - startTime).toDouble()
                CodeWhispererTelemetryService.getInstance().sendSecurityScanEvent(
                    CodeScanTelemetryEvent(codeScanResponseContext, duration, codeScanStatus, getProjectSize.await()?.toDouble(), connection)
                )
            }
        }
    }

    fun handleException(coroutineContext: CoroutineContext, e: Exception): String {
        val errorMessage = when (e) {
            is CodeWhispererException -> e.awsErrorDetails().errorMessage() ?: message("codewhisperer.codescan.service_error")
            is CodeWhispererCodeScanException -> e.message
            is WaiterTimeoutException, is TimeoutCancellationException -> message("codewhisperer.codescan.scan_timed_out")
            is CancellationException -> "Code scan job cancelled by user."
            else -> null
        } ?: message("codewhisperer.codescan.run_scan_error")

        val errorCode = (e as? CodeWhispererException)?.awsErrorDetails()?.errorCode()
        val requestId = if (e is CodeWhispererException) e.requestId() else null

        if (!coroutineContext.isActive) {
            codeScanResultsPanel.setDefaultUI()
        } else {
            codeScanResultsPanel.showError(errorMessage)
        }

        if (
            e is ThrottlingException &&
            e.message == CodeWhispererConstants.THROTTLING_MESSAGE
        ) {
            CodeWhispererExplorerActionManager.getInstance().setSuspended(project)
            CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit(project, isCodeScan = true)
        }

        LOG.error {
            "Failed to run security scan and display results. Caused by: $errorMessage, status code: $errorCode, " +
                "exception: ${e::class.simpleName}, request ID: $requestId " +
                "Jetbrains IDE: ${ApplicationInfo.getInstance().fullApplicationName}, " +
                "IDE version: ${ApplicationInfo.getInstance().apiVersion}, " +
                "stacktrace: ${e.stackTrace.contentDeepToString()}"
        }
        return errorMessage
    }

    /**
     * The initial landing UI for the code scan results view panel.
     * This method adds code content to the problems view if not already added.
     * When [setSelected] is true, code scan panel is set to be in focus.
     */
    fun addCodeScanUI(setSelected: Boolean = false) = runInEdt {
        reset()
        val problemsWindow = getProblemsWindow()
        if (!problemsWindow.contentManager.contents.contains(codeScanIssuesContent)) {
            problemsWindow.contentManager.addContent(codeScanIssuesContent)
        }
        codeScanIssuesContent.displayName = message("codewhisperer.codescan.scan_display")
        if (setSelected) {
            problemsWindow.contentManager.setSelectedContent(codeScanIssuesContent)
            problemsWindow.show()
        }
    }

    fun removeCodeScanUI() = runInEdt {
        val problemsWindow = getProblemsWindow()
        if (problemsWindow.contentManager.contents.contains(codeScanIssuesContent)) {
            problemsWindow.contentManager.removeContent(codeScanIssuesContent, false)
        }
    }

    fun getScanNodesInRange(file: VirtualFile, startOffset: Int): List<DefaultMutableTreeNode> =
        getOverlappingScanNodes(file, TextRange.create(startOffset, startOffset + 1))

    fun getOverlappingScanNodes(file: VirtualFile, range: TextRange): List<DefaultMutableTreeNode> = synchronized(scanNodesLookup) {
        scanNodesLookup[file]?.mapNotNull { node ->
            val issue = node.userObject as CodeWhispererCodeScanIssue
            if (issue.textRange?.overlaps(range) == true) node else null
        } ?: listOf()
    }

    fun getScanTree(): Tree = codeScanResultsPanel.getCodeScanTree()

    /**
     * Updates the scan nodes in a [file] with the new text range.
     */
    fun updateScanNodes(file: VirtualFile) {
        scanNodesLookup[file]?.forEach { node ->
            val issue = node.userObject as CodeWhispererCodeScanIssue
            val newRange = issue.rangeHighlighter?.range
            val oldRange = issue.textRange
            // Check if the location of the issue is changed and only update the valid nodes.
            if (newRange != null && oldRange != newRange && !issue.isInvalid) {
                val newIssue = issue.copyRange(newRange)
                synchronized(node) {
                    getScanTree().model.valueForPathChanged(TreePath(node.path), newIssue)
                    node.userObject = newIssue
                }
            }
        }
    }

    private fun CodeWhispererCodeScanIssue.copyRange(newRange: TextRange): CodeWhispererCodeScanIssue {
        val newStartLine = document.getLineNumber(newRange.startOffset)
        val newStartCol = newRange.startOffset - document.getLineStartOffset(newStartLine)
        val newEndLine = document.getLineNumber(newRange.endOffset)
        val newEndCol = newRange.endOffset - document.getLineStartOffset(newEndLine)
        return copy(
            startLine = newStartLine + 1,
            startCol = newStartCol + 1,
            endLine = newEndLine + 1,
            endCol = newEndCol + 1
        )
    }

    private fun getProblemsWindow() = ProblemsView.getToolWindow(project)
        ?: error(message("codewhisperer.codescan.problems_window_not_found"))

    private fun reset() = runInEdt {
        // Remove previous document listeners before starting a new scan.
        removeListeners()
        fileNodeLookup.clear()
        // Erase all range highlighter before cleaning up.
        scanNodesLookup.apply {
            forEach { (_, nodes) ->
                nodes.forEach { node ->
                    val issue = node.userObject as CodeWhispererCodeScanIssue
                    issue.rangeHighlighter?.dispose()
                }
            }
            clear()
        }
    }

    private fun addListeners() {
        fileNodeLookup.keys.forEach { file ->
            runInEdt {
                val document = FileDocumentManager.getInstance().getDocument(file)
                if (document == null) {
                    LOG.error { message("codewhisperer.codescan.file_not_found", file.path) }
                    return@runInEdt
                }
                document.addDocumentListener(documentListener, codeScanIssuesContent)
            }
        }
        EditorFactory.getInstance().eventMulticaster.addEditorMouseMotionListener(
            editorMouseListener,
            codeScanIssuesContent
        )
    }

    private fun removeListeners() {
        fileNodeLookup.keys.forEach { file ->
            runInEdt {
                val document = FileDocumentManager.getInstance().getDocument(file)
                if (document == null) {
                    LOG.error { message("codewhisperer.codescan.file_not_found", file.path) }
                    return@runInEdt
                }
                document.removeDocumentListener(documentListener)
            }
        }
        EditorFactory.getInstance().eventMulticaster.removeEditorMouseMotionListener(editorMouseListener)
    }

    private fun beforeCodeScan() {
        // Refresh CodeWhisperer Explorer tree node to reflect scan in progress.
        project.refreshDevToolTree()
        addCodeScanUI(setSelected = true)
        // Show in progress indicator
        codeScanResultsPanel.showInProgressIndicator()
        (FileDocumentManagerImpl.getInstance() as FileDocumentManagerImpl).saveAllDocuments(false)
        LOG.info { "Starting security scan on package ${project.name}..." }
    }

    private fun afterCodeScan() {
        isCodeScanInProgress.set(false)
        project.refreshDevToolTree()
    }

    /**
     * Creates a CodeWhisperer code scan issues tree.
     * For each scan node:
     *   1. (Add file node if not already present and) add scan node to the file node.
     *   2. Update the lookups - [fileNodeLookup] for efficiently adding scan nodes and
     *   [scanNodesLookup] for receiving the editor events and updating the corresponding scan nodes.
     */
    private fun createCodeScanIssuesTree(codeScanIssues: List<CodeWhispererCodeScanIssue>): DefaultMutableTreeNode {
        LOG.debug { "Rendering response from the scan API" }

        val codeScanTreeNodeRoot = DefaultMutableTreeNode("CodeWhisperer Code scan results")
        codeScanIssues.forEach { issue ->
            val fileNode = synchronized(fileNodeLookup) {
                fileNodeLookup.getOrPut(issue.file) {
                    val node = DefaultMutableTreeNode(issue.file)
                    synchronized(codeScanTreeNodeRoot) {
                        codeScanTreeNodeRoot.add(node)
                    }
                    node
                }
            }

            val scanNode = DefaultMutableTreeNode(issue)
            fileNode.add(scanNode)
            scanNodesLookup.getOrPut(issue.file) {
                mutableListOf()
            }.add(scanNode)
        }
        // Add document and editor listeners to the documents having scan issues.
        addListeners()
        return codeScanTreeNodeRoot
    }

    suspend fun renderResponseOnUIThread(issues: List<CodeWhispererCodeScanIssue>, scannedFiles: List<VirtualFile>, isProjectTruncated: Boolean) {
        withContext(getCoroutineUiContext()) {
            val root = createCodeScanIssuesTree(issues)
            val codeScanTreeModel = CodeWhispererCodeScanTreeModel(root)
            val totalIssuesCount = codeScanTreeModel.getTotalIssuesCount()
            if (totalIssuesCount > 0) {
                codeScanIssuesContent.displayName =
                    message("codewhisperer.codescan.scan_display_with_issues", totalIssuesCount, INACTIVE_TEXT_COLOR)
            }
            codeScanResultsPanel.updateAndDisplayScanResults(codeScanTreeModel, scannedFiles, isProjectTruncated)
        }
    }

    @TestOnly
    suspend fun testRenderResponseOnUIThread(issues: List<CodeWhispererCodeScanIssue>, scannedFiles: List<VirtualFile>, isProjectTruncated: Boolean) {
        assert(ApplicationManager.getApplication().isUnitTestMode)
        renderResponseOnUIThread(issues, scannedFiles, isProjectTruncated)
    }

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanManager>()
        fun getInstance(project: Project): CodeWhispererCodeScanManager = project.service()
    }
}

/**
 * Wrapper Data class representing a CodeWhisperer code scan issue.
 * @param title is shown in the code scan tree in the `CodeWhisperer Security Issues` tab.
 * @param description is shown in the tooltip of the scan node and also shown when the mouse
 * is hovered over the highlighted text in the editor.
 */
data class CodeWhispererCodeScanIssue(
    val project: Project,
    val file: VirtualFile,
    val startLine: Int,
    val startCol: Int,
    val endLine: Int,
    val endCol: Int,
    val title: @InspectionMessage String,
    val description: Description,
    val issueSeverity: HighlightDisplayLevel = HighlightDisplayLevel.WARNING,
    val isInvalid: Boolean = false,
    var rangeHighlighter: RangeHighlighterEx? = null
) {
    override fun toString(): String = title

    val document = runReadAction {
        FileDocumentManager.getInstance().getDocument(file)
            ?: cannotFindFile(file.path)
    }

    /**
     * Immutable value of the textRange at the time the issue was constructed.
     */
    val textRange = toTextRange()

    fun displayTextRange() = "[$startLine:$startCol-$endLine:$endCol]"

    /**
     * Adds a range highlighter for the corresponding code scan issue with the given markup model.
     * Note that the default markup model which is fetched from [DocumentMarkupModel] can be null.
     * Must be run in [runInEdt].
     */
    fun addRangeHighlighter(
        markupModel: MarkupModel? =
            DocumentMarkupModel.forDocument(document, project, false)
    ): RangeHighlighterEx? {
        if (!ApplicationManager.getApplication().isDispatchThread) return null
        return markupModel?.let {
            textRange ?: return null
            it.addRangeHighlighter(
                textRange.startOffset,
                textRange.endOffset,
                HighlighterLayer.LAST + 1,
                ISSUE_HIGHLIGHT_TEXT_ATTRIBUTES,
                HighlighterTargetArea.EXACT_RANGE
            ) as RangeHighlighterEx
        }
    }

    private fun toTextRange(): TextRange? {
        if (startLine < 1 || endLine > document.lineCount) return null
        val startOffset = document.getLineStartOffset(startLine - 1) + startCol - 1
        val endOffset = document.getLineStartOffset(endLine - 1) + endCol - 1
        if (startOffset < 0 || endOffset > document.textLength || startOffset > endOffset) return null
        return TextRange.create(startOffset, endOffset)
    }
}


// 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.icons.AllIcons
import com.intellij.ide.BrowserUtil
import com.intellij.openapi.actionSystem.ActionGroup
import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.actionSystem.ActionToolbar
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.ui.AnimatedIcon
import com.intellij.ui.OnePixelSplitter
import com.intellij.ui.ScrollPaneFactory
import com.intellij.ui.border.CustomLineBorder
import com.intellij.ui.components.ActionLink
import com.intellij.ui.treeStructure.Tree
import com.intellij.util.ui.JBUI
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.listeners.CodeWhispererCodeScanTreeMouseListener
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.layout.CodeWhispererLayoutConfig.addHorizontalGlue
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.INACTIVE_TEXT_COLOR
import software.aws.toolkits.resources.message
import java.awt.BorderLayout
import java.awt.Component
import java.awt.GridBagConstraints
import java.awt.GridBagLayout
import java.net.URI
import java.time.Instant
import java.time.format.DateTimeFormatter
import javax.swing.BorderFactory
import javax.swing.JButton
import javax.swing.JLabel
import javax.swing.JPanel
import javax.swing.JTree
import javax.swing.tree.DefaultMutableTreeNode
import javax.swing.tree.TreeCellRenderer

/**
 * Create a Code Scan results view that displays the code scan results.
 */
internal class CodeWhispererCodeScanResultsView(private val project: Project) : JPanel(BorderLayout()) {

    private val codeScanTree: Tree = Tree().apply {
        isRootVisible = false
        CodeWhispererCodeScanTreeMouseListener(project).installOn(this)
        cellRenderer = ColoredTreeCellRenderer()
    }

    private val scrollPane = ScrollPaneFactory.createScrollPane(codeScanTree, true)
    private val splitter = OnePixelSplitter(CODE_SCAN_SPLITTER_PROPORTION_KEY, 1.0f).apply {
        firstComponent = scrollPane
    }

    private val toolbar = createToolbar().apply {
        setTargetComponent(this@CodeWhispererCodeScanResultsView)
        component.border = BorderFactory.createCompoundBorder(
            CustomLineBorder(JBUI.insetsRight(1)),
            component.border
        )
    }

    private val infoLabelInitialText = message("codewhisperer.codescan.run_scan_info")
    private val infoLabelPrefix = JLabel(infoLabelInitialText, JLabel.LEFT).apply {
        icon = AllIcons.General.BalloonInformation
    }
    private val scannedFilesLabelLink = ActionLink().apply {
        border = BorderFactory.createEmptyBorder(0, 7, 0, 0)
    }
    private val learnMoreLabelLink = ActionLink().apply {
        border = BorderFactory.createEmptyBorder(0, 7, 0, 0)
    }

    private val completeInfoLabel = JPanel(GridBagLayout()).apply {
        layout = GridBagLayout()
        border = BorderFactory.createCompoundBorder(
            CustomLineBorder(JBUI.insetsBottom(1)),
            BorderFactory.createEmptyBorder(7, 11, 8, 11)
        )
        add(infoLabelPrefix, CodeWhispererLayoutConfig.inlineLabelConstraints)
        add(scannedFilesLabelLink, CodeWhispererLayoutConfig.inlineLabelConstraints)
        add(learnMoreLabelLink, CodeWhispererLayoutConfig.inlineLabelConstraints)
        addHorizontalGlue()
    }

    private val progressIndicatorLabel = JLabel(message("codewhisperer.codescan.scan_in_progress"), AnimatedIcon.Default(), JLabel.CENTER).apply {
        border = BorderFactory.createEmptyBorder(7, 7, 7, 7)
    }
    private val stopCodeScanButton = JButton(message("codewhisperer.codescan.stop_scan")).apply {
        addActionListener {
            CodeWhispererCodeScanManager.getInstance(project).stopCodeScan()
        }
    }

    private val progressIndicator = JPanel(GridBagLayout()).apply {
        add(progressIndicatorLabel, GridBagConstraints())
        add(stopCodeScanButton, GridBagConstraints().apply { gridy = 1 })
    }

    // Results panel containing info label and progressIndicator/scrollPane
    private val resultsPanel = JPanel(BorderLayout()).apply {
        add(BorderLayout.NORTH, completeInfoLabel)
    }

    init {
        add(BorderLayout.WEST, toolbar.component)
        add(BorderLayout.CENTER, resultsPanel)
    }

    fun getCodeScanTree() = codeScanTree

    /**
     * Updates the [codeScanTree] with the new tree model root and displays the same on the UI.
     */
    fun updateAndDisplayScanResults(scanTreeModel: CodeWhispererCodeScanTreeModel, scannedFiles: List<VirtualFile>, isProjectTruncated: Boolean) {
        codeScanTree.apply {
            model = scanTreeModel
            repaint()
        }

        scannedFilesLabelLink.addActionListener {
            showScannedFiles(scannedFiles)
        }

        if (isProjectTruncated) {
            learnMoreLabelLink.addActionListener {
                // TODO: Change this URL to point to updated security scan documentation
                BrowserUtil.browse(URI("https://docs.aws.amazon.com/codewhisperer/latest/userguide/security-scans.html"))
            }
        }

        resultsPanel.apply {
            if (components.contains(progressIndicator)) remove(progressIndicator)
            add(BorderLayout.CENTER, splitter)
            splitter.proportion = 1.0f
            splitter.secondComponent = null
            revalidate()
            repaint()
        }

        changeInfoLabelToDisplayScanCompleted(scannedFiles.size, isProjectTruncated)
    }

    fun setStoppingCodeScan() {
        completeInfoLabel.isVisible = false
        stopCodeScanButton.isVisible = false
        resultsPanel.apply {
            if (components.contains(splitter)) remove(splitter)
            progressIndicatorLabel.apply {
                text = message("codewhisperer.codescan.stopping_scan")
                revalidate()
                repaint()
            }
            add(BorderLayout.CENTER, progressIndicator)
            revalidate()
            repaint()
        }
    }

    fun setDefaultUI() {
        completeInfoLabel.isVisible = true
        infoLabelPrefix.apply {
            text = infoLabelInitialText
            icon = AllIcons.General.BalloonInformation
            isVisible = true
        }
        scannedFilesLabelLink.apply {
            text = ""
            isVisible = false
        }
        learnMoreLabelLink.apply {
            text = ""
            isVisible = false
        }
        resultsPanel.apply {
            removeAll()
            add(BorderLayout.NORTH, completeInfoLabel)
            revalidate()
            repaint()
        }
    }

    /**
     * Shows in progress indicator indicating that the scan is in progress.
     */
    fun showInProgressIndicator() {
        completeInfoLabel.isVisible = false

        stopCodeScanButton.isVisible = true
        progressIndicatorLabel.text = message("codewhisperer.codescan.scan_in_progress")
        resultsPanel.apply {
            if (components.contains(splitter)) remove(splitter)
            add(BorderLayout.CENTER, progressIndicator)
            revalidate()
            repaint()
        }
    }

    /**
     * Sets info label to show error in case a runtime exception is encountered while running a code scan.
     */
    fun showError(errorMsg: String) {
        completeInfoLabel.isVisible = true
        infoLabelPrefix.apply {
            text = errorMsg
            icon = AllIcons.General.Error
            isVisible = true
            repaint()
        }
        scannedFilesLabelLink.text = ""
        learnMoreLabelLink.text = ""
        resultsPanel.apply {
            if (components.contains(splitter)) remove(splitter)
            if (components.contains(progressIndicator)) remove(progressIndicator)
            revalidate()
            repaint()
        }
    }
    private fun showScannedFiles(files: List<VirtualFile>) {
        val scannedFilesViewPanel = CodeWhispererCodeScanHighlightingFilesPanel(project, files)
        scannedFilesViewPanel.apply {
            isVisible = true
            revalidate()
        }
        splitter.apply {
            secondComponent = scannedFilesViewPanel
            proportion = 0.5f
            revalidate()
            repaint()
        }
    }

    private fun changeInfoLabelToDisplayScanCompleted(numScannedFiles: Int, isProjectTruncated: Boolean) {
        completeInfoLabel.isVisible = true
        infoLabelPrefix.icon = AllIcons.Actions.Commit
        infoLabelPrefix.text = message(
            "codewhisperer.codescan.run_scan_complete",
            numScannedFiles,
            (codeScanTree.model as CodeWhispererCodeScanTreeModel).getTotalIssuesCount(),
            project.name,
            if (isProjectTruncated) 1 else 0,
            INACTIVE_TEXT_COLOR,
            DateTimeFormatter.ISO_INSTANT.format(Instant.now())
        )
        infoLabelPrefix.repaint()
        infoLabelPrefix.isVisible = true
        scannedFilesLabelLink.text = message("codewhisperer.codescan.view_scanned_files", numScannedFiles)
        scannedFilesLabelLink.isVisible = true
        if (isProjectTruncated) {
            learnMoreLabelLink.apply {
                text = message("aws.settings.learn_more")
                isVisible = true
            }
        }
    }

    private fun createToolbar(): ActionToolbar {
        val actionManager = ActionManager.getInstance()
        val group = actionManager.getAction("aws.toolkit.codewhisperer.toolbar.security") as ActionGroup
        return actionManager.createActionToolbar(ACTION_PLACE, group, false)
    }

    private class ColoredTreeCellRenderer : TreeCellRenderer {
        override fun getTreeCellRendererComponent(
            tree: JTree?,
            value: Any?,
            selected: Boolean,
            expanded: Boolean,
            leaf: Boolean,
            row: Int,
            hasFocus: Boolean
        ): Component {
            value as DefaultMutableTreeNode
            val cell = JLabel()
            synchronized(value) {
                when (val obj = value.userObject) {
                    is VirtualFile -> {
                        cell.text = message("codewhisperer.codescan.file_name_issues_count", obj.name, obj.path, value.childCount, INACTIVE_TEXT_COLOR)
                        cell.icon = obj.fileType.icon
                    }
                    is CodeWhispererCodeScanIssue -> {
                        val cellText = "${obj.title}: ${obj.description.text}"
                        if (obj.isInvalid) {
                            cell.text = message("codewhisperer.codescan.scan_recommendation_invalid", obj.title, obj.displayTextRange(), INACTIVE_TEXT_COLOR)
                            cell.toolTipText = message("codewhisperer.codescan.scan_recommendation_invalid.tooltip_text")
                            cell.icon = AllIcons.General.Information
                        } else {
                            cell.text = message("codewhisperer.codescan.scan_recommendation", cellText, obj.displayTextRange(), INACTIVE_TEXT_COLOR)
                            cell.toolTipText = cellText
                            cell.icon = obj.issueSeverity.icon
                        }
                    }
                }
            }
            return cell
        }
    }

    private companion object {
        const val ACTION_PLACE = "CodeScanResultsPanel"
        const val CODE_SCAN_SPLITTER_PROPORTION_KEY = "CODE_SCAN_SPLITTER_PROPORTION"
    }
}


// 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.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.fasterxml.jackson.module.kotlin.readValue
import com.intellij.openapi.application.ApplicationInfo
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.util.TimeoutUtil.sleep
import com.intellij.util.io.HttpRequests
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.time.withTimeout
import org.apache.commons.codec.digest.DigestUtils
import software.amazon.awssdk.services.codewhisperer.model.ArtifactType
import software.amazon.awssdk.services.codewhisperer.model.CodeScanFindingsSchema
import software.amazon.awssdk.services.codewhisperer.model.CodeScanStatus
import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanRequest
import software.amazon.awssdk.services.codewhisperer.model.GetCodeScanResponse
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsRequest
import software.amazon.awssdk.services.codewhisperer.model.ListCodeScanFindingsResponse
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
import software.amazon.awssdk.utils.IoUtils
import software.aws.toolkits.core.utils.Waiters.waitUntil
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.core.AwsClientManager
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanResponseContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanServiceInvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODE_SCAN_POLLING_INTERVAL_IN_SECONDS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_BYTES_IN_KB
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_MILLIS_IN_SECOND
import software.aws.toolkits.jetbrains.utils.assertIsNonDispatchThread
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.File
import java.io.FileInputStream
import java.io.IOException
import java.net.HttpURLConnection
import java.nio.file.Path
import java.time.Duration
import java.time.Instant
import java.util.Base64
import java.util.UUID
import kotlin.coroutines.coroutineContext

class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
    private val clientToken: UUID = UUID.randomUUID()
    private val urlResponse = mutableMapOf<ArtifactType, CreateUploadUrlResponse>()

    private val clientAdaptor = CodeWhispererClientAdaptor.getInstance(sessionContext.project)

    private fun now() = Instant.now().toEpochMilli()

    /**
     * Note that this function makes network calls and needs to be run from a background thread.
     * Runs a code scan session which comprises the following steps:
     *  1. Generate truncation (zip files) based on the truncation in the session context.
     *  2. CreateUploadURL to upload the context.
     *  3. Upload the zip files using the URL
     *  4. Call createCodeScan to start a code scan
     *  5. Keep polling the API GetCodeScan to wait for results for a given timeout period.
     *  6. Return the results from the ListCodeScan API.
     */
    suspend fun run(): CodeScanResponse {
        var issues: List<CodeWhispererCodeScanIssue> = listOf()
        var codeScanResponseContext = defaultCodeScanResponseContext()
        val currentCoroutineContext = coroutineContext
        try {
            assertIsNonDispatchThread()
            currentCoroutineContext.ensureActive()
            val startTime = now()
            val (payloadContext, sourceZip) = withTimeout(Duration.ofSeconds(sessionContext.sessionConfig.createPayloadTimeoutInSeconds())) {
                runReadAction { sessionContext.sessionConfig.createPayload() }
            }

            LOG.debug {
                "Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
                    "Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
                    "Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
                    "Total number of lines scanned: ${payloadContext.totalLines} \n" +
                    "Total number of files included in payload: ${payloadContext.totalFiles} \n" +
                    "Total time taken for creating payload: ${payloadContext.totalTimeInMilliseconds * 1.0 / TOTAL_MILLIS_IN_SECOND} seconds\n" +
                    "Payload context language: ${payloadContext.language}"
            }
            codeScanResponseContext = codeScanResponseContext.copy(payloadContext = payloadContext)

            // 2 & 3. CreateUploadURL and upload the context.
            currentCoroutineContext.ensureActive()
            LOG.debug { "Uploading source zip located at ${sourceZip.path} to s3" }
            val artifactsUploadStartTime = now()
            val sourceZipUploadResponse = createUploadUrlAndUpload(sourceZip, "SourceCode")
            LOG.debug {
                "Successfully uploaded source zip to s3: " +
                    "Upload id: ${sourceZipUploadResponse.uploadId()} " +
                    "Request id: ${sourceZipUploadResponse.responseMetadata().requestId()}"
            }
            urlResponse[ArtifactType.SOURCE_CODE] = sourceZipUploadResponse
            currentCoroutineContext.ensureActive()
            val artifactsUploadDuration = now() - artifactsUploadStartTime
            codeScanResponseContext = codeScanResponseContext.copy(
                serviceInvocationContext = codeScanResponseContext.serviceInvocationContext.copy(artifactsUploadDuration = artifactsUploadDuration)
            )

            // 4. Call createCodeScan to start a code scan
            currentCoroutineContext.ensureActive()
            LOG.debug { "Requesting security scan for the uploaded artifacts, language: ${payloadContext.language}" }
            val serviceInvocationStartTime = now()
            val createCodeScanResponse = createCodeScan(payloadContext.language.toString())
            LOG.debug {
                "Successfully created security scan with " +
                    "status: ${createCodeScanResponse.status()} " +
                    "for request id: ${createCodeScanResponse.responseMetadata().requestId()}"
            }
            var codeScanStatus = createCodeScanResponse.status()
            if (codeScanStatus == CodeScanStatus.FAILED) {
                LOG.debug {
                    "CodeWhisperer service error occurred. Something went wrong when creating a security scan: $createCodeScanResponse " +
                        "Status: ${createCodeScanResponse.status()} for request id: ${createCodeScanResponse.responseMetadata().requestId()}"
                }
                codeScanFailed()
            }
            val jobId = createCodeScanResponse.jobId()
            codeScanResponseContext = codeScanResponseContext.copy(codeScanJobId = jobId)

            // 5. Keep polling the API GetCodeScan to wait for results for a given timeout period.
            waitUntil(
                succeedOn = { codeScanStatus == CodeScanStatus.COMPLETED },
                maxDuration = Duration.ofSeconds(sessionContext.sessionConfig.overallJobTimeoutInSeconds())
            ) {
                currentCoroutineContext.ensureActive()
                val elapsedTime = (now() - startTime) * 1.0 / TOTAL_MILLIS_IN_SECOND
                LOG.debug { "Waiting for security scan to complete. Elapsed time: $elapsedTime sec." }
                val getCodeScanResponse = getCodeScan(jobId)
                codeScanStatus = getCodeScanResponse.status()
                LOG.debug {
                    "Get security scan status: ${getCodeScanResponse.status()}, " +
                        "request id: ${getCodeScanResponse.responseMetadata().requestId()}"
                }
                sleepThread()
                if (codeScanStatus == CodeScanStatus.FAILED) {
                    LOG.debug {
                        "CodeWhisperer service error occurred. Something went wrong fetching results for security scan: $getCodeScanResponse " +
                            "Status: ${getCodeScanResponse.status()} for request id: ${getCodeScanResponse.responseMetadata().requestId()}"
                    }
                    codeScanFailed()
                }
            }

            LOG.debug { "Security scan completed successfully by CodeWhisperer." }

            // 6. Return the results from the ListCodeScan API.
            currentCoroutineContext.ensureActive()
            LOG.debug { "Fetching results for the completed security scan" }
            var listCodeScanFindingsResponse = listCodeScanFindings(jobId)
            LOG.debug {
                "Successfully fetched results for security scan with " +
                    "request id: ${listCodeScanFindingsResponse.responseMetadata().requestId()}"
            }
            val serviceInvocationDuration = now() - serviceInvocationStartTime
            codeScanResponseContext = codeScanResponseContext.copy(
                serviceInvocationContext = codeScanResponseContext.serviceInvocationContext.copy(serviceInvocationDuration = serviceInvocationDuration)
            )

            val documents = mutableListOf<String>()
            documents.add(listCodeScanFindingsResponse.codeScanFindings())
            while (listCodeScanFindingsResponse.nextToken() != null) {
                documents.add(listCodeScanFindingsResponse.codeScanFindings())
                listCodeScanFindingsResponse = listCodeScanFindings(jobId)
            }
            LOG.debug { "Successfully fetched results for the security scan." }
            LOG.debug { "Code scan findings: ${listCodeScanFindingsResponse.codeScanFindings()}" }
            LOG.debug { "Rendering response to display security scan results." }
            currentCoroutineContext.ensureActive()
            issues = mapToCodeScanIssues(documents)
            codeScanResponseContext = codeScanResponseContext.copy(codeScanTotalIssues = issues.count())
            codeScanResponseContext = codeScanResponseContext.copy(reason = "Succeeded")
            return CodeScanResponse.Success(issues, codeScanResponseContext)
        } catch (e: Exception) {
            val errorCode = (e as? CodeWhispererException)?.awsErrorDetails()?.errorCode()
            val requestId = if (e is CodeWhispererException) e.requestId() else null
            LOG.error {
                "Failed to run security scan and display results. Caused by: ${e.message}, status code: $errorCode, " +
                    "exception: ${e::class.simpleName}, request ID: $requestId " +
                    "Jetbrains IDE: ${ApplicationInfo.getInstance().fullApplicationName}, " +
                    "IDE version: ${ApplicationInfo.getInstance().apiVersion}, " +
                    "stacktrace: ${e.stackTrace.contentDeepToString()}"
            }
            return CodeScanResponse.Failure(issues, codeScanResponseContext, e)
        }
    }

    /**
     * Creates an upload URL and uplaods the zip file to the presigned URL
     */
    fun createUploadUrlAndUpload(zipFile: File, artifactType: String): CreateUploadUrlResponse = try {
        val fileMd5: String = Base64.getEncoder().encodeToString(DigestUtils.md5(FileInputStream(zipFile)))
        LOG.debug { "Fetching presigned URL for uploading $artifactType." }
        val createUploadUrlResponse = createUploadUrl(fileMd5, artifactType)
        LOG.debug { "Successfully fetched presigned URL for uploading $artifactType." }
        val url = createUploadUrlResponse.uploadUrl()
        LOG.debug { "Uploading $artifactType using the presigned URL." }
        uploadArtifactToS3(url, createUploadUrlResponse.uploadId(), zipFile, fileMd5, createUploadUrlResponse.kmsKeyArn())
        createUploadUrlResponse
    } catch (e: Exception) {
        LOG.debug { "Security scan failed. Something went wrong uploading artifacts: ${e.message}" }
        throw e
    }

    fun createUploadUrl(md5Content: String, artifactType: String): CreateUploadUrlResponse = clientAdaptor.createUploadUrl(
        CreateUploadUrlRequest.builder()
            .contentMd5(md5Content)
            .artifactType(artifactType)
            .build()
    )

    @Throws(IOException::class)
    fun uploadArtifactToS3(url: String, uploadId: String, fileToUpload: File, md5: String, kmsArn: String?) {
        val uploadIdJson = """{"uploadId":"$uploadId"}"""
        HttpRequests.put(url, "application/zip").userAgent(AwsClientManager.userAgent).tuner {
            it.setRequestProperty(CONTENT_MD5, md5)
            it.setRequestProperty(SERVER_SIDE_ENCRYPTION, AWS_KMS)
            it.setRequestProperty(CONTENT_TYPE, APPLICATION_ZIP)
            if (kmsArn?.isNotEmpty() == true) {
                it.setRequestProperty(SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID, kmsArn)
            }
            it.setRequestProperty(SERVER_SIDE_ENCRYPTION_CONTEXT, Base64.getEncoder().encodeToString(uploadIdJson.toByteArray()))
        }.connect {
            val connection = it.connection as HttpURLConnection
            connection.setFixedLengthStreamingMode(fileToUpload.length())
            IoUtils.copy(fileToUpload.inputStream(), connection.outputStream)
        }
    }

    fun createCodeScan(language: String): CreateCodeScanResponse {
        val artifactsMap = mapOf(
            ArtifactType.SOURCE_CODE to urlResponse[ArtifactType.SOURCE_CODE]?.uploadId(),
            ArtifactType.BUILT_JARS to urlResponse[ArtifactType.BUILT_JARS]?.uploadId()
        ).filter { (_, v) -> v != null }

        try {
            return clientAdaptor.createCodeScan(
                CreateCodeScanRequest.builder()
                    .clientToken(clientToken.toString())
                    .programmingLanguage { it.languageName(language) }
                    .artifacts(artifactsMap)
                    .build()
            )
        } catch (e: Exception) {
            LOG.debug { "Creating security scan failed: ${e.message}" }
            throw e
        }
    }

    fun getCodeScan(jobId: String): GetCodeScanResponse = try {
        clientAdaptor.getCodeScan(
            GetCodeScanRequest.builder()
                .jobId(jobId)
                .build()
        )
    } catch (e: Exception) {
        LOG.debug { "Getting security scan failed: ${e.message}" }
        throw e
    }

    fun listCodeScanFindings(jobId: String): ListCodeScanFindingsResponse = try {
        clientAdaptor.listCodeScanFindings(
            ListCodeScanFindingsRequest.builder()
                .jobId(jobId)
                .codeScanFindingsSchema(CodeScanFindingsSchema.CODESCAN_FINDINGS_1_0)
                .build()
        )
    } catch (e: Exception) {
        LOG.debug { "Listing security scan failed: ${e.message}" }
        throw e
    }

    fun mapToCodeScanIssues(recommendations: List<String>): List<CodeWhispererCodeScanIssue> {
        val scanRecommendations: List<CodeScanRecommendation> = recommendations.map {
            val value: List<CodeScanRecommendation> = MAPPER.readValue(it)
            value
        }.flatten()
        LOG.debug { "Total code scan issues returned from service: ${scanRecommendations.size}" }
        return scanRecommendations.mapNotNull {
            val file = try {
                LocalFileSystem.getInstance().findFileByIoFile(
                    Path.of(sessionContext.sessionConfig.projectRoot.path, it.filePath).toFile()
                )
            } catch (e: Exception) {
                LOG.debug { "Cannot find file at location ${it.filePath}" }
                null
            }
            when (file?.isDirectory) {
                false -> {
                    runReadAction {
                        FileDocumentManager.getInstance().getDocument(file)
                    }?.let { document ->
                        val endCol = document.getLineEndOffset(it.endLine - 1) - document.getLineStartOffset(it.endLine - 1) + 1
                        CodeWhispererCodeScanIssue(
                            startLine = it.startLine,
                            startCol = 1,
                            endLine = it.endLine,
                            endCol = endCol,
                            file = file,
                            project = sessionContext.project,
                            title = it.title,
                            description = it.description
                        )
                    }
                }
                else -> null
            }
        }.onEach { issue ->
            // Add range highlighters for all the issues found.
            runInEdt {
                issue.rangeHighlighter = issue.addRangeHighlighter()
            }
        }
    }

    fun sleepThread() {
        sleep(CODE_SCAN_POLLING_INTERVAL_IN_SECONDS * TOTAL_MILLIS_IN_SECOND)
    }

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanSession>()
        private val MAPPER = jacksonObjectMapper()
            .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
        const val AWS_KMS = "aws:kms"
        const val APPLICATION_ZIP = "application/zip"
        const val CONTENT_MD5 = "Content-MD5"
        const val CONTENT_TYPE = "Content-Type"
        const val SERVER_SIDE_ENCRYPTION = "x-amz-server-side-encryption"
        const val SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID = "x-amz-server-side-encryption-aws-kms-key-id"
        const val SERVER_SIDE_ENCRYPTION_CONTEXT = "x-amz-server-side-encryption-context"
    }
}

sealed class CodeScanResponse {
    abstract val issues: List<CodeWhispererCodeScanIssue>
    abstract val responseContext: CodeScanResponseContext

    data class Success(
        override val issues: List<CodeWhispererCodeScanIssue>,
        override val responseContext: CodeScanResponseContext
    ) : CodeScanResponse()

    data class Failure(
        override val issues: List<CodeWhispererCodeScanIssue>,
        override val responseContext: CodeScanResponseContext,
        val failureReason: Throwable
    ) : CodeScanResponse()
}

internal data class CodeScanRecommendation(
    val filePath: String,
    val startLine: Int,
    val endLine: Int,
    val title: String,
    val description: Description
)

data class Description(val text: String, val markdown: String)

data class CodeScanSessionContext(
    val project: Project,
    val sessionConfig: CodeScanSessionConfig
)

internal fun defaultPayloadContext() = PayloadContext(CodewhispererLanguage.Unknown, 0, 0, 0, listOf(), 0, 0)

internal fun defaultServiceInvocationContext() = CodeScanServiceInvocationContext(0, 0)

internal fun defaultCodeScanResponseContext() = CodeScanResponseContext(defaultPayloadContext(), defaultServiceInvocationContext())


// 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 javax.swing.tree.DefaultMutableTreeNode
import javax.swing.tree.DefaultTreeModel

internal class CodeWhispererCodeScanTreeModel(
    private val codeScanTreeNodeRoot: DefaultMutableTreeNode = DefaultMutableTreeNode("CodeWhisperer security scan results")
) : DefaultTreeModel(codeScanTreeNodeRoot) {

    override fun getRoot(): Any = codeScanTreeNodeRoot

    override fun getChild(parent: Any?, index: Int): Any {
        parent as DefaultMutableTreeNode
        return synchronized(parent) { parent.getChildAt(index) }
    }

    override fun getChildCount(parent: Any?): Int {
        parent as DefaultMutableTreeNode
        return synchronized(parent) { parent.childCount }
    }

    override fun isLeaf(node: Any?): Boolean {
        node as DefaultMutableTreeNode
        return synchronized(node) { node.isLeaf }
    }

    override fun getIndexOfChild(parent: Any?, child: Any?): Int {
        parent as DefaultMutableTreeNode
        child as DefaultMutableTreeNode
        return synchronized(parent) { parent.getIndex(child) }
    }

    fun getTotalIssuesCount(): Int = synchronized(codeScanTreeNodeRoot) {
        codeScanTreeNodeRoot.children().asSequence().sumBy { it.childCount }
    }
}


// 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.actions

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.resources.message

class CodeWhispererCodeScanRunAction : DumbAwareAction(
    message("codewhisperer.codescan.run_scan"),
    null,
    AllIcons.Actions.Execute
) {
    override fun update(event: AnActionEvent) {
        val project = event.project ?: return
        event.presentation.isEnabledAndVisible = isCodeWhispererEnabled(project)
        val scanManager = CodeWhispererCodeScanManager.getInstance(project)
        event.presentation.icon = scanManager.getRunActionButtonIcon()
    }

    override fun actionPerformed(event: AnActionEvent) {
        val project = event.project ?: return
        CodeWhispererCodeScanManager.getInstance(project).runCodeScan()
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.project.DumbAwareAction
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.resources.message

class CodeWhispererStopCodeScanAction : DumbAwareAction(
    message("codewhisperer.codescan.stop_scan"),
    null,
    AllIcons.Actions.Suspend
) {
    override fun update(event: AnActionEvent) {
        val project = event.project ?: return
        val scanManager = CodeWhispererCodeScanManager.getInstance(project)
        event.presentation.isEnabled = scanManager.isCodeScanJobActive()
    }

    override fun actionPerformed(event: AnActionEvent) {
        val project = event.project ?: return
        CodeWhispererCodeScanManager.getInstance(project).stopCodeScan()
    }
}


// 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.listeners

import com.intellij.openapi.editor.event.DocumentEvent
import com.intellij.openapi.editor.event.DocumentListener
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.TextRange
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanIssue
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import javax.swing.tree.TreePath

internal class CodeWhispererCodeScanDocumentListener(val project: Project) : DocumentListener {

    override fun documentChanged(event: DocumentEvent) {
        val scanManager = CodeWhispererCodeScanManager.getInstance(project)
        val treeModel = scanManager.getScanTree().model
        val file = FileDocumentManager.getInstance().getFile(event.document)
        if (file == null) {
            LOG.error { "Cannot find file for document ${event.document}" }
            return
        }
        val editedTextRange = TextRange.create(event.offset, event.offset + event.oldLength)
        val nodes = scanManager.getOverlappingScanNodes(file, editedTextRange)
        nodes.forEach {
            val issue = it.userObject as CodeWhispererCodeScanIssue
            synchronized(it) {
                treeModel.valueForPathChanged(TreePath(it.path), issue.copy(isInvalid = true))
            }
            issue.rangeHighlighter?.dispose()
            issue.rangeHighlighter?.textAttributes = null
        }
        scanManager.updateScanNodes(file)
    }

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanDocumentListener>()
    }
}


// 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.listeners

import com.intellij.ide.BrowserUtil
import com.intellij.openapi.editor.colors.EditorFontType
import com.intellij.openapi.editor.event.EditorMouseEvent
import com.intellij.openapi.editor.event.EditorMouseEventArea
import com.intellij.openapi.editor.event.EditorMouseMotionListener
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.openapi.ui.popup.JBPopupFactory
import com.intellij.ui.awt.RelativePoint
import com.intellij.ui.components.JBScrollPane
import com.intellij.ui.popup.AbstractPopup
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanIssue
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import software.aws.toolkits.jetbrains.utils.convertMarkdownToHTML
import java.awt.Dimension
import javax.swing.BorderFactory
import javax.swing.JEditorPane
import javax.swing.ScrollPaneConstants
import javax.swing.event.HyperlinkEvent

class CodeWhispererCodeScanEditorMouseMotionListener(private val project: Project) : EditorMouseMotionListener {
    /**
     * Current context for popup is still being shown.
     */
    private var currentPopupContext: ScanIssuePopupContext? = null

    private fun hidePopup() {
        currentPopupContext?.popup?.cancel()
        currentPopupContext = null
    }

    private fun showPopup(issue: CodeWhispererCodeScanIssue?, e: EditorMouseEvent) {
        if (issue == null) {
            LOG.debug {
                "Unable to show popup issue at ${e.logicalPosition} as the issue was null"
            }
            return
        }
        val description = convertMarkdownToHTML(issue.description.markdown)

        val editorPane = JEditorPane("text/html", description).apply {
            putClientProperty(JEditorPane.HONOR_DISPLAY_PROPERTIES, true)
            border = BorderFactory.createCompoundBorder(
                BorderFactory.createEmptyBorder(),
                BorderFactory.createEmptyBorder(7, 11, 8, 11)
            )
            font = e.editor.colorsScheme.getFont(EditorFontType.PLAIN)
            isEditable = false
            addHyperlinkListener { he ->
                if (he.eventType == HyperlinkEvent.EventType.ACTIVATED) {
                    BrowserUtil.browse(he.url)
                }
            }
        }
        val scrollPane = JBScrollPane(editorPane).apply {
            verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
            horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
            preferredSize = Dimension(480, 150)
        }

        val popup = JBPopupFactory.getInstance().createComponentPopupBuilder(scrollPane, null).setFocusable(true)
            .setTitle(issue.title)
            .createPopup()
        // Set the currently shown issue popup context as this issue
        popup.size = (popup as AbstractPopup).preferredContentSize
        popup.content.apply {
            size = preferredSize
        }

        currentPopupContext = ScanIssuePopupContext(issue, popup)

        popup.show(RelativePoint(e.mouseEvent))
    }

    override fun mouseMoved(e: EditorMouseEvent) {
        val scanManager = CodeWhispererCodeScanManager.getInstance(project)
        if (e.area != EditorMouseEventArea.EDITING_AREA || !e.isOverText) {
            hidePopup()
            return
        }
        val offset = e.offset
        val file = FileDocumentManager.getInstance().getFile(e.editor.document)
        if (file == null) {
            LOG.error { "Cannot find file for the document ${e.editor.document}" }
            return
        }
        val issuesInRange = scanManager.getScanNodesInRange(file, offset).map {
            it.userObject as CodeWhispererCodeScanIssue
        }
        if (issuesInRange.isEmpty()) {
            hidePopup()
            return
        }
        if (issuesInRange.contains(currentPopupContext?.issue)) return

        // No popups should be visible at this point.
        hidePopup()
        // Show popup for only the first issue found.
        // Only add popup if the issue is still valid. If the issue has gone stale or invalid because
        // the user has made some edits, we don't need to show the popup for the stale or invalid issues.
        if (!issuesInRange.first().isInvalid) showPopup(issuesInRange.first(), e)
    }

    private data class ScanIssuePopupContext(val issue: CodeWhispererCodeScanIssue, val popup: JBPopup)

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanEditorMouseMotionListener>()
    }
}


// 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.listeners

import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.fileEditor.OpenFileDescriptor
import com.intellij.openapi.project.Project
import com.intellij.ui.DoubleClickListener
import com.intellij.ui.treeStructure.Tree
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanIssue
import java.awt.event.MouseEvent
import javax.swing.tree.DefaultMutableTreeNode

class CodeWhispererCodeScanTreeMouseListener(private val project: Project) : DoubleClickListener() {
    override fun onDoubleClick(e: MouseEvent): Boolean {
        val codeScanNode = (e.source as Tree).selectionPath?.lastPathComponent as? DefaultMutableTreeNode ?: return false
        var status = true
        synchronized(codeScanNode) {
            if (codeScanNode.userObject !is CodeWhispererCodeScanIssue) return false
            val codeScanIssue = codeScanNode.userObject as CodeWhispererCodeScanIssue
            val textRange = codeScanIssue.textRange ?: return false
            val startOffset = textRange.startOffset

            if (codeScanIssue.isInvalid) return false

            runInEdt {
                val editor = FileEditorManager.getInstance(project).openTextEditor(
                    OpenFileDescriptor(project, codeScanIssue.file, startOffset),
                    true
                )
                if (editor == null) {
                    LOG.error { "Cannot fetch editor for the file ${codeScanIssue.file.path}" }
                    status = false
                    return@runInEdt
                }
                // If the codeScanIssue is still valid and rangehighlighter was not added previously for some reason,
                // try adding again.
                if (!codeScanIssue.isInvalid && codeScanIssue.rangeHighlighter == null) {
                    codeScanIssue.rangeHighlighter = codeScanIssue.addRangeHighlighter(editor.markupModel)
                }
            }
        }
        return status
    }

    companion object {
        private val LOG = getLogger<CodeWhispererCodeScanTreeMouseListener>()
    }
}


// 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.sessionconfig

import com.intellij.openapi.project.Project
import com.intellij.openapi.project.guessProjectDir
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VFileProperty
import com.intellij.openapi.vfs.VfsUtil
import com.intellij.openapi.vfs.VirtualFile
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.time.withTimeout
import software.aws.toolkits.core.utils.createTemporaryZipFile
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.putNextEntry
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.fileFormatNotSupported
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.fileTooLarge
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODE_SCAN_CREATE_PAYLOAD_TIMEOUT_IN_SECONDS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_BYTES_IN_KB
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_BYTES_IN_MB
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.File
import java.nio.file.Files
import java.nio.file.Path
import java.time.Duration
import java.time.Instant
import kotlin.io.path.relativeTo

sealed class CodeScanSessionConfig(
    private val selectedFile: VirtualFile,
    private val project: Project
) {
    val projectRoot = project.guessProjectDir() ?: error("Cannot guess base directory for project ${project.name}")

    abstract val sourceExt: String

    private var isProjectTruncated = false

    /**
     * Timeout for the overall job - "Run Security Scan".
     */
    abstract fun overallJobTimeoutInSeconds(): Long

    abstract fun getPayloadLimitInBytes(): Int

    protected fun willExceedPayloadLimit(currentTotalFileSize: Long, currentFileSize: Long): Boolean {
        val exceedsLimit = currentTotalFileSize > getPayloadLimitInBytes() - currentFileSize
        isProjectTruncated = isProjectTruncated || exceedsLimit
        return exceedsLimit
    }

    open fun getImportedFiles(file: VirtualFile, includedSourceFiles: Set<String>): List<String> = listOf()

    open fun createPayload(): Payload {
        // Fail fast if the selected file size is greater than the payload limit.
        if (selectedFile.length > getPayloadLimitInBytes()) {
            fileTooLarge(getPresentablePayloadLimit())
        }

        val start = Instant.now().toEpochMilli()

        LOG.debug { "Creating payload. File selected as root for the context truncation: ${selectedFile.path}" }

        val (includedSourceFiles, payloadSize, totalLines, _) = includeDependencies()

        // Copy all the included source files to the source zip
        val srcZip = zipFiles(includedSourceFiles.map { Path.of(it) })
        val payloadContext = PayloadContext(
            selectedFile.programmingLanguage().toTelemetryType(),
            totalLines,
            includedSourceFiles.size,
            Instant.now().toEpochMilli() - start,
            includedSourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
            payloadSize,
            srcZip.length()
        )

        return Payload(payloadContext, srcZip)
    }

    open fun includeDependencies(): PayloadMetadata {
        val includedSourceFiles = mutableSetOf<String>()
        var currentTotalFileSize = 0L
        var currentTotalLines = 0L
        val files = getSourceFilesUnderProjectRoot(selectedFile)
        val queue = ArrayDeque<String>()

        files.forEach { pivotFile ->
            val filePath = pivotFile.path
            queue.addLast(filePath)

            // BFS
            while (queue.isNotEmpty()) {
                if (currentTotalFileSize.equals(getPayloadLimitInBytes())) {
                    return PayloadMetadata(includedSourceFiles, currentTotalFileSize, currentTotalLines)
                }

                val currentFilePath = queue.removeFirst()
                val currentFile = File(currentFilePath).toVirtualFile()
                if (includedSourceFiles.contains(currentFilePath) ||
                    currentFile == null ||
                    willExceedPayloadLimit(currentTotalFileSize, currentFile.length)
                ) {
                    continue
                }

                val currentFileSize = currentFile.length

                currentTotalFileSize += currentFileSize
                currentTotalLines += Files.lines(currentFile.toNioPath()).count()
                includedSourceFiles.add(currentFilePath)

                getImportedFiles(currentFile, includedSourceFiles).forEach {
                    if (!includedSourceFiles.contains(it)) queue.addLast(it)
                }
            }
        }

        return PayloadMetadata(includedSourceFiles, currentTotalFileSize, currentTotalLines)
    }

    /**
     * Timeout for creating the payload [createPayload]
     */
    open fun createPayloadTimeoutInSeconds(): Long = CODE_SCAN_CREATE_PAYLOAD_TIMEOUT_IN_SECONDS

    open fun getPresentablePayloadLimit(): String = when (getPayloadLimitInBytes() >= TOTAL_BYTES_IN_MB) {
        true -> "${getPayloadLimitInBytes() / TOTAL_BYTES_IN_MB}MB"
        false -> "${getPayloadLimitInBytes() / TOTAL_BYTES_IN_KB}KB"
    }

    open suspend fun getTotalProjectSizeInBytes(): Long {
        var totalSize = 0L
        try {
            withTimeout(Duration.ofSeconds(TELEMETRY_TIMEOUT_IN_SECONDS)) {
                VfsUtil.collectChildrenRecursively(projectRoot).filter {
                    !it.isDirectory && !it.`is`((VFileProperty.SYMLINK)) && it.path.endsWith(sourceExt)
                }.fold(0L) { acc, next ->
                    totalSize = acc + next.length
                    totalSize
                }
            }
        } catch (e: TimeoutCancellationException) {
            // Do nothing
        }
        return totalSize
    }

    protected fun zipFiles(files: List<Path>): File = createTemporaryZipFile {
        files.forEach { file ->
            val relativePath = file.relativeTo(projectRoot.toNioPath())
            LOG.debug { "Selected file for truncation: $file" }
            it.putNextEntry(relativePath.toString(), file)
        }
    }.toFile()

    /**
     * Returns all the source files for a given payload type.
     */
    open fun getSourceFilesUnderProjectRoot(selectedFile: VirtualFile): List<VirtualFile> {
        // Include the current selected file
        val files = mutableListOf(selectedFile)
        // Include other files only if the current file is in the project.
        if (selectedFile.path.startsWith(projectRoot.path)) {
            files.addAll(
                VfsUtil.collectChildrenRecursively(projectRoot).filter {
                    it.path.endsWith(sourceExt) && it != selectedFile
                }
            )
        }
        return files
    }

    open fun isProjectTruncated() = isProjectTruncated

    protected fun getPath(root: String, relativePath: String = ""): Path? = try {
        Path.of(root, relativePath).normalize()
    } catch (e: Exception) {
        LOG.debug { "Cannot find file at path $relativePath relative to the root $root" }
        null
    }

    protected fun File.toVirtualFile() = LocalFileSystem.getInstance().findFileByIoFile(this)

    companion object {
        private val LOG = getLogger<CodeScanSessionConfig>()
        private const val TELEMETRY_TIMEOUT_IN_SECONDS: Long = 10
        const val FILE_SEPARATOR = '/'
        fun create(file: VirtualFile, project: Project): CodeScanSessionConfig = when (file.programmingLanguage().toTelemetryType()) {
            CodewhispererLanguage.Java -> JavaCodeScanSessionConfig(file, project)
            CodewhispererLanguage.Python -> PythonCodeScanSessionConfig(file, project)
            CodewhispererLanguage.Javascript -> JavaScriptCodeScanSessionConfig(file, project)
            else -> fileFormatNotSupported(file.extension ?: "")
        }
    }
}

data class Payload(
    val context: PayloadContext,
    val srcZip: File
)

data class PayloadContext(
    val language: CodewhispererLanguage,
    val totalLines: Long,
    val totalFiles: Int,
    val totalTimeInMilliseconds: Long,
    val scannedFiles: List<VirtualFile>,
    val srcPayloadSize: Long,
    val srcZipFileSize: Long,
    val buildPayloadSize: Long? = null
)

data class PayloadMetadata(
    val sourceFiles: Set<String>,
    val payloadSize: Long,
    val linesScanned: Long,
    val buildPaths: Set<String> = setOf()
)