// 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.compiler.CompilerPaths
import com.intellij.openapi.module.ModuleManager
import com.intellij.openapi.module.ModuleUtil
import com.intellij.openapi.project.Project
import com.intellij.openapi.project.rootManager
import com.intellij.openapi.roots.ProjectRootManager
import com.intellij.openapi.vfs.VfsUtil
import com.intellij.openapi.vfs.VirtualFile
import org.jetbrains.jps.model.java.JavaSourceRootType
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.exists
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.cannotFindBuildArtifacts
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.fileTooLarge
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.JAVA_CODE_SCAN_TIMEOUT_IN_SECONDS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.JAVA_PAYLOAD_LIMIT_IN_BYTES
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.CodewhispererLanguage
import java.io.IOException
import java.lang.IndexOutOfBoundsException
import java.nio.file.Files
import java.nio.file.Path
import java.time.Instant

internal class JavaCodeScanSessionConfig(
    private val selectedFile: VirtualFile,
    private val project: Project
) : CodeScanSessionConfig(selectedFile, project) {

    private val packageRegex = Regex("package\\s+([\\w.]+)\\s*;")
    private val importRegex = Regex("import\\s+([\\w.]+[*]?)\\s*;")
    private val buildExt = ".class"
    override val sourceExt = ".java"

    data class JavaImportsInfo(val imports: List<String>, val packagePath: String)

    override fun overallJobTimeoutInSeconds(): Long = JAVA_CODE_SCAN_TIMEOUT_IN_SECONDS

    override fun getPayloadLimitInBytes(): Int = JAVA_PAYLOAD_LIMIT_IN_BYTES

    override 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}" }

        // Include all the dependencies using BFS
        val (sourceFiles, srcPayloadSize, totalLines, buildPaths) = includeDependencies()

        var noClassFilesFound = true
        val outputPaths = CompilerPaths.getOutputPaths(ModuleManager.getInstance(project).modules)
        var totalBuildPayloadSize = 0L
        val buildFiles = buildPaths.mapNotNull { relativePath ->
            val classFile = findClassFile(relativePath, outputPaths)
            if (classFile == null) {
                LOG.debug { "Cannot find class file for $relativePath" }
            } else {
                noClassFilesFound = false
                totalBuildPayloadSize += classFile.toFile().length()
            }
            classFile
        }
        LOG.debug { "Total build files sent in payload: ${buildFiles.size}" }
        if (noClassFilesFound) cannotFindBuildArtifacts()

        // Copy all the included source and build files to the source zip
        val srcZip = zipFiles(sourceFiles.mapNotNull { getPath(it) } + buildFiles)

        val payloadContext = PayloadContext(
            CodewhispererLanguage.Java,
            totalLines,
            sourceFiles.size,
            Instant.now().toEpochMilli() - start,
            sourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
            srcPayloadSize,
            srcZip.length(),
            totalBuildPayloadSize
        )
        return Payload(payloadContext, srcZip)
    }

    private fun findClassFile(relativePath: String, outputPaths: Array<String>): Path? {
        outputPaths.forEach { outputPath ->
            val classFile = getPath(outputPath, relativePath)
            if (classFile?.exists() == true) return classFile
        }
        return null
    }

    fun parseImports(file: VirtualFile): JavaImportsInfo {
        val imports = mutableSetOf<String>()
        val inputStream = file.inputStream
        var packagePath = ""
        try {
            inputStream.use {
                it.bufferedReader().lines().forEach { line ->
                    val importMatcher = importRegex.toPattern().matcher(line)
                    val packageMatcher = packageRegex.toPattern().matcher(line)
                    if (importMatcher.find()) {
                        val import = importMatcher.group(1).replace('.', FILE_SEPARATOR)
                        imports.add(import)
                    }
                    if (packageMatcher.find()) {
                        packagePath = packageMatcher.group(1).replace('.', FILE_SEPARATOR)
                    }
                }
            }
        } catch (e: IOException) {
            LOG.error { message("codewhisperer.codescan.cannot_read_file", file.path) }
        }
        return JavaImportsInfo(imports.toList(), packagePath)
    }

    /**
     * Gets a relative build path for file and package
     */
    private fun getRelativeBuildPath(file: VirtualFile, packagePath: String): String? {
        val sourceFilePath = file.path
        return if (packagePath.isEmpty()) {
            file.nameWithoutExtension
        } else {
            val start = sourceFilePath.lastIndexOf(packagePath)
            val end = sourceFilePath.lastIndexOf('.')
            try {
                sourceFilePath.substring(start, end)
            } catch (e: IndexOutOfBoundsException) {
                return null
            }
        } + buildExt
    }

    override fun getSourceFilesUnderProjectRoot(selectedFile: VirtualFile): List<VirtualFile> {
        val files = mutableListOf<VirtualFile>()
        val sourceRoots = ProjectRootManager.getInstance(project).getModuleSourceRoots(setOf(JavaSourceRootType.SOURCE))
        files.add(selectedFile)
        sourceRoots.forEach { vFile ->
            files.addAll(
                VfsUtil.collectChildrenRecursively(vFile).filter {
                    it.path.endsWith(sourceExt) && it != selectedFile
                }
            )
        }
        return files
    }

    override fun includeDependencies(): PayloadMetadata {
        val sourceFiles = mutableSetOf<VirtualFile>()
        val buildPaths = mutableSetOf<String?>()
        var currentTotalFileSize = 0L
        var currentTotalLines = 0L
        val files = getSourceFilesUnderProjectRoot(selectedFile)
        val queue = ArrayDeque<VirtualFile>()

        files.forEach { file ->
            queue.add(file)

            // BFS
            while (queue.isNotEmpty()) {
                if (currentTotalFileSize.equals(getPayloadLimitInBytes())) {
                    return PayloadMetadata(sourceFiles.map { it.path }.toSet(), currentTotalFileSize, currentTotalLines, buildPaths.filterNotNull().toSet())
                }

                val currentFile = queue.removeFirst()
                if (!currentFile.path.startsWith(projectRoot.path) ||
                    sourceFiles.contains(currentFile) ||
                    willExceedPayloadLimit(currentTotalFileSize, currentFile.length)
                ) {
                    if (!currentFile.path.startsWith(projectRoot.path)) {
                        LOG.error { "Invalid workspace: Current file ${currentFile.path} is not under the project root ${projectRoot.path}" }
                    }
                    continue
                }

                val currentFileSize = currentFile.length

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

                // Get all imports from the file
                val importsInfo = parseImports(currentFile)
                importsInfo.imports.forEach { importPath ->
                    val importedFiles = getSourceFilesForImport(currentFile, importPath)
                    importedFiles.forEach { importedFile ->
                        if (!sourceFiles.contains(importedFile)) queue.addLast(importedFile)
                    }
                }
                buildPaths.add(getRelativeBuildPath(currentFile, importsInfo.packagePath))
            }
        }

        return PayloadMetadata(sourceFiles.map { it.path }.toSet(), currentTotalFileSize, currentTotalLines, buildPaths.filterNotNull().toSet())
    }

    private fun getImportedFile(currentFile: VirtualFile, importPath: String): VirtualFile? {
        // Handle '*' imports
        val resolvedImportPath = if (importPath.contains('*')) {
            importPath.substring(0, importPath.indexOfFirst { it == '*' } - 1)
        } else {
            importPath + sourceExt
        }

        // First try searching the module containing the current file
        ModuleUtil.findModuleForFile(currentFile, project)?.rootManager?.getSourceRoots(JavaSourceRootType.SOURCE)?.forEach { srcRoot ->
            val path = getPath(srcRoot.path, resolvedImportPath)
            path?.toFile()?.toVirtualFile()?.let {
                return it
            }
        }

        // Fallback to all other java source roots.
        val projectSourceRoots = ProjectRootManager.getInstance(project).contentSourceRoots
        projectSourceRoots.forEach { srcRoot ->
            val path = getPath(srcRoot.path, resolvedImportPath)
            path?.toFile()?.toVirtualFile()?.let {
                return it
            }
        }
        return null
    }

    /**
     * Get source files for import statement. If the import is a star import, include all the files in the package directory.
     */
    private fun getSourceFilesForImport(currentFile: VirtualFile, importPath: String): List<VirtualFile> {
        val importedFile = getImportedFile(currentFile, importPath) ?: return listOf()
        if (!importedFile.isDirectory) {
            return listOf(importedFile)
        }
        return VfsUtil.collectChildrenRecursively(importedFile).filter { it.name.endsWith(sourceExt) }
    }

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


// 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.vfs.VirtualFile
import com.intellij.util.containers.addIfNotNull
import software.aws.toolkits.core.utils.exists
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.JS_CODE_SCAN_TIMEOUT_IN_SECONDS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.JS_PAYLOAD_LIMIT_IN_BYTES
import software.aws.toolkits.resources.message
import java.io.IOException
import java.nio.file.Path

