refactor: update pure context for design

Author: phodal

java/src/main/kotlin/cc/unitmesh/idea/service/JavaRelatedContext.kt

@@ -0,0 +1,96 @@
+package cc.unitmesh.idea.service
+
+import com.intellij.openapi.roots.ProjectFileIndex
+import com.intellij.psi.*
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.PsiUtil
+import org.jetbrains.annotations.NotNull
+import org.jetbrains.annotations.Nullable
+
+object JavaRelatedContext {
+    fun findRelatedClassesAndCleanUp(method: PsiMethod): List<PsiClass> {
+        val relatedClasses = findRelatedClasses(method)
+        val cleanedUpClasses = relatedClasses.flatMap { findSuperClasses(it) }.map { cleanUp(it) }
+        return cleanedUpClasses.toList()
+    }
+
+    fun createDescription(psiClass: PsiClass?): String {
+        if (psiClass == null) {
+            return ""
+        }
+
+        val containingFile = psiClass.containingFile
+        val fileHeader = if (containingFile is PsiJavaFile) {
+            val packageStatement = containingFile.packageStatement?.text ?: ""
+            val importList = containingFile.importList?.text ?: ""
+            "$packageStatement\n  \n$importList\n\n"
+        } else {
+            ""
+        }
+
+        val name = containingFile.name
+        val classDescription = "File: $name\n```Java\n$fileHeader${psiClass.text}\n```\n"
+        return classDescription
+    }
+
+    fun prepareContainingClassContext(@NotNull method: PsiMethod): PsiClass? {
+        val file = method.containingFile.copy() as PsiFile
+        val containingClass = method.containingClass ?: return null
+        val psiElement = PsiTreeUtil.findSameElementInCopy(method, file)
+        val createCommentFromText = PsiElementFactory.getInstance(method.project)
+            .createCommentFromText("// TODO: write unit tests", null)
+        psiElement.parent?.addBefore(createCommentFromText, psiElement)
+        return containingClass
+    }
+
+    private fun cleanUp(psiClass: PsiClass): PsiClass {
+        val psiElement = psiClass.copy() as PsiClass
+        psiElement.containingFile.setName(psiClass.containingFile.name)
+        val members = PsiTreeUtil.findChildrenOfType(psiElement, PsiMember::class.java)
+
+        for (element in members) {
+            if (element is PsiMethod) {
+                element.body?.delete()
+                element.docComment?.delete()
+            }
+        }
+
+        val removableMembers = members.filter { canBeRemoved(it) }
+        removableMembers.forEach { it.delete() }
+
+        psiElement.docComment?.delete()
+        return psiElement
+    }
+
+    private fun findSuperClasses(psiClass: PsiClass): List<PsiClass> {
+        return generateSequence(psiClass) { it.superClass }.toList()
+    }
+
+    private fun canBeRemoved(member: PsiMember): Boolean {
+        val modifierList = member.modifierList
+        if (modifierList?.hasModifierProperty("public") != true) {
+            val annotations = member.annotations
+            if (annotations.isEmpty()) {
+                return true
+            }
+        }
+
+        return false
+    }
+
+    private fun findRelatedClasses(method: PsiMethod): List<PsiClass> {
+        val parameters = method.parameterList.parameters
+        val parameterTypes = parameters.map { it.type }
+        val genericTypes = parameters.flatMap { (it.type as? PsiClassType)?.parameters?.toList() ?: emptyList() }
+        val mentionedTypes = parameterTypes + method.returnType + genericTypes
+        val relatedTypes = mentionedTypes.filterIsInstance<PsiClassType>()
+        val resolvedClasses = relatedTypes.mapNotNull { it.resolve() }
+        val projectContentClasses = resolvedClasses.filter { isProjectContent(it) }
+        return projectContentClasses.toList()
+    }
+
+    private fun isProjectContent(element: PsiElement): Boolean {
+        val virtualFile = PsiUtil.getVirtualFile(element)
+        return virtualFile == null || ProjectFileIndex.getInstance(element.project).isInContent(virtualFile)
+    }
+}

java/src/main/kotlin/cc/unitmesh/idea/service/JavaWriteTestService.kt

@@ -94,12 +94,20 @@ class JavaWriteTestService : WriteTestService() {
             val elements = mutableListOf<ClassContext>()
             val projectPath = project.guessProjectDir()?.path
 
-            val resolvedClasses = JavaTypeUtil.resolveByMethod(element)
+            val resolvedClasses: MutableMap<String, PsiClass?> = mutableMapOf()
             resolvedClasses.putAll(JavaTypeUtil.resolveByField(element))
 
-            if (element is PsiClass) {
-                element.methods.forEach { method ->
-                    resolvedClasses.putAll(JavaTypeUtil.resolveByMethod(method))
+            when (element) {
+                is PsiClass -> {
+                    element.methods.forEach { method ->
+                        resolvedClasses.putAll(JavaTypeUtil.resolveByMethod(method))
+                    }
+                }
+
+                is PsiMethod -> {
+                    JavaRelatedContext.findRelatedClassesAndCleanUp(element).forEach { psiClass ->
+                        resolvedClasses[psiClass.name!!] = psiClass
+                    }
                 }
             }