/*
 * Minecraft Development for IntelliJ
 *
 * https://mcdev.io/
 *
 * Copyright (C) 2023 minecraft-dev
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, version 3.0 only.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 
import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
import com.demonwav.mcdev.platform.mixin.reference.MixinSelectorParser
import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.AT
import com.demonwav.mcdev.platform.mixin.util.shortName
import com.demonwav.mcdev.util.MemberReference
import com.demonwav.mcdev.util.constantStringValue
import com.demonwav.mcdev.util.descriptor
import com.demonwav.mcdev.util.fullQualifiedName
import com.demonwav.mcdev.util.internalName
import com.demonwav.mcdev.util.mapToArray
import com.demonwav.mcdev.util.shortName
import com.intellij.codeInsight.completion.JavaLookupElementBuilder
import com.intellij.codeInsight.lookup.LookupElementBuilder
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.text.StringUtil
import com.intellij.psi.PsiAnnotation
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiMember
import com.intellij.psi.PsiMethod
import com.intellij.psi.PsiModifier
import com.intellij.psi.PsiNewExpression
import com.intellij.psi.PsiSubstitutor
import com.intellij.psi.util.parentOfType
import org.objectweb.asm.Opcodes
import org.objectweb.asm.Type
import org.objectweb.asm.tree.AbstractInsnNode
import org.objectweb.asm.tree.ClassNode
import org.objectweb.asm.tree.MethodInsnNode
import org.objectweb.asm.tree.MethodNode
import org.objectweb.asm.tree.TypeInsnNode
 
class NewInsnInjectionPoint : InjectionPoint<PsiMember>() {
    private fun getTarget(at: PsiAnnotation, target: MixinSelector?): MixinSelector? {
        if (target != null) {
            return target
        }
        val clazz = AtResolver.getArgs(at)["class"] ?: return null
        return classToMemberReference(clazz)
    }
 
    override fun createNavigationVisitor(
        at: PsiAnnotation,
        target: MixinSelector?,
        targetClass: PsiClass,
    ): NavigationVisitor? {
        return getTarget(at, target)?.let { MyNavigationVisitor(it) }
    }
 
    override fun doCreateCollectVisitor(
        at: PsiAnnotation,
        target: MixinSelector?,
        targetClass: ClassNode,
        mode: CollectVisitor.Mode,
    ): CollectVisitor<PsiMember>? {
        if (mode == CollectVisitor.Mode.COMPLETION) {
            return MyCollectVisitor(mode, at.project, MemberReference(""))
        }
        return getTarget(at, target)?.let { MyCollectVisitor(mode, at.project, it) }
    }
 
    override fun createLookup(targetClass: ClassNode, result: CollectVisitor.Result<PsiMember>): LookupElementBuilder? {
        when (val target = result.target) {
            is PsiClass -> {
                return JavaLookupElementBuilder.forClass(target, target.internalName)
                    .withPresentableText(target.shortName ?: return null)
            }
            is PsiMethod -> {
                val ownerName = result.qualifier?.substringAfterLast('.')?.replace('$', '.') ?: targetClass.shortName
                val descriptorArgs = target.descriptor?.dropLast(1) ?: return null
                val qualifierInternalName = result.qualifier?.replace('.', '/')
                return JavaLookupElementBuilder.forMethod(
                    target,
                    "${descriptorArgs}L$qualifierInternalName;",
                    PsiSubstitutor.EMPTY,
                    null,
                )
                    .setBoldIfInClass(target, targetClass)
                    .withPresentableText(ownerName + "." + target.internalName)
                    .withLookupString(target.internalName)
            }
            else -> return null
        }
    }
 
    private class MyNavigationVisitor(
        private val selector: MixinSelector,
    ) : NavigationVisitor() {
        override fun visitNewExpression(expression: PsiNewExpression) {
            val anonymousName = expression.anonymousClass?.fullQualifiedName?.replace('.', '/')
            if (anonymousName != null) {
                // guess descriptor
                val hasThis = expression.parentOfType<PsiMethod>()?.hasModifierProperty(PsiModifier.STATIC) == false
                val thisType = if (hasThis) expression.parentOfType<PsiClass>()?.internalName else null
                val argTypes = expression.argumentList?.expressionTypes?.map { it.descriptor } ?: emptyList()
                val bytecodeArgTypes = if (thisType != null) listOf(thisType) + argTypes else argTypes
                val methodDesc = Type.getMethodDescriptor(
                    Type.VOID_TYPE,
                    *bytecodeArgTypes.mapToArray { Type.getType(it) },
                )
                if (selector.matchMethod(anonymousName, "<init>", methodDesc)) {
                    addResult(expression)
                }
            } else {
                val ctor = expression.resolveConstructor()
                val containingClass = ctor?.containingClass
                if (ctor != null && containingClass != null) {
                    if (selector.matchMethod(ctor, containingClass)) {
                        addResult(expression)
                    }
                }
            }
            super.visitNewExpression(expression)
        }
    }
 
    private class MyCollectVisitor(
        mode: Mode,
        private val project: Project,
        private val selector: MixinSelector,
    ) : CollectVisitor<PsiMember>(mode) {
        override fun accept(methodNode: MethodNode) {
            val insns = methodNode.instructions ?: return
            insns.iterator().forEachRemaining { insn ->
                if (insn !is TypeInsnNode) return@forEachRemaining
                if (insn.opcode != Opcodes.NEW) return@forEachRemaining
                val initCall = findInitCall(insn) ?: return@forEachRemaining
 
                val sourceMethod = nodeMatchesSelector(initCall, mode, selector, project) ?: return@forEachRemaining
                addResult(
                    insn,
                    sourceMethod,
                    qualifier = initCall.owner.replace('/', '.'),
                )
            }
        }
    }
 
    companion object {
        fun findInitCall(newInsn: TypeInsnNode): MethodInsnNode? {
            var newInsns = 0
            var insn: AbstractInsnNode? = newInsn
            while (insn != null) {
                when (insn) {
                    is TypeInsnNode -> {
                        if (insn.opcode == Opcodes.NEW) {
                            newInsns++
                        }
                    }
                    is MethodInsnNode -> {
                        if (insn.opcode == Opcodes.INVOKESPECIAL && insn.name == "<init>") {
                            newInsns--
                            if (newInsns == 0) {
                                return insn
                            }
                        }
                    }
                }
                insn = insn.next
            }
 
            return null
        }
    }
}
 
class NewInsnSelectorParser : MixinSelectorParser {
    override fun parse(value: String, context: PsiElement): MixinSelector? {
        // check we're inside NEW
        val at = context.parentOfType<PsiAnnotation>() ?: return null
        if (!at.hasQualifiedName(AT)) return null
        if (at.findAttributeValue("value")?.constantStringValue != "NEW") return null
 
        return NewInsnSelector(value)
    }
}
 
private class NewInsnSelector(
    override val methodDescriptor: String,
) : MixinSelector {
    override fun matchField(owner: String, name: String, desc: String): Boolean = false
 
    override fun matchMethod(owner: String, name: String, desc: String): Boolean {
        if (name != "<init>" || desc.last() != 'V') {
            return false
        }
 
        val lastParen = methodDescriptor.lastIndexOf(')')
        val argsDesc = methodDescriptor.substring(0, lastParen + 1)
        val descRet = methodDescriptor.substringAfterLast(')').removeSurrounding("L", ";")
        return desc.dropLast(1) == argsDesc && descRet == owner
    }
 
    override val owner = null
    override val fieldDescriptor = null
    override val displayName = methodDescriptor
}
 
private fun classToMemberReference(value: String): MemberReference? {
    val fqn = value.replace('/', '.').replace('$', '.')
    if (fqn.isNotEmpty() && !fqn.startsWith('.') && !fqn.endsWith('.') && !fqn.contains("..")) {
        if (StringUtil.isJavaIdentifier(fqn.replace('.', '_'))) {
            return MemberReference("<init>", owner = fqn)
        }
    }
 
    return null
}

      

        
/*
 * Minecraft Development for IntelliJ
 *
 * https://mcdev.io/
 *
 * Copyright (C) 2023 minecraft-dev
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, version 3.0 only.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 
import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
import com.demonwav.mcdev.platform.mixin.reference.MixinSelectorParser
import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.AT
import com.demonwav.mcdev.platform.mixin.util.shortName
import com.demonwav.mcdev.util.MemberReference
import com.demonwav.mcdev.util.constantStringValue
import com.demonwav.mcdev.util.descriptor
import com.demonwav.mcdev.util.fullQualifiedName
import com.demonwav.mcdev.util.internalName
import com.demonwav.mcdev.util.mapToArray
import com.demonwav.mcdev.util.shortName
import com.intellij.codeInsight.completion.JavaLookupElementBuilder
import com.intellij.codeInsight.lookup.LookupElementBuilder
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.text.StringUtil
import com.intellij.psi.PsiAnnotation
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiMember
import com.intellij.psi.PsiMethod
import com.intellij.psi.PsiModifier
import com.intellij.psi.PsiNewExpression
import com.intellij.psi.PsiSubstitutor
import com.intellij.psi.util.parentOfType
import org.objectweb.asm.Opcodes
import org.objectweb.asm.Type
import org.objectweb.asm.tree.AbstractInsnNode
import org.objectweb.asm.tree.ClassNode
import org.objectweb.asm.tree.MethodInsnNode
import org.objectweb.asm.tree.MethodNode
import org.objectweb.asm.tree.TypeInsnNode
 
class NewInsnInjectionPoint : InjectionPoint<PsiMember>() {
    private fun getTarget(at: PsiAnnotation, target: MixinSelector?): MixinSelector? {
        if (target != null) {
            return target
        }
        val clazz = AtResolver.getArgs(at)["class"] ?: return null
        return classToMemberReference(clazz)
    }
 
    override fun createNavigationVisitor(
        at: PsiAnnotation,
        target: MixinSelector?,
        targetClass: PsiClass,
    ): NavigationVisitor? {
        return getTarget(at, target)?.let { MyNavigationVisitor(it) }
    }
 
    override fun doCreateCollectVisitor(
        at: PsiAnnotation,
        target: MixinSelector?,
        targetClass: ClassNode,
        mode: CollectVisitor.Mode,
    ): CollectVisitor<PsiMember>? {
        if (mode == CollectVisitor.Mode.COMPLETION) {
            return MyCollectVisitor(mode, at.project, MemberReference(""))
        }
        return getTarget(at, target)?.let { MyCollectVisitor(mode, at.project, it) }
    }
 
    override fun createLookup(targetClass: ClassNode, result: CollectVisitor.Result<PsiMember>): LookupElementBuilder? {
        when (val target = result.target) {
            is PsiClass -> {
                return JavaLookupElementBuilder.forClass(target, target.internalName)
                    .withPresentableText(target.shortName ?: return null)
            }
            is PsiMethod -> {
                val ownerName = result.qualifier?.substringAfterLast('.')?.replace('$', '.') ?: targetClass.shortName
                val descriptorArgs = target.descriptor?.dropLast(1) ?: return null
                val qualifierInternalName = result.qualifier?.replace('.', '/')
                return JavaLookupElementBuilder.forMethod(
                    target,
                    "${descriptorArgs}L$qualifierInternalName;",
                    PsiSubstitutor.EMPTY,
                    null,
                )
                    .setBoldIfInClass(target, targetClass)
                    .withPresentableText(ownerName + "." + target.internalName)
                    .withLookupString(target.internalName)
            }
            else -> return null
        }
    }
 
    private class MyNavigationVisitor(
        private val selector: MixinSelector,
    ) : NavigationVisitor() {
        override fun visitNewExpression(expression: PsiNewExpression) {
            val anonymousName = expression.anonymousClass?.fullQualifiedName?.replace('.', '/')
            if (anonymousName != null) {
                // guess descriptor
                val hasThis = expression.parentOfType<PsiMethod>()?.hasModifierProperty(PsiModifier.STATIC) == false
                val thisType = if (hasThis) expression.parentOfType<PsiClass>()?.internalName else null
                val argTypes = expression.argumentList?.expressionTypes?.map { it.descriptor } ?: emptyList()
                val bytecodeArgTypes = if (thisType != null) listOf(thisType) + argTypes else argTypes
                val methodDesc = Type.getMethodDescriptor(
                    Type.VOID_TYPE,
                    *bytecodeArgTypes.mapToArray { Type.getType(it) },
                )
                if (selector.matchMethod(anonymousName, "<init>", methodDesc)) {
                    addResult(expression)
                }
            } else {
                val ctor = expression.resolveConstructor()
                val containingClass = ctor?.containingClass
                if (ctor != null && containingClass != null) {
                    if (selector.matchMethod(ctor, containingClass)) {
                        addResult(expression)
                    }
                }
            }
            super.visitNewExpression(expression)
        }
    }
 
    private class MyCollectVisitor(
        mode: Mode,
        private val project: Project,
        private val selector: MixinSelector,
    ) : CollectVisitor<PsiMember>(mode) {
        override fun accept(methodNode: MethodNode) {
            val insns = methodNode.instructions ?: return
            insns.iterator().forEachRemaining { insn ->
                if (insn !is TypeInsnNode) return@forEachRemaining
                if (insn.opcode != Opcodes.NEW) return@forEachRemaining
                val initCall = findInitCall(insn) ?: return@forEachRemaining
 
                val sourceMethod = nodeMatchesSelector(initCall, mode, selector, project) ?: return@forEachRemaining
                addResult(
                    insn,
                    sourceMethod,
                    qualifier = initCall.owner.replace('/', '.'),
                )
            }
        }
    }
 
    companion object {
        fun findInitCall(newInsn: TypeInsnNode): MethodInsnNode? {
            var newInsns = 0
            var insn: AbstractInsnNode? = newInsn
            while (insn != null) {
                when (insn) {
                    is TypeInsnNode -> {
                        if (insn.opcode == Opcodes.NEW) {
                            newInsns++
                        }
                    }
                    is MethodInsnNode -> {
                        if (insn.opcode == Opcodes.INVOKESPECIAL && insn.name == "<init>") {
                            newInsns--
                            if (newInsns == 0) {
                                return insn
                            }
                        }
                    }
                }
                insn = insn.next
            }
 
            return null
        }
    }
}
 
class NewInsnSelectorParser : MixinSelectorParser {
    override fun parse(value: String, context: PsiElement): MixinSelector? {
        // check we're inside NEW
        val at = context.parentOfType<PsiAnnotation>() ?: return null
        if (!at.hasQualifiedName(AT)) return null
        if (at.findAttributeValue("value")?.constantStringValue != "NEW") return null
 
        return NewInsnSelector(value)
    }
}
 
private class NewInsnSelector(
    override val methodDescriptor: String,
) : MixinSelector {
    override fun matchField(owner: String, name: String, desc: String): Boolean = false
 
    override fun matchMethod(owner: String, name: String, desc: String): Boolean {
        if (name != "<init>" || desc.last() != 'V') {
            return false
        }
 
        val lastParen = methodDescriptor.lastIndexOf(')')
        val argsDesc = methodDescriptor.substring(0, lastParen + 1)
        val descRet = methodDescriptor.substringAfterLast(')').removeSurrounding("L", ";")
        return desc.dropLast(1) == argsDesc && descRet == owner
    }
 
    override val owner = null
    override val fieldDescriptor = null
    override val displayName = methodDescriptor
}
 
private fun classToMemberReference(value: String): MemberReference? {
    val fqn = value.replace('/', '.')
    if (fqn.isNotEmpty() && !fqn.startsWith('.') && !fqn.endsWith('.') && !fqn.contains("..")) {
        if (StringUtil.isJavaIdentifier(fqn.replace('.', '_'))) {
            return MemberReference("<init>", owner = fqn)
        }
    }
 
    return null
}