internal class JavaScriptCodeScanSessionConfig(
    private val selectedFile: VirtualFile,
    private val project: Project
) : CodeScanSessionConfig(selectedFile, project) {

    private val importRegex = Regex("^import.*(?:[\"'](.+)[\"']);?\$")
    private val requireRegex = Regex("^.+require\\(['\"](.+)['\"]\\)[ \\t]*;?")
    override val sourceExt = ".js"

    override fun overallJobTimeoutInSeconds(): Long = JS_CODE_SCAN_TIMEOUT_IN_SECONDS

    override fun getPayloadLimitInBytes(): Int = JS_PAYLOAD_LIMIT_IN_BYTES

    fun parseImports(file: VirtualFile): List<String> {
        val imports = mutableSetOf<String>()
        try {
            file.inputStream.use {
                it.bufferedReader().lines().forEach { line ->
                    val importMatcher = importRegex.toPattern().matcher(line)
                    val moduleName = when (importMatcher.find()) {
                        true -> importMatcher.group(1)
                        false -> {
                            val requireMatcher = requireRegex.toPattern().matcher(line)
                            if (requireMatcher.find()) {
                                requireMatcher.group(1)
                            } else {
                                ""
                            }
                        }
                    }.trim()
                    if (moduleName.isNotEmpty()) {
                        when (moduleName.endsWith(sourceExt)) {
                            true -> imports.add(moduleName)
                            false -> {
                                imports.add(moduleName + sourceExt)
                            }
                        }
                    }
                }
            }
        } catch (e: IOException) {
            error(message("codewhisperer.codescan.cannot_read_file", file.path))
        }
        return imports.toList()
    }

    override fun getImportedFiles(file: VirtualFile, includedSourceFiles: Set<String>): List<String> {
        val importedFiles = mutableListOf<String>()
        val imports = parseImports(file)
        val importedFilePaths = mutableListOf<String>()
        imports.forEach { importPath ->
            if (getPath(importPath)?.exists() == true) {
                importedFilePaths.add(Path.of(importPath).normalize().toString())
            } else {
                val importedFilePath = getPath(file.parent.path, importPath)
                if (importedFilePath?.exists() == true) {
                    importedFilePaths.addIfNotNull(importedFilePath.toFile().toVirtualFile()?.path)
                }
            }
        }
        val validSourceFiles = importedFilePaths.filter { !includedSourceFiles.contains(it) }
        validSourceFiles.forEach { validFile ->
            importedFiles.add(validFile)
        }
        return importedFiles
    }
}


// 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.roots.ProjectRootManager
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.util.containers.addIfNotNull
import software.aws.toolkits.core.utils.exists
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.PYTHON_CODE_SCAN_TIMEOUT_IN_SECONDS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.PYTHON_PAYLOAD_LIMIT_IN_BYTES
import software.aws.toolkits.resources.message
import java.io.IOException

internal class PythonCodeScanSessionConfig(
    private val selectedFile: VirtualFile,
    private val project: Project
) : CodeScanSessionConfig(selectedFile, project) {

    private val importRegex = Regex("^(?:from\\s+(\\S+)\\s+)?(?:import\\s+((?:\\S+(?:\\s+as\\s+\\S+)?\\s*[,]?\\s*)+))\$")
    private val projectContentRoots = ProjectRootManager.getInstance(project).contentRoots
    override val sourceExt = ".py"

    override fun overallJobTimeoutInSeconds(): Long = PYTHON_CODE_SCAN_TIMEOUT_IN_SECONDS

    override fun getPayloadLimitInBytes(): Int = PYTHON_PAYLOAD_LIMIT_IN_BYTES

    fun parseImports(file: VirtualFile): List<String> {
        val imports = mutableSetOf<String>()
        try {
            file.inputStream.use {
                it.bufferedReader().lines().forEach { line ->
                    val importMatcher = importRegex.toPattern().matcher(line)
                    if (importMatcher.find()) {
                        // Group(1) is the 'from' module in the import statement.
                        // For E.g. in "from <Module1> import xyz", import module is Module1
                        val fromModule = importMatcher.group(1)?.plus(FILE_SEPARATOR) ?: ""
                        // Group(2) is the "<Module1> as <asName1>, <Module2> as <asName2>, <Module3>,..." statement
                        val importStatements = importMatcher.group(2)
                        importStatements.split(",").forEach { statement ->
                            // Just get the first word in <module> [as <name>] statement
                            val importModule = statement.trim().split(" ").first()
                            val importPath = fromModule + importModule.replace(".", FILE_SEPARATOR.toString()) + sourceExt
                            imports.add(importPath)
                        }
                    }
                }
            }
        } catch (e: IOException) {
            error(message("codewhisperer.codescan.cannot_read_file", file.path))
        }
        return imports.toList()
    }

    override fun getImportedFiles(file: VirtualFile, includedSourceFiles: Set<String>): List<String> {
        val importedFiles = mutableListOf<String>()
        val imports = parseImports(file)
        val importedFilePaths = mutableListOf<String>()
        projectContentRoots.forEach { root ->
            imports.forEach { importPath ->
                val importedFilePath = getPath(root.path, importPath)
                if (importedFilePath?.exists() == true) {
                    importedFilePaths.addIfNotNull(importedFilePath.toFile().toVirtualFile()?.path)
                }
            }
        }
        val validSourceFiles = importedFilePaths.filter { !includedSourceFiles.contains(it) }
        validSourceFiles.forEach { validFile ->
            importedFiles.add(validFile)
        }
        return importedFiles
    }
}


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

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

import com.intellij.openapi.Disposable
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClient
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.CodeWhispererRuntimeClient
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.AwsClientManager
import software.aws.toolkits.jetbrains.core.awsClient
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.CredentialManager.Companion.CREDENTIALS_CHANGED
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.transform
import kotlin.reflect.KProperty0
import kotlin.reflect.jvm.isAccessible

// TODO: move this file to package "/client"
// As the connection is project-level, we need to make this project-level too
interface CodeWhispererClientAdaptor : Disposable {
    val project: Project

    fun generateCompletionsPaginator(
        firstRequest: GenerateCompletionsRequest,
    ): Sequence<GenerateCompletionsResponse>

    fun createUploadUrl(
        request: CreateUploadUrlRequest
    ): CreateUploadUrlResponse

    fun createCodeScan(
        request: CreateCodeScanRequest,
        isSigv4: Boolean = shouldUseSigv4Client(project)
    ): CreateCodeScanResponse

    fun getCodeScan(
        request: GetCodeScanRequest,
        isSigv4: Boolean = shouldUseSigv4Client(project)
    ): GetCodeScanResponse

    fun listCodeScanFindings(
        request: ListCodeScanFindingsRequest,
        isSigv4: Boolean = shouldUseSigv4Client(project)
    ): ListCodeScanFindingsResponse

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

        private fun shouldUseSigv4Client(project: Project) =
            CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Accountless
    }
}

open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor {
    private val mySigv4Client by lazy { createUnmanagedSigv4Client() }

    @Volatile private var myBearerClient: CodeWhispererRuntimeClient? = null

    private val KProperty0<*>.isLazyInitialized: Boolean
        get() {
            isAccessible = true
            return (getDelegate() as Lazy<*>).isInitialized()
        }

    init {
        initClientUpdateListener()
    }

    private fun initClientUpdateListener() {
        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
            ToolkitConnectionManagerListener.TOPIC,
            object : ToolkitConnectionManagerListener {
                override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                    if (newConnection is AwsBearerTokenConnection) {
                        myBearerClient = getBearerClient(newConnection.getConnectionSettings().providerId)
                    }
                }
            }
        )
    }

    private fun bearerClient(): CodeWhispererRuntimeClient {
        if (myBearerClient != null) return myBearerClient as CodeWhispererRuntimeClient
        myBearerClient = getBearerClient()
        return myBearerClient as CodeWhispererRuntimeClient
    }

    override fun generateCompletionsPaginator(firstRequest: GenerateCompletionsRequest) = sequence<GenerateCompletionsResponse> {
        var nextToken: String? = firstRequest.nextToken()
        do {
            val response = bearerClient().generateCompletions(firstRequest.copy { it.nextToken(nextToken) })
            nextToken = response.nextToken()
            yield(response)
        } while (!nextToken.isNullOrEmpty())
    }

    override fun createUploadUrl(request: CreateUploadUrlRequest): CreateUploadUrlResponse =
        bearerClient().createUploadUrl(request)

    override fun createCodeScan(request: CreateCodeScanRequest, isSigv4: Boolean): CreateCodeScanResponse =
        if (isSigv4) {
            mySigv4Client.createCodeScan(request)
        } else {
            bearerClient().startCodeAnalysis(request.transform()).transform()
        }

    override fun getCodeScan(request: GetCodeScanRequest, isSigv4: Boolean): GetCodeScanResponse =
        if (isSigv4) {
            mySigv4Client.getCodeScan(request)
        } else {
            bearerClient().getCodeAnalysis(request.transform()).transform()
        }

    override fun listCodeScanFindings(request: ListCodeScanFindingsRequest, isSigv4: Boolean): ListCodeScanFindingsResponse =
        if (isSigv4) {
            mySigv4Client.listCodeScanFindings(request)
        } else {
            bearerClient().listCodeAnalysisFindings(request.transform()).transform()
        }

    override fun dispose() {
        if (this::mySigv4Client.isLazyInitialized) {
            mySigv4Client.close()
        }
        myBearerClient?.close()
    }

    /**
     * Every different SSO/AWS Builder ID connection requires a new client which has its corresponding bearer token provider,
     * thus we have to create them dynamically.
     * Invalidate and recycle the old client first, and create a new client with the new connection.
     * This makes sure when we invoke CW, we always use the up-to-date connection.
     * In case this fails to close the client, myBearerClient is already set to null thus next time when we invoke CW,
     * it will go through this again which should get the current up-to-date connection. This stale client would be
     * unused and stay in memory for a while until eventually closed by ToolkitClientManager.
     */
    open fun getBearerClient(oldProviderIdToRemove: String = ""): CodeWhispererRuntimeClient? {
        myBearerClient = null
        ApplicationManager.getApplication().messageBus.syncPublisher(CREDENTIALS_CHANGED)
            .providerRemoved(oldProviderIdToRemove)

        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
        connection as? AwsBearerTokenConnection ?: run {
            LOG.warn { "$connection is not a bearer token connection" }
            return null
        }

        return AwsClientManager.getInstance().getClient<CodeWhispererRuntimeClient>(connection.getConnectionSettings())
    }

    companion object {
        private val LOG = getLogger<CodeWhispererClientAdaptorImpl>()
        private fun createUnmanagedSigv4Client(): CodeWhispererClient = AwsClientManager.getInstance().createUnmanagedClient(
            AnonymousCredentialsProvider.create(),
            CodeWhispererConstants.Config.Sigv4ClientRegion,
            CodeWhispererConstants.Config.CODEWHISPERER_ENDPOINT
        )
    }
}

class MockCodeWhispererClientAdaptor(override val project: Project) : CodeWhispererClientAdaptorImpl(project) {
    override fun getBearerClient(oldProviderIdToRemove: String): CodeWhispererRuntimeClient = project.awsClient()

    override fun dispose() {}
}


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

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

import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.core.credentials.ConnectionDialogCustomizer
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAddConnectionDialog
import software.aws.toolkits.jetbrains.core.credentials.sono.CODEWHISPERER_SCOPES
import software.aws.toolkits.jetbrains.core.help.HelpIds
import software.aws.toolkits.resources.message

// TODO: doesn't need to be a class, but need to minimize file deltas
class CodeWhispererLoginDialog(project: Project) : ToolkitAddConnectionDialog(
    project,
    customizer = ConnectionDialogCustomizer(
        title = message("codewhisperer.credential.login.dialog.title"),
        header = message("codewhisperer.credential.login.dialog.prompt"),
        helpId = HelpIds.CODEWHISPERER_LOGIN_DIALOG,
        replaceIamComment = message("codewhisperer.credential.login.dialog.iam.description"),
        scopes = CODEWHISPERER_SCOPES
    )
)


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

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

enum class CodeWhispererLoginType(val displayName: String) {
    SSO("IAM Identity Center"),
    Sono("Builder ID"),
    Accountless("Access Code"),
    Logout("Logout"),
    Expired("Expired")
}


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

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

import com.intellij.openapi.editor.event.BulkAwareDocumentListener
import com.intellij.openapi.editor.event.DocumentEvent
import com.intellij.openapi.editor.event.EditorFactoryEvent
import com.intellij.openapi.editor.event.EditorFactoryListener
import com.intellij.openapi.editor.impl.EditorImpl
import com.intellij.psi.PsiDocumentManager
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererCodeCoverageTracker

class CodeWhispererEditorListener : EditorFactoryListener {
    override fun editorCreated(event: EditorFactoryEvent) {
        val editor = (event.editor as? EditorImpl) ?: return
        editor.project?.let { project ->
            PsiDocumentManager.getInstance(project).getPsiFile(editor.document)?.programmingLanguage() ?. let { language ->
                // If language is not supported by CodeWhisperer, no action needed
                if (!language.isCodeCompletionSupported()) return
                // If language is supported, install document listener for CodeWhisperer service
                editor.document.addDocumentListener(
                    object : BulkAwareDocumentListener {
                        override fun documentChanged(event: DocumentEvent) {
                            if (!isCodeWhispererEnabled(project)) return
                            CodeWhispererInvocationStatus.getInstance().documentChanged()
                            CodeWhispererCodeCoverageTracker.getInstance(language).apply {
                                activateTrackerIfNotActive()
                                documentChanged(event)
                            }
                        }
                    },
                    editor.disposable
                )
            }
        }
    }
}


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

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

import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.util.TextRange
import com.intellij.psi.PsiDocumentManager
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.PAIRED_BRACKETS
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.PAIRED_QUOTES
import java.time.Instant
import java.util.Stack

class CodeWhispererEditorManager {
    fun updateEditorWithRecommendation(states: InvocationContext, sessionContext: SessionContext) {
        val (requestContext, responseContext, recommendationContext) = states
        val (project, editor) = requestContext
        val document = editor.document
        val primaryCaret = editor.caretModel.primaryCaret
        val selectedIndex = sessionContext.selectedIndex
        val typeahead = sessionContext.typeahead
        val detail = recommendationContext.details[selectedIndex]
        val reformatted = CodeWhispererPopupManager.getInstance().getReformattedRecommendation(
            detail,
            recommendationContext.userInputSinceInvocation
        )
        val remainingRecommendation = reformatted.substring(typeahead.length)
        val originalOffset = primaryCaret.offset - typeahead.length

        val endOffset = primaryCaret.offset + remainingRecommendation.length

        val insertEndOffset = sessionContext.insertEndOffset
        val endOffsetToReplace = if (insertEndOffset != -1) insertEndOffset else primaryCaret.offset

        WriteCommandAction.runWriteCommandAction(project) {
            document.replaceString(originalOffset, endOffsetToReplace, reformatted)
            PsiDocumentManager.getInstance(project).commitDocument(document)
            primaryCaret.moveToOffset(endOffset + detail.rightOverlap.length)
        }

        ApplicationManager.getApplication().invokeLater {
            WriteCommandAction.runWriteCommandAction(project) {
                val rangeMarker = document.createRangeMarker(originalOffset, endOffset, true)

                CodeWhispererTelemetryService.getInstance().enqueueAcceptedSuggestionEntry(
                    detail.requestId,
                    requestContext,
                    responseContext,
                    Instant.now(),
                    PsiDocumentManager.getInstance(project).getPsiFile(document)?.virtualFile,
                    rangeMarker,
                    remainingRecommendation,
                    selectedIndex
                )

                ApplicationManager.getApplication().messageBus.syncPublisher(
                    CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED,
                ).afterAccept(states, sessionContext, rangeMarker)
            }
        }
    }

    private fun isMatchingSymbol(symbol: Char): Boolean =
        PAIRED_BRACKETS.containsKey(symbol) || PAIRED_BRACKETS.containsValue(symbol) || PAIRED_QUOTES.contains(symbol) ||
            symbol.isWhitespace()

    fun getUserInputSinceInvocation(editor: Editor, invocationOffset: Int): String {
        val currentOffset = editor.caretModel.primaryCaret.offset
        return editor.document.getText(TextRange(invocationOffset, currentOffset))
    }

    fun getCaretMovement(editor: Editor, caretPosition: CaretPosition): CaretMovement {
        val oldOffset = caretPosition.offset
        val newOffset = editor.caretModel.primaryCaret.offset
        return when {
            oldOffset < newOffset -> CaretMovement.MOVE_FORWARD
            oldOffset > newOffset -> CaretMovement.MOVE_BACKWARD
            else -> CaretMovement.NO_CHANGE
        }
    }

    fun getMatchingSymbolsFromRecommendation(
        editor: Editor,
        recommendation: String,
        isTruncatedOnRight: Boolean,
        sessionContext: SessionContext
    ): List<Pair<Int, Int>> {
        val result = mutableListOf<Pair<Int, Int>>()
        val bracketsStack = Stack<Char>()
        val quotesStack = Stack<Pair<Char, Pair<Int, Int>>>()
        val caretOffset = editor.caretModel.primaryCaret.offset
        val document = editor.document
        val lineEndOffset = document.getLineEndOffset(document.getLineNumber(caretOffset))
        val lineText = document.charsSequence.subSequence(caretOffset, lineEndOffset)

        var totalDocLengthChecked = 0
        var current = 0

        result.add(0 to caretOffset)
        result.add(recommendation.length + 1 to lineEndOffset)

        if (isTruncatedOnRight) return result

        while (current < recommendation.length &&
            totalDocLengthChecked < lineText.length &&
            totalDocLengthChecked < recommendation.length
        ) {
            val currentDocChar = lineText[totalDocLengthChecked]
            if (!isMatchingSymbol(currentDocChar)) {
                // currentDocChar is not a matching symbol, so we try to compare the remaining strings as a last step to match
                val recommendationRemaining = recommendation.substring(current)
                val rightContextRemaining = lineText.subSequence(totalDocLengthChecked, lineText.length).toString()
                if (recommendationRemaining == rightContextRemaining) {
                    for (i in 1..recommendation.length - current) {
                        result.add(current + i to caretOffset + totalDocLengthChecked + i)
                    }
                    result.sortBy { it.first }
                }
                break
            }
            totalDocLengthChecked++

            // find symbol in the recommendation that will match this
            while (current < recommendation.length) {
                val char = recommendation[current]
                current++

                // if char isn't a paired symbol, or it is, but it's not the matching currentDocChar or
                // the opening version of it, then we're done
                if (!isMatchingSymbol(char) || (char != currentDocChar && PAIRED_BRACKETS[char] != currentDocChar)) {
                    continue
                }

                // if char is an opening bracket, push it to the stack
                if (PAIRED_BRACKETS[char] == currentDocChar) {
                    bracketsStack.push(char)
                    continue
                }

                // char is currentDocChar, it's one of a bracket, a quote, or a whitespace character.
                // If it's a whitespace character, directly add it to the result,
                // if it's a bracket or a quote, check if this char is already having a matching opening symbol
                // on the stack
                if (char.isWhitespace()) {
                    result.add(current to caretOffset + totalDocLengthChecked)
                    break
                } else if (bracketsStack.isNotEmpty() && PAIRED_BRACKETS[bracketsStack.peek()] == char) {
                    bracketsStack.pop()
                } else if (quotesStack.isNotEmpty() && quotesStack.peek().first == char) {
                    result.add(quotesStack.pop().second)
                    result.add(current to caretOffset + totalDocLengthChecked)
                    break
                } else {
                    // char does not have a matching opening symbol in the stack, if it's a (opening) bracket,
                    // immediately add it to the result; if it's a quote, push it to the stack
                    if (PAIRED_QUOTES.contains(char)) {
                        quotesStack.push(char to (current to caretOffset + totalDocLengthChecked))
                    } else {
                        result.add(current to caretOffset + totalDocLengthChecked)
                    }
                    break
                }
            }
        }

        // if there are any symbols left in the stack, add them to the result
        quotesStack.forEach { result.add(it.second) }
        result.sortBy { it.first }

        sessionContext.insertEndOffset = result[result.size - 2].second

        return result
    }

    // example:         recommendation:         document
    //                  line1
    //                  line2
    //                  line3                   line3
    //                                          line4
    //                                          ...
    // number of lines overlapping would be one, and it will be line 3
    fun findOverLappingLines(
        editor: Editor,
        recommendationLines: List<String>,
        isTruncatedOnRight: Boolean,
        sessionContext: SessionContext
    ): Int {
        val caretOffset = editor.caretModel.offset
        if (isTruncatedOnRight) {
            // insertEndOffset value only makes sense when there are matching closing brackets, if there's right context
            // resolution applied, set this value to the current caret offset
            sessionContext.insertEndOffset = caretOffset
            return 0
        }

        val text = editor.document.charsSequence
        val document = editor.document
        val textLines = mutableListOf<Pair<String, Int>>()
        val caretLine = document.getLineNumber(caretOffset)
        var currentLineNum = caretLine + 1
        val recommendationLinesNotBlank = recommendationLines.filter { it.isNotBlank() }
        while (currentLineNum < document.lineCount && textLines.size < recommendationLinesNotBlank.size) {
            val currentLine = text.subSequence(
                document.getLineStartOffset(currentLineNum),
                document.getLineEndOffset(currentLineNum)
            )
            if (currentLine.isNotBlank()) {
                textLines.add(currentLine.toString() to document.getLineEndOffset(currentLineNum))
            }
            currentLineNum++
        }

        val numOfNonEmptyLinesMatching = countNonEmptyLinesMatching(recommendationLinesNotBlank, textLines)
        val numOfLinesMatching = countLinesMatching(recommendationLines, numOfNonEmptyLinesMatching)
        if (numOfNonEmptyLinesMatching > 0) {
            sessionContext.insertEndOffset = textLines[numOfNonEmptyLinesMatching - 1].second
        } else if (recommendationLines.isNotEmpty()) {
            sessionContext.insertEndOffset = document.getLineEndOffset(caretLine)
        }

        return numOfLinesMatching
    }

    private fun countLinesMatching(lines: List<String>, targetNonEmptyLines: Int): Int {
        var count = 0
        var nonEmptyCount = 0

        for (line in lines.asReversed()) {
            if (nonEmptyCount == targetNonEmptyLines) {
                break
            }
            if (line.isNotBlank()) {
                nonEmptyCount++
            }
            count++
        }
        return count
    }

    private fun countNonEmptyLinesMatching(recommendationLines: List<String>, textLines: List<Pair<String, Int>>): Int {
        // i lines we want to match
        for (i in textLines.size downTo 1) {
            val recommendationStart = recommendationLines.size - i
            var matching = true
            for (j in 0 until i) {
                if (recommendationLines[recommendationStart + j].trimEnd() != textLines[j].first.trimEnd()) {
                    matching = false
                    break
                }
            }
            if (matching) {
                return i
            }
        }
        return 0
    }

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


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

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

import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.roots.ProjectFileIndex
import com.intellij.openapi.ui.popup.JBPopup
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.vfs.VfsUtilCore
import com.intellij.psi.PsiFile
import com.intellij.ui.popup.AbstractPopup
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.LEFT_CONTEXT_ON_CURRENT_LINE
import java.awt.Point
import kotlin.math.max
import kotlin.math.min

object CodeWhispererEditorUtil {
    fun getFileContextInfo(editor: Editor, psiFile: PsiFile): FileContextInfo {
        val caretContext = extractCaretContext(editor)
        val fileName = getFileName(psiFile)
        val programmingLanguage = psiFile.programmingLanguage()
        return FileContextInfo(caretContext, fileName, programmingLanguage)
    }

    fun extractCaretContext(editor: Editor): CaretContext {
        val document = editor.document
        val caretOffset = editor.caretModel.primaryCaret.offset
        val totalCharLength = editor.document.textLength

        val caretLeftFileContext = document.getText(
            TextRange(
                CodeWhispererConstants.BEGINNING_OF_FILE.coerceAtLeast(caretOffset - CodeWhispererConstants.CHARACTERS_LIMIT),
                caretOffset
            )
        )

        val caretRightFileContext = document.getText(
            TextRange(
                caretOffset,
                totalCharLength.coerceAtMost(CodeWhispererConstants.CHARACTERS_LIMIT + caretOffset)
            )
        )

        val lineNumber = document.getLineNumber(caretOffset)
        val startOffset = document.getLineStartOffset(lineNumber)
        var leftContextOnCurrentLine = document.getText(TextRange(startOffset, caretOffset))
        leftContextOnCurrentLine = leftContextOnCurrentLine.substring(
            leftContextOnCurrentLine.length - leftContextOnCurrentLine.length.coerceAtMost(LEFT_CONTEXT_ON_CURRENT_LINE)
        )

        return CaretContext(caretLeftFileContext, caretRightFileContext, leftContextOnCurrentLine)
    }

    fun getCaretPosition(editor: Editor): CaretPosition {
        val offset = editor.caretModel.primaryCaret.offset
        val line = editor.caretModel.primaryCaret.visualPosition.line
        return CaretPosition(offset, line)
    }

    private fun getFileName(psiFile: PsiFile): String =
        psiFile.name.substring(0, psiFile.name.length.coerceAtMost(CodeWhispererConstants.FILENAME_CHARS_LIMIT))

    fun getRelativePathToContentRoot(editor: Editor): String? =
        editor.project?.let { project ->
            FileDocumentManager.getInstance().getFile(editor.document)?.let { vFile ->
                val fileIndex = ProjectFileIndex.getInstance(project)
                val contentRoot = runReadAction { fileIndex.getContentRootForFile(vFile) }
                contentRoot?.let {
                    VfsUtilCore.getRelativePath(vFile, it)
                }
            }
        }

    fun getPopupPositionAboveText(editor: Editor, popup: JBPopup, offset: Int): Point {
        val textAbsolutePosition = editor.offsetToXY(offset)
        val editorLocation = editor.component.locationOnScreen
        val editorContentLocation = editor.contentComponent.locationOnScreen
        return Point(
            editorContentLocation.x + textAbsolutePosition.x,
            editorLocation.y + textAbsolutePosition.y - editor.scrollingModel.verticalScrollOffset -
                (popup as AbstractPopup).preferredContentSize.height
        )
    }

    /**
     * Checks if the [otherRange] overlaps this TextRange. Note that the comparison is `<` because the endOffset of TextRange is exclusive.
     */
    fun TextRange.overlaps(otherRange: TextRange): Boolean =
        if (otherRange.isEmpty) {
            // Handle case when otherRange is empty and within the range
            otherRange.startOffset in startOffset until endOffset
        } else {
            max(startOffset, otherRange.startOffset) < min(endOffset, otherRange.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.editor

import com.intellij.codeInsight.editorActions.EnterHandler
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.editor.Caret
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.actionSystem.EditorActionHandler
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutoTriggerService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType

class CodeWhispererEnterHandler(private val originalHandler: EditorActionHandler) : EnterHandler(originalHandler) {
    override fun executeWriteAction(editor: Editor, caret: Caret?, dataContext: DataContext?) {
        originalHandler.execute(editor, caret, dataContext)

        ApplicationManager.getApplication().executeOnPooledThread {
            CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.Enter())
        }
    }
}


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

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

import com.intellij.codeInsight.editorActions.TypedHandlerDelegate
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiFile
import kotlinx.coroutines.Job
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererLanguageManager
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutoTriggerService
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants

class CodeWhispererTypedHandler : TypedHandlerDelegate() {
    private var triggerOnIdle: Job? = null
    override fun charTyped(c: Char, project: Project, editor: Editor, psiFiles: PsiFile): Result {
        triggerOnIdle?.cancel()

        // Special Char
        if (CodeWhispererConstants.SPECIAL_CHARACTERS_LIST.contains(c.toString())) {
            CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.SpecialChar(c))
            return Result.CONTINUE
        }

        val language = CodeWhispererLanguageManager.getInstance().getLanguage(psiFiles)

        if ((
                CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.Classifier &&
                    language.isClassifierSupported()
                ) || language.isAllClassifier()
        ) {
            CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.Classifier())
        } else {
            triggerOnIdle = CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.IdleTime())
        }

        return Result.CONTINUE
    }
}


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

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

import com.intellij.openapi.components.BaseState
import com.intellij.openapi.components.PersistentStateComponent
import com.intellij.openapi.components.State
import com.intellij.openapi.components.Storage
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.util.xmlb.annotations.Property
import org.jetbrains.annotations.ApiStatus.ScheduledForRemoval
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
import software.aws.toolkits.jetbrains.core.explorer.refreshDevToolTree
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getConnectionStartUrl
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.isAccessTokenExpired
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.isRefreshTokenExpired
import software.aws.toolkits.telemetry.AwsTelemetry
import java.time.LocalDateTime

// TODO: refactor this class, now it's managing action and state
@State(name = "codewhispererStates", storages = [Storage("aws.xml")])
class CodeWhispererExplorerActionManager : PersistentStateComponent<CodeWhispererExploreActionState> {
    private val actionState = CodeWhispererExploreActionState()
    private val suspendedConnections = mutableSetOf<String>()

    fun isSuspended(project: Project): Boolean {
        val startUrl = getCodeWhispererConnectionStartUrl(project)
        return suspendedConnections.contains(startUrl)
    }

    fun setSuspended(project: Project) {
        val startUrl = getCodeWhispererConnectionStartUrl(project)
        if (!suspendedConnections.add(startUrl)) {
            return
        }
        project.refreshDevToolTree()
    }

    private fun getCodeWhispererConnectionStartUrl(project: Project): String {
        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
        return getConnectionStartUrl(connection) ?: CodeWhispererConstants.ACCOUNTLESS_START_URL
    }

    fun isAutoEnabled(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.IsAutoEnabled, true)

    fun setAutoEnabled(isAutoEnabled: Boolean) {
        actionState.value[CodeWhispererExploreStateType.IsAutoEnabled] = isAutoEnabled
    }

    fun hasShownHowToUseCodeWhisperer(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.HasShownHowToUseCodeWhisperer, false)

    fun setHasShownHowToUseCodeWhisperer(hasShownHowToUseCodeWhisperer: Boolean) {
        actionState.value[CodeWhispererExploreStateType.HasShownHowToUseCodeWhisperer] = hasShownHowToUseCodeWhisperer
    }

    fun setAccountlessNotificationWarnTimestamp() {
        actionState.accountlessWarnTimestamp = LocalDateTime.now().format(CodeWhispererConstants.TIMESTAMP_FORMATTER)
    }

    fun setAccountlessNotificationErrorTimestamp() {
        actionState.accountlessErrorTimestamp = LocalDateTime.now().format(CodeWhispererConstants.TIMESTAMP_FORMATTER)
    }

    fun getAccountlessWarnNotificationTimestamp(): String? = actionState.accountlessWarnTimestamp

    fun getAccountlessErrorNotificationTimestamp(): String? = actionState.accountlessErrorTimestamp

    fun getDoNotShowAgainWarn(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.DoNotShowAgainWarn, false)

    fun setDoNotShowAgainWarn(doNotShowAgain: Boolean) {
        actionState.value[CodeWhispererExploreStateType.DoNotShowAgainWarn] = doNotShowAgain
    }

    fun getDoNotShowAgainError(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.DoNotShowAgainError, false)

    fun setDoNotShowAgainError(doNotShowAgain: Boolean) {
        actionState.value[CodeWhispererExploreStateType.DoNotShowAgainError] = doNotShowAgain
    }

    fun getConnectionExpiredDoNotShowAgain(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.ConnectionExpiredDoNotShowAgain, false)

    fun setConnectionExpiredDoNotShowAgain(doNotShowAgain: Boolean) {
        actionState.value[CodeWhispererExploreStateType.ConnectionExpiredDoNotShowAgain] = doNotShowAgain
    }

    fun getAccountlessNullified(): Boolean = actionState.value.getOrDefault(CodeWhispererExploreStateType.AccountlessNullified, false)

    fun setAccountlessNullified(accountlessNullified: Boolean) {
        actionState.value[CodeWhispererExploreStateType.AccountlessNullified] = accountlessNullified
    }

    fun setAutoSuggestion(project: Project, isAutoEnabled: Boolean) {
        setAutoEnabled(isAutoEnabled)
        val autoSuggestionState = if (isAutoEnabled) CodeWhispererConstants.AutoSuggestion.ACTIVATED else CodeWhispererConstants.AutoSuggestion.DEACTIVATED
        AwsTelemetry.modifySetting(project, settingId = CodeWhispererConstants.AutoSuggestion.SETTING_ID, settingState = autoSuggestionState)
        project.refreshDevToolTree()
    }

    @Deprecated("Accountless credential will be removed soon")
    @ScheduledForRemoval
    // Will keep it for existing accountless users
    /**
     * Will be called from CodeWhispererService.showRecommendationInPopup()
     * Caller (e.x. CodeWhispererService) should take care if null value returned, popup a notification/hint window or dialog etc.
     */
    fun resolveAccessToken(): String? {
        if (actionState.token == null) {
            LOG.warn { "Logical Error: Try to get access token before token initialization" }
        }
        return actionState.token
    }

    fun checkActiveCodeWhispererConnectionType(project: Project) = when {
        actionState.token != null -> CodeWhispererLoginType.Accountless
        isAccessTokenExpired(project) || isRefreshTokenExpired(project) -> CodeWhispererLoginType.Expired
        else -> {
            val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
            if (conn != null) {
                if (conn.isSono()) {
                    CodeWhispererLoginType.Sono
                } else {
                    CodeWhispererLoginType.SSO
                }
            } else {
                CodeWhispererLoginType.Logout
            }
        }
    }

    fun nullifyAccountlessCredentialIfNeeded() {
        if (actionState.token != null) {
            setAccountlessNullified(true)
            actionState.token = null
        }
    }

    override fun getState(): CodeWhispererExploreActionState = CodeWhispererExploreActionState().apply {
        value.putAll(actionState.value)
        token = actionState.token
        accountlessWarnTimestamp = actionState.accountlessWarnTimestamp
        accountlessErrorTimestamp = actionState.accountlessErrorTimestamp
    }

    override fun loadState(state: CodeWhispererExploreActionState) {
        actionState.value.clear()
        actionState.token = state.token
        actionState.value.putAll(state.value)
        actionState.accountlessWarnTimestamp = state.accountlessWarnTimestamp
        actionState.accountlessErrorTimestamp = state.accountlessErrorTimestamp
    }

    companion object {
        @JvmStatic
        fun getInstance(): CodeWhispererExplorerActionManager = service()

        private val LOG = getLogger<CodeWhispererExplorerActionManager>()
    }
}

class CodeWhispererExploreActionState : BaseState() {
    @get:Property
    val value by map<CodeWhispererExploreStateType, Boolean>()

    // can not remove this as we want to support existing accountless users
    @get:Property
    var token by string()

    @get:Property
    var accountlessWarnTimestamp by string()

    @get:Property
    var accountlessErrorTimestamp by string()
}

// TODO: Don't remove IsManualEnabled
enum class CodeWhispererExploreStateType {
    IsAutoEnabled,
    IsManualEnabled,
    HasAcceptedTermsOfServices,
    HasShownHowToUseCodeWhisperer,
    DoNotShowAgainWarn,
    DoNotShowAgainError,
    AccountlessNullified,
    ConnectionExpiredDoNotShowAgain
}

interface CodeWhispererActivationChangedListener {
    fun activationChanged(value: Boolean) {}
}

fun isCodeWhispererEnabled(project: Project) = with(CodeWhispererExplorerActionManager.getInstance()) {
    checkActiveCodeWhispererConnectionType(project) != CodeWhispererLoginType.Logout
}

fun isCodeWhispererExpired(project: Project) = with(CodeWhispererExplorerActionManager.getInstance()) {
    checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Expired
}


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

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

import com.intellij.ide.util.treeView.AbstractTreeNode
import software.aws.toolkits.jetbrains.core.explorer.devToolsTab.DevToolsTreeStructureProvider
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.CodeWhispererActionNode

class CodeWhispererExplorerTreeStructureProvider : DevToolsTreeStructureProvider() {
    override fun modify(parent: AbstractTreeNode<*>, children: MutableCollection<AbstractTreeNode<*>>): MutableCollection<AbstractTreeNode<*>> =
        when (parent) {
            is CodeWhispererServiceNode ->
                children
                    .sortedWith { x, y ->
                        val order1 = (x as? CodeWhispererActionNode)?.order ?: Int.MAX_VALUE
                        val order2 = (y as? CodeWhispererActionNode)?.order ?: Int.MAX_VALUE
                        order1.compareTo(order2)
                    }.toMutableList()
            else -> children
        }
}


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

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

import com.intellij.ide.projectView.PresentationData
import com.intellij.ide.util.treeView.AbstractTreeNode
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.project.Project
import com.intellij.ui.SimpleTextAttributes
import com.intellij.util.text.DateTimeFormatManager
import software.aws.toolkits.core.utils.tryOrNull
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
import software.aws.toolkits.jetbrains.core.explorer.devToolsTab.nodes.AbstractActionTreeNode
import software.aws.toolkits.jetbrains.core.explorer.devToolsTab.nodes.ActionGroupOnRightClick
import software.aws.toolkits.jetbrains.core.explorer.devToolsTab.nodes.PinnedConnectionNode
import software.aws.toolkits.jetbrains.core.explorer.refreshDevToolTree
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.CodeWhispererReconnectNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.FreeTierUsageLimitHitNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.GetStartedNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.OpenCodeReferenceNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.PauseCodeWhispererNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.ResumeCodeWhispererNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.RunCodeScanNode
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.nodes.WhatIsCodeWhispererNode
import software.aws.toolkits.jetbrains.utils.isRunningOnRemoteBackend
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent
import java.time.LocalDate
import java.time.format.DateTimeFormatter
import java.time.temporal.TemporalAdjusters

class CodeWhispererServiceNode(
    project: Project,
    value: String,
) : AbstractActionTreeNode(project, value, null), ActionGroupOnRightClick, PinnedConnectionNode {
    private val nodeProject
        get() = myProject
    private val pauseCodeWhispererNode by lazy { PauseCodeWhispererNode(nodeProject) }
    private val resumeCodeWhispererNode by lazy { ResumeCodeWhispererNode(nodeProject) }
    private val whatIsCodeWhispererNode by lazy { WhatIsCodeWhispererNode(nodeProject) }
    private val getStartedCodeWhispererNode by lazy { GetStartedNode(nodeProject) }
    private val openCodeReferenceNode by lazy { OpenCodeReferenceNode(nodeProject) }
    private val runCodeScanNode by lazy { RunCodeScanNode(nodeProject) }
    private val codeWhispererReconnectNode by lazy { CodeWhispererReconnectNode(nodeProject) }
    private val freeTierUsageLimitHitNode by lazy {
        // we should probably build the text dynamically in case the format setting changes,
        // but that shouldn't happen often enough for us to care
        val formatter = tryOrNull {
            DateTimeFormatter.ofPattern(DateTimeFormatManager.getInstance().dateFormatPattern)
        } ?: DateTimeFormatter.ofPattern(DateTimeFormatManager.DEFAULT_DATE_FORMAT)
        val date = LocalDate.now().with(TemporalAdjusters.firstDayOfNextMonth())

        FreeTierUsageLimitHitNode(nodeProject, formatter.format(date))
    }

    init {
        ApplicationManager.getApplication().messageBus.connect().subscribe(
            ToolkitConnectionManagerListener.TOPIC,
            object : ToolkitConnectionManagerListener {
                override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                    // TODO: Move this IF block into nullifyAccountlessCredentialIfNeeded()
                    if (newConnection is AwsBearerTokenConnection) {
                        CodeWhispererExplorerActionManager.getInstance().nullifyAccountlessCredentialIfNeeded()
                    }
                    project.refreshDevToolTree()
                }
            }
        )
    }

    override fun onDoubleClick(event: MouseEvent) {}

    override fun getChildren(): Collection<AbstractTreeNode<*>> {
        if (isRunningOnRemoteBackend()) {
            return emptyList()
        }

        val manager = CodeWhispererExplorerActionManager.getInstance()
        val activeConnectionType = manager.checkActiveCodeWhispererConnectionType(project)

        return when (activeConnectionType) {
            CodeWhispererLoginType.Logout -> listOf(whatIsCodeWhispererNode, getStartedCodeWhispererNode)
            CodeWhispererLoginType.Expired -> listOf(codeWhispererReconnectNode, whatIsCodeWhispererNode)

            else -> {
                if (manager.isSuspended(nodeProject)) {
                    listOf(freeTierUsageLimitHitNode, runCodeScanNode, openCodeReferenceNode)
                } else if (manager.isAutoEnabled()) {
                    listOf(pauseCodeWhispererNode, runCodeScanNode, openCodeReferenceNode)
                } else {
                    listOf(resumeCodeWhispererNode, runCodeScanNode, openCodeReferenceNode)
                }
            }
        }
    }

    override fun update(presentation: PresentationData) {
        super.update(presentation)
        if (isRunningOnRemoteBackend()) {
            presentation.addText(message("codewhisperer.explorer.root_node.unavailable"), SimpleTextAttributes.GRAY_ATTRIBUTES)
            return
        }

        val connectionType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project)
        when (connectionType) {
            CodeWhispererLoginType.Expired -> {
                presentation.addText(message("codewhisperer.explorer.root_node.login_type.expired"), SimpleTextAttributes.GRAY_ATTRIBUTES)
            }

            CodeWhispererLoginType.Accountless -> {
                presentation.addText(message("codewhisperer.explorer.root_node.login_type.accountless"), SimpleTextAttributes.GRAY_ATTRIBUTES)
            }

            CodeWhispererLoginType.SSO -> {
                presentation.addText(message("codewhisperer.explorer.root_node.login_type.sso"), SimpleTextAttributes.GRAY_ATTRIBUTES)
            }

            CodeWhispererLoginType.Sono -> {
                presentation.addText(message("codewhisperer.explorer.root_node.login_type.aws_builder_id"), SimpleTextAttributes.GRAY_ATTRIBUTES)
            }
            else -> {}
        }
    }

    override fun actionGroupName(): String = "aws.toolkit.explorer.codewhisperer"

    override fun feature() = CodeWhispererConnection.getInstance()
}


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

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

import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.core.explorer.devToolsTab.nodes.AbstractActionTreeNode
import software.aws.toolkits.jetbrains.core.explorer.nodes.AwsExplorerNode
import javax.swing.Icon

abstract class CodeWhispererActionNode(
    project: Project,
    actionName: String,
    val order: Int,
    icon: Icon
) : AbstractActionTreeNode(
    project,
    actionName,
    icon
) {
    override fun getChildren(): List<AwsExplorerNode<*>> = emptyList()
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.reconnectCodeWhisperer
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent

class CodeWhispererReconnectNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.reconnect"),
    0,
    AllIcons.Actions.Execute
) {
    override fun onDoubleClick(event: MouseEvent) {
        reconnectCodeWhisperer(project)
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.project.Project
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent

class FreeTierUsageLimitHitNode(nodeProject: Project, val date: String) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.usage_limit_hit", date),
    -1,
    AllIcons.Actions.Suspend
) {
    override fun onDoubleClick(event: MouseEvent) {}
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.ide.plugins.PluginManagerCore
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.extensions.PluginId
import com.intellij.openapi.fileEditor.FileEditorManager
import com.intellij.openapi.fileEditor.OpenFileDescriptor
import com.intellij.openapi.fileEditor.TextEditorWithPreview
import com.intellij.openapi.project.Project
import com.intellij.openapi.startup.StartupActivity
import com.intellij.openapi.vfs.VfsUtil
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.AwsToolkit
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.credentials.CodeWhispererLoginDialog
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.startup.CodeWhispererProjectStartupActivity
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.UiTelemetry
import java.awt.event.MouseEvent

class GetStartedNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.enable"),
    0,
    AllIcons.Actions.Execute
) {
    override fun onDoubleClick(event: MouseEvent) {
        enableCodeWhisperer(project)
        UiTelemetry.click(project, "cw_signUp_Cta")
    }

    /**
     * 2 cases
     * (1) User who don't have SSO based connection click on CodeWhisperer Start node
     * (2) User who already have SSO based connection from previous operation via i.g. Toolkit Add Connection click on CodeWhisperer Start node
     */
    private fun enableCodeWhisperer(project: Project) {
        val explorerActionManager = CodeWhispererExplorerActionManager.getInstance()
        val connectionManager = ToolkitConnectionManager.getInstance(project)
        connectionManager.activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let {
            project.refreshDevToolTree()
        } ?: run {
            runInEdt {
                // Start from scratch if no active connection
                if (CodeWhispererLoginDialog(project).showAndGet()) {
                    project.refreshDevToolTree()
                }
            }
        }

        if (isCodeWhispererEnabled(project)) {
            StartupActivity.POST_STARTUP_ACTIVITY.findExtension(CodeWhispererProjectStartupActivity::class.java)?.let {
                it.runActivity(project)
            }
            if (!explorerActionManager.hasShownHowToUseCodeWhisperer()) {
                showHowToUseCodeWhispererPage(project)
            }
        }
    }

    private fun showHowToUseCodeWhispererPage(project: Project) {
        val plugin = PluginManagerCore.getPlugin(PluginId.getId(AwsToolkit.PLUGIN_ID)) ?: return
        val path = plugin.pluginPath.resolve("assets").resolve("WelcomeToCodeWhisperer.md") ?: return
        VfsUtil.findFile(path, true)?.let { readme ->
            readme.putUserData(TextEditorWithPreview.DEFAULT_LAYOUT_FOR_FILE, TextEditorWithPreview.Layout.SHOW_PREVIEW)

            val fileEditorManager = FileEditorManager.getInstance(project)
            ApplicationManager.getApplication().invokeLater {
                val editor = fileEditorManager.openTextEditor(OpenFileDescriptor(project, readme), true)
                if (editor == null) {
                    LOG.warn { "Failed to open WelcomeToCodeWhisperer.md" }
                } else {
                    CodeWhispererExplorerActionManager.getInstance().setHasShownHowToUseCodeWhisperer(true)
                }
            }
        }
    }

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


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent

class OpenCodeReferenceNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.code_reference.open"),
    3,
    AllIcons.Actions.Preview
) {
    override fun onDoubleClick(event: MouseEvent) {
        CodeWhispererCodeReferenceManager.getInstance(project).showCodeReferencePanel()
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent

class PauseCodeWhispererNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.pause_auto"),
    1,
    AllIcons.Actions.Pause
) {
    override fun onDoubleClick(event: MouseEvent) {
        CodeWhispererExplorerActionManager.getInstance().setAutoSuggestion(project, false)
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
import software.aws.toolkits.resources.message
import java.awt.event.MouseEvent

class ResumeCodeWhispererNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.resume_auto"),
    1,
    AllIcons.Actions.Resume
) {
    override fun onDoubleClick(event: MouseEvent) {
        CodeWhispererExplorerActionManager.getInstance().setAutoSuggestion(project, 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.explorer.nodes

import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
import java.awt.event.MouseEvent

class RunCodeScanNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    CodeWhispererCodeScanManager.getInstance(nodeProject).getActionButtonText(),
    2,
    CodeWhispererCodeScanManager.getInstance(nodeProject).getActionButtonIconForExplorerNode()
) {

    private val codeScanManager = CodeWhispererCodeScanManager.getInstance(project)

    override fun onDoubleClick(event: MouseEvent) {
        if (codeScanManager.isCodeScanInProgress()) {
            codeScanManager.stopCodeScan()
        } else {
            codeScanManager.runCodeScan()
        }
    }
}


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

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

import com.intellij.icons.AllIcons
import com.intellij.ide.BrowserUtil
import com.intellij.openapi.project.Project
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
import software.aws.toolkits.resources.message
import software.aws.toolkits.telemetry.UiTelemetry
import java.awt.event.MouseEvent
import java.net.URI

class WhatIsCodeWhispererNode(nodeProject: Project) : CodeWhispererActionNode(
    nodeProject,
    message("codewhisperer.explorer.what_is"),
    1,
    AllIcons.Actions.Help
) {
    override fun onDoubleClick(event: MouseEvent) {
        BrowserUtil.browse(URI(CodeWhispererConstants.CODEWHISPERER_LEARN_MORE_URI))
        UiTelemetry.click(project, "cw_learnMore_Cta")
    }
}


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

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

import com.intellij.openapi.editor.Editor
import com.intellij.openapi.fileTypes.PlainTextLanguage
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiFileFactory
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererUnknownLanguage

/**
 * This class provides a fallback implementation of the CodeWhispererImportAdder abstract class for handling
 * import statements for Java, Python, and JavaScript files in cases when the respective language plugins are not
 * installed. It is meant to be generic and create bare import elements and insert them to the top of the file.
 */
class CodeWhispererFallbackImportAdder : CodeWhispererImportAdder() {
    override val supportedLanguages: List<CodeWhispererProgrammingLanguage> = listOf(
        CodeWhispererUnknownLanguage.INSTANCE
    )
    override val dummyFileName: String = "test.txt"

    override fun createNewImportPsiElement(psiFile: PsiFile, statement: String): PsiElement? {
        val statementWithNewLine = if (statement.endsWith("\n")) statement else "$statement\n"
        val project = psiFile.project
        val fileFactory = PsiFileFactory.getInstance(project)
        val dummyFile = fileFactory.createFileFromText(dummyFileName, PlainTextLanguage.INSTANCE, statementWithNewLine)
            ?: return null
        return dummyFile.firstChild
    }

    /**
     * Always returns null, as duplicates are not checked in the fallback implementation.
     */
    override fun hasDuplicatedImportsHelper(newImport: PsiElement, existingImports: List<PsiElement>): PsiElement? = null

    /**
     * Always returns an empty list, as top-level imports are not checked in the fallback implementation.
     */
    override fun getTopLevelImports(psiFile: PsiFile, editor: Editor): List<PsiElement> = emptyList()

    /**
     * Always returns an empty list, as local imports are not checked in the fallback implementation.
     */
    override fun getLocalImports(psiFile: PsiFile, editor: Editor): List<PsiElement> = emptyList()

    /**
     * Adds the new import statement to the beginning of the file.
     */
    override fun addImport(psiFile: PsiFile, editor: Editor, newImport: PsiElement): Boolean {
        val first = psiFile.firstChild
        if (first == null) {
            psiFile.add(newImport)
        } else {
            psiFile.addBefore(newImport, first)
        }
        return true
    }
}


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

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

import com.intellij.openapi.editor.Editor
import com.intellij.openapi.extensions.ExtensionPointName
import com.intellij.psi.PsiDocumentManager
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import software.amazon.awssdk.services.codewhispererruntime.model.Import
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext

abstract class CodeWhispererImportAdder {
    abstract val supportedLanguages: List<CodeWhispererProgrammingLanguage>
    abstract val dummyFileName: String

    fun insertImportStatements(states: InvocationContext, sessionContext: SessionContext) {
        val imports = states.recommendationContext.details[sessionContext.selectedIndex]
            .recommendation.mostRelevantMissingImports()
        LOG.info { "Adding ${imports.size} imports for completions, sessionId: ${states.responseContext.sessionId}" }
        imports.forEach {
            insertImportStatement(states, it)
        }
    }

    private fun insertImportStatement(states: InvocationContext, import: Import) {
        val project = states.requestContext.project
        val editor = states.requestContext.editor
        val document = editor.document
        val psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document) ?: return

        val statement = import.statement()
        LOG.info { "Import statement to be added: $statement" }
        val newImport = createNewImportPsiElement(psiFile, statement)
        if (newImport == null) {
            LOG.debug { "Failed to create the import element using the import string" }
            return
        }

        if (!isSupportedImportStyle(newImport)) {
            LOG.debug { "Import statement \"${newImport.text}\" is not supported" }
            return
        }

        LOG.debug { "Checking duplicates with existing imports" }
        val hasDuplicate = hasDuplicatedImports(psiFile, editor, newImport)
        if (hasDuplicate) {
            LOG.debug { "Found duplicates with existing imports, not adding the new import" }
            return
        } else {
            LOG.debug { "Found no duplicates with existing imports" }
        }

        val added = addImport(psiFile, editor, newImport)
        LOG.info { "Added import: $added" }
    }

    abstract fun createNewImportPsiElement(psiFile: PsiFile, statement: String): PsiElement?

    open fun isSupportedImportStyle(newImport: PsiElement) = true

    // Currently if the new import is 'from a import b, c', a duplicate match to any of the import element
    // will return as a valid duplicate to the whole import statement.
    open fun hasDuplicatedImports(psiFile: PsiFile, editor: Editor, newImport: PsiElement): Boolean {
        val topLevelImports = getTopLevelImports(psiFile, editor)
        LOG.debug {
            "Checking top-level imports: [${topLevelImports.map { it.text }.reduceOrNull { acc, s -> "$acc, $s" }.orEmpty()}]"
        }

        var duplicate = hasDuplicatedImportsHelper(newImport, topLevelImports)
        if (duplicate != null) {
            LOG.debug { "Found duplicates from top-level imports \"${duplicate?.text}\"" }
            return true
        } else {
            LOG.debug { "Found no duplicates from top-level imports" }
        }

        val localImports = getLocalImports(psiFile, editor)
        LOG.debug {
            "Checking local imports: [${localImports.map { it.text }.reduceOrNull { acc, s -> "$acc, $s" }.orEmpty()}]"
        }
        duplicate = hasDuplicatedImportsHelper(newImport, localImports)
        if (duplicate != null) {
            LOG.debug { "Found duplicates from local imports \"${duplicate.text}\"" }
            return true
        } else {
            LOG.debug { "Found no duplicates from local imports" }
        }

        return false
    }

    abstract fun hasDuplicatedImportsHelper(newImport: PsiElement, existingImports: List<PsiElement>): PsiElement?

    abstract fun getTopLevelImports(psiFile: PsiFile, editor: Editor): List<PsiElement>

    abstract fun getLocalImports(psiFile: PsiFile, editor: Editor): List<PsiElement>

    abstract fun addImport(psiFile: PsiFile, editor: Editor, newImport: PsiElement): Boolean

    companion object {
        private val EP = ExtensionPointName.create<CodeWhispererImportAdder>("aws.toolkit.codewhisperer.importAdder")
        internal val LOG = getLogger<CodeWhispererImportAdder>()

        fun get(language: CodeWhispererProgrammingLanguage): CodeWhispererImportAdder? =
            EP.extensionList.firstOrNull { language in it.supportedLanguages }
                ?: EP.extensionList.find { it is CodeWhispererFallbackImportAdder }
    }
}


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

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

import com.intellij.openapi.editor.RangeMarker
import software.aws.toolkits.core.utils.debug
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings

object CodeWhispererImportAdderListener : CodeWhispererUserActionListener {
    internal val LOG = getLogger<CodeWhispererImportAdderListener>()
    override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {
        if (!CodeWhispererSettings.getInstance().isImportAdderEnabled()) {
            LOG.debug { "Import adder not enabled in user settings" }
            return
        }
        val language = states.requestContext.fileContextInfo.programmingLanguage
        if (!language.isImportAdderSupported()) {
            LOG.debug { "Import adder is not supported for $language" }
            return
        }
        val importAdder = CodeWhispererImportAdder.get(language)
        if (importAdder == null) {
            LOG.debug { "No import adder found for $language" }
            return
        }
        importAdder.insertImportStatements(states, sessionContext)
    }
}


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

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

import com.intellij.lang.java.JavaLanguage
import com.intellij.openapi.editor.Editor
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiFileFactory
import com.intellij.psi.PsiImportStatement
import com.intellij.psi.PsiImportStatementBase
import com.intellij.psi.PsiJavaFile
import com.intellij.util.IncorrectOperationException
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava

class CodeWhispererJavaImportAdder : CodeWhispererImportAdder() {
    override val supportedLanguages: List<CodeWhispererProgrammingLanguage> = listOf(CodeWhispererJava.INSTANCE)
    override val dummyFileName: String = "dummy.java"

    override fun createNewImportPsiElement(psiFile: PsiFile, statement: String): PsiElement? {
        val project = psiFile.project
        val fileFactory = PsiFileFactory.getInstance(project)
        val dummyFile = fileFactory.createFileFromText(dummyFileName, JavaLanguage.INSTANCE, statement)
            as? PsiJavaFile ?: return null
        return dummyFile.importList?.allImportStatements?.getOrNull(0)
    }

    override fun hasDuplicatedImportsHelper(newImport: PsiElement, existingImports: List<PsiElement>): PsiElement? {
        if (newImport !is PsiImportStatement) return newImport
        existingImports.forEach {
            if (it !is PsiImportStatementBase) return@forEach
            if (it.text == newImport.text) return it
        }
        return null
    }

    override fun getTopLevelImports(psiFile: PsiFile, editor: Editor): List<PsiElement> {
        if (psiFile !is PsiJavaFile) return emptyList()
        return psiFile.importList?.allImportStatements?.toList().orEmpty()
    }

    override fun getLocalImports(psiFile: PsiFile, editor: Editor): List<PsiElement> = emptyList()

    override fun addImport(psiFile: PsiFile, editor: Editor, newImport: PsiElement): Boolean {
        if (psiFile !is PsiJavaFile) return false
        if (newImport !is PsiImportStatement) return false

        val addedImport =
            try {
                psiFile.importList?.add(newImport) ?: return false
            } catch (e: IncorrectOperationException) {
                return false
            }
        if (addedImport !is PsiImportStatement) return false
        return true
    }
}


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

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

import com.intellij.openapi.editor.Editor
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiFileFactory
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.QualifiedName
import com.intellij.refactoring.suggested.startOffset
import com.jetbrains.python.PythonLanguage
import com.jetbrains.python.codeInsight.PyCodeInsightSettings
import com.jetbrains.python.codeInsight.imports.AddImportHelper
import com.jetbrains.python.psi.PyFile
import com.jetbrains.python.psi.PyFromImportStatement
import com.jetbrains.python.psi.PyImportStatement
import com.jetbrains.python.psi.PyImportStatementBase
import com.jetbrains.python.psi.PyStatementList
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython

class CodeWhispererPythonImportAdder : CodeWhispererImportAdder() {
    override val supportedLanguages: List<CodeWhispererProgrammingLanguage> = listOf(CodeWhispererPython.INSTANCE)
    override val dummyFileName: String = "dummy.py"

    override fun createNewImportPsiElement(psiFile: PsiFile, statement: String): PsiElement? {
        val fileFactory = PsiFileFactory.getInstance(psiFile.project)
        val dummyFile = fileFactory.createFileFromText(dummyFileName, PythonLanguage.INSTANCE, statement)
            as? PyFile ?: return null
        return dummyFile.importBlock.firstOrNull()
    }

    override fun hasDuplicatedImportsHelper(newImport: PsiElement, existingImports: List<PsiElement>): PsiElement? {
        if (newImport !is PyImportStatementBase) return newImport

        newImport.importElements.map { it.importedQName }.forEachIndexed outer@{ newI, newImportedQName ->
            if (newImportedQName == null) return@outer
            existingImports.forEach { existingImport ->
                if (existingImport !is PyImportStatementBase) return@outer
                if (existingImport::class.java != newImport::class.java) return@outer
                existingImport.importElements.map { it.importedQName }.forEachIndexed inner@{ existingI, existingImportedQName ->
                    if (existingImportedQName == null) return@inner
                    val existingImportAsName = existingImport.importElements[existingI].asName
                    val newImportAsName = newImport.importElements[newI].asName
                    if (existingImportAsName != newImportAsName) return@inner

                    val newFullyQName = QualifiedName.fromDottedString(newImport.fullyQualifiedObjectNames[newI])
                    val existingFullyQName = QualifiedName.fromDottedString(existingImport.fullyQualifiedObjectNames[existingI])
                    if (newImport is PyImportStatement) {
                        if (newImportAsName != null) {
                            if (existingFullyQName == newFullyQName) return existingImport
                        } else {
                            if (existingFullyQName.matchesPrefix(newFullyQName)) return existingImport
                        }
                    } else if (newImport is PyFromImportStatement) {
                        if (newImportedQName == existingImportedQName && newFullyQName == existingFullyQName) return existingImport
                    }
                }
            }
        }
        return null
    }

    override fun getTopLevelImports(psiFile: PsiFile, editor: Editor): List<PsiElement> =
        if (psiFile !is PyFile) emptyList() else psiFile.importBlock

    override fun getLocalImports(psiFile: PsiFile, editor: Editor): List<PsiElement> {
        val localImports = mutableListOf<PsiElement>()
        val offset = editor.caretModel.offset
        val element = psiFile.findElementAt(offset) ?: return emptyList()
        val elementStartOffset = element.startOffset
        var block: PsiElement? = PsiTreeUtil.getParentOfType(element, PyStatementList::class.java)
        while (block != null) {
            localImports.addAll(
                block.children.filter {
                    (it is PyFromImportStatement || it is PyImportStatement) && it.startOffset < elementStartOffset
                }
            )
            block = PsiTreeUtil.getParentOfType(block, PyStatementList::class.java)
        }
        return localImports
    }

    override fun addImport(psiFile: PsiFile, editor: Editor, newImport: PsiElement): Boolean {
        if (newImport !is PyImportStatementBase) return false
        if (newImport is PyFromImportStatement && newImport.isStarImport) {
            val from = newImport.importSourceQName.toString()
            AddImportHelper.addOrUpdateFromImportStatement(psiFile, from, "*", null, null, null)
            return true
        }
        newImport.importElements.forEach {
            val path = it.importedQName.toString()
            val asName = it.asName

            if (!PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT || newImport !is PyFromImportStatement) {
                AddImportHelper.addImportStatement(psiFile, path, asName, null, null)
            } else {
                val from = newImport.importSourceQName.toString()
                AddImportHelper.addOrUpdateFromImportStatement(psiFile, from, path, asName, null, null)
            }
        }
        return 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.inlay

import com.intellij.openapi.editor.Inlay
import com.intellij.openapi.editor.markup.TextAttributes
import java.awt.Graphics
import java.awt.Rectangle
import kotlin.math.max

class CodeWhispererInlayBlockRenderer(myValue: String) : CodeWhispererInlayRenderer(myValue) {
    private val myLines: List<String>
    init {
        myLines = myValue.split("\n")
    }

    override fun calcWidthInPixels(inlay: Inlay<*>): Int {
        val fontMetrics = getFontInfo(inlay.editor).fontMetrics()
        var maxWidthForSingleLine = fontMetrics.stringWidth(myLines[0])
        for (i in myLines.indices) {
            maxWidthForSingleLine = max(maxWidthForSingleLine, fontMetrics.stringWidth(myLines[i]))
        }
        return maxWidthForSingleLine
    }

    override fun calcHeightInPixels(inlay: Inlay<*>): Int = myLines.size * inlay.editor.lineHeight

    override fun paint(inlay: Inlay<*>, g: Graphics, targetRegion: Rectangle, textAttributes: TextAttributes) {
        applyCodeWhispererColorAndFontSettings(inlay.editor, g)
        for (i in myLines.indices) {
            g.drawString(myLines[i], 0, targetRegion.y + i * inlay.editor.lineHeight + inlay.editor.ascent)
        }
    }
}


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

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

import com.intellij.openapi.editor.Inlay
import com.intellij.openapi.editor.markup.TextAttributes
import java.awt.Graphics
import java.awt.Rectangle

class CodeWhispererInlayInlineRenderer(myValue: String) : CodeWhispererInlayRenderer(myValue) {
    override fun calcWidthInPixels(inlay: Inlay<*>): Int {
        val fontInfo = getFontInfo(inlay.editor)
        return if (myValue.isEmpty()) {
            1
        } else {
            fontInfo.fontMetrics().stringWidth(myValue)
        }
    }

    override fun paint(inlay: Inlay<*>, g: Graphics, targetRegion: Rectangle, textAttributes: TextAttributes) {
        applyCodeWhispererColorAndFontSettings(inlay.editor, g)
        g.drawString(myValue, targetRegion.x, targetRegion.y + inlay.editor.ascent)
    }
}