// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.intellij.codeInsight;

import com.intellij.lang.ASTNode;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.Key;
import com.intellij.psi.*;
import com.intellij.psi.util.*;
import com.intellij.refactoring.util.RefactoringChangeUtil;
import com.intellij.util.IncorrectOperationException;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Objects;

public final class ChangeContextUtil {
  private static final Logger LOG = Logger.getInstance(ChangeContextUtil.class);

  private static final Key<ASTNode> HARD_REF_TO_AST = Key.create("HARD_REF_TO_AST");
  private static final Key<String> ENCODED_KEY = Key.create("ENCODED_KEY");
  private static final Key<PsiClass> THIS_QUALIFIER_CLASS_KEY = Key.create("THIS_QUALIFIER_CLASS_KEY");
  public static final Key<PsiMember> REF_MEMBER_KEY = Key.create("REF_MEMBER_KEY");
  public static final Key<Boolean> CAN_REMOVE_QUALIFIER_KEY = Key.create("CAN_REMOVE_QUALIFIER_KEY");
  public static final Key<PsiClass> REF_CLASS_KEY = Key.create("REF_CLASS_KEY");
  private static final Key<PsiClass> REF_MEMBER_THIS_CLASS_KEY = Key.create("REF_MEMBER_THIS_CLASS_KEY");

  private ChangeContextUtil() {}

  public static void encodeContextInfo(PsiElement scope, boolean includeRefClasses) {
    encodeContextInfo(scope, scope, includeRefClasses, true);
  }

  public static void encodeContextInfo(PsiElement scope, boolean includeRefClasses, boolean canChangeQualifier) {
    encodeContextInfo(scope, scope, includeRefClasses, canChangeQualifier);
  }

  private static void encodeContextInfo(PsiElement scope,
                                        PsiElement topLevelScope,
                                        boolean includeRefClasses,
                                        boolean canChangeQualifier) {
    if (scope instanceof StubBasedPsiElement) {
      // as long as "scope" is reachable, don't let GC collect AST together with all the copyable user data
      scope.putUserData(HARD_REF_TO_AST, scope.getNode());
    }

    if (scope instanceof PsiThisExpression thisExpr){
      scope.putCopyableUserData(ENCODED_KEY, "");

      final PsiJavaCodeReferenceElement qualifier = thisExpr.getQualifier();
      if (qualifier == null){
        PsiClass thisClass = RefactoringChangeUtil.getThisClass(thisExpr);
        if (thisClass != null && !(thisClass instanceof PsiAnonymousClass)){
          thisExpr.putCopyableUserData(THIS_QUALIFIER_CLASS_KEY, thisClass);
        }
      }
      else {
        final PsiElement resolved = qualifier.resolve();
        if (resolved instanceof PsiClass && resolved == topLevelScope) {
          thisExpr.putCopyableUserData(THIS_QUALIFIER_CLASS_KEY, (PsiClass)topLevelScope);
        }
      }
    }
    else if (scope instanceof PsiReferenceExpression refExpr){
      scope.putCopyableUserData(ENCODED_KEY, "");

      PsiExpression qualifier = refExpr.getQualifierExpression();
      if (qualifier == null){
        final JavaResolveResult resolveResult = refExpr.advancedResolve(false);
        final PsiElement refElement = resolveResult.getElement();
        if (refElement != null && !PsiTreeUtil.isAncestor(topLevelScope, refElement, false)){
          if (refElement instanceof PsiClass){
            if (includeRefClasses){
              refExpr.putCopyableUserData(REF_CLASS_KEY, (PsiClass)refElement);
            }
          }
          else if (refElement instanceof PsiMember && 
                   !Objects.equals(((PsiMember)refElement).getContainingClass(), topLevelScope)) {
            refExpr.putCopyableUserData(REF_MEMBER_KEY, ( (PsiMember)refElement));
            final PsiElement resolveScope = resolveResult.getCurrentFileResolveScope();
            if (resolveScope instanceof PsiClass && !PsiTreeUtil.isAncestor(topLevelScope, resolveScope, false)) {
              refExpr.putCopyableUserData(REF_MEMBER_THIS_CLASS_KEY, (PsiClass)resolveScope);
            }
          }
        }
      }
      else if (canChangeQualifier) {
        refExpr.putCopyableUserData(CAN_REMOVE_QUALIFIER_KEY, canRemoveQualifier(refExpr));
      }
    }
    else if (includeRefClasses) {
      PsiReference ref = scope.getReference();
      if (ref != null){
        scope.putCopyableUserData(ENCODED_KEY, "");

        PsiElement refElement = ref.resolve();
        if (refElement instanceof PsiClass && !PsiTreeUtil.isAncestor(topLevelScope, refElement, false)){
          scope.putCopyableUserData(REF_CLASS_KEY, (PsiClass)refElement);
        }
      }
    }

    for(PsiElement child = scope.getFirstChild(); child != null; child = child.getNextSibling()){
      encodeContextInfo(child, topLevelScope, includeRefClasses, canChangeQualifier);
    }
  }

  public static @NotNull PsiElement decodeContextInfo(@NotNull PsiElement scope,
                                                      @Nullable PsiClass thisClass,
                                                      @Nullable PsiExpression thisAccessExpr) throws IncorrectOperationException {
    if (scope instanceof StubBasedPsiElement) {
      scope.putUserData(HARD_REF_TO_AST, null);
    }
    if (scope.getCopyableUserData(ENCODED_KEY) != null) {
      scope.putCopyableUserData(ENCODED_KEY, null);

      if (scope instanceof PsiThisExpression thisExpr) {
        scope = decodeThisExpression(thisExpr, thisClass, thisAccessExpr);
      }
      else if (scope instanceof PsiReferenceExpression) {
        scope = decodeReferenceExpression((PsiReferenceExpression)scope, thisAccessExpr, thisClass);
      }
      else {
        PsiClass refClass = scope.getCopyableUserData(REF_CLASS_KEY);
        scope.putCopyableUserData(REF_CLASS_KEY, null);

        if (refClass != null && refClass.isValid()) {
          PsiReference ref = scope.getReference();
          if (ref != null) {
            final String qualifiedName = refClass.getQualifiedName();
            if (qualifiedName != null) {
              if (JavaPsiFacade.getInstance(refClass.getProject()).findClass(qualifiedName, scope.getResolveScope()) != null) {
                scope = ref.bindToElement(refClass);
              }
            }
          }
        }
      }
    }

    if (scope instanceof PsiClass) {
      if (thisAccessExpr != null) {
        thisAccessExpr = (PsiExpression)qualifyThis(thisAccessExpr, thisClass);
      }
    }

    PsiElement child = scope.getFirstChild();
    while (child != null) {
      child = decodeContextInfo(child, thisClass, thisAccessExpr).getNextSibling();
    }

    return scope;
  }

  private static PsiElement decodeThisExpression(PsiThisExpression thisExpr,
                                                 PsiClass thisClass,
                                                 PsiExpression thisAccessExpr) throws IncorrectOperationException {
    final PsiJavaCodeReferenceElement qualifier = thisExpr.getQualifier();
    PsiClass encodedQualifierClass = thisExpr.getCopyableUserData(THIS_QUALIFIER_CLASS_KEY);
    thisExpr.putCopyableUserData(THIS_QUALIFIER_CLASS_KEY, null);
    if (qualifier == null){
      if (encodedQualifierClass != null && encodedQualifierClass.isValid()){
        if (encodedQualifierClass.equals(thisClass) && thisAccessExpr != null && thisAccessExpr.isValid()){
          if (thisAccessExpr instanceof PsiThisExpression) {
            PsiJavaCodeReferenceElement thisAccessQualifier = ((PsiThisExpression)thisAccessExpr).getQualifier();
            PsiElement resolve = thisAccessQualifier != null ? thisAccessQualifier.resolve() : null;
            if (PsiTreeUtil.getParentOfType(thisExpr, PsiClass.class) == resolve) {
              return thisExpr;
            }
          }
          return maybeRemoveCast((PsiExpression)thisExpr.replace(thisAccessExpr));
        }
      }
    }
    else {
      PsiClass qualifierClass = (PsiClass)qualifier.resolve();
      if (encodedQualifierClass == qualifierClass && thisClass != null) {
        qualifier.bindToElement(thisClass);
      }
      else {
        if (qualifierClass != null) {
          if (qualifierClass.equals(thisClass) && thisAccessExpr != null && thisAccessExpr.isValid()) {
            return maybeRemoveCast((PsiExpression)thisExpr.replace(thisAccessExpr));
          }
        }
      }
    }
    return thisExpr;
  }

  private static PsiExpression maybeRemoveCast(@NotNull PsiExpression expression) {
    if (PsiUtil.skipParenthesizedExprDown(expression) instanceof PsiTypeCastExpression cast &&
        RedundantCastUtil.isCastRedundant(cast)) {
      expression = (PsiExpression)expression.replace(Objects.requireNonNull(cast.getOperand()));
    }
    if (expression.getParent() instanceof PsiParenthesizedExpression parens &&
        parens.getParent() instanceof PsiExpression grandParent &&
        PsiPrecedenceUtil.areParenthesesNeeded(grandParent, expression, false)) {
      return (PsiExpression)parens.replace(expression);
    }
    return expression;
  }

  private static PsiReferenceExpression decodeReferenceExpression(@NotNull PsiReferenceExpression refExpr,
                                                                  PsiExpression thisAccessExpr,
                                                                  PsiClass thisClass) {
    PsiManager manager = refExpr.getManager();
    PsiElementFactory factory = JavaPsiFacade.getElementFactory(manager.getProject());

    PsiExpression qualifier = refExpr.getQualifierExpression();
    if (qualifier == null) {
      PsiMember refMember = refExpr.getCopyableUserData(REF_MEMBER_KEY);
      refExpr.putCopyableUserData(REF_MEMBER_KEY, null);

      if (refMember != null && refMember.isValid()) {
        PsiClass containingClass = refMember.getContainingClass();
        if (containingClass != null && containingClass.isValid() && refMember.hasModifierProperty(PsiModifier.STATIC)) {
          PsiElement refElement = refExpr.resolve();
          if (!manager.areElementsEquivalent(refMember, refElement)) {
            final PsiClass currentClass = PsiTreeUtil.getParentOfType(refExpr, PsiClass.class);
            if (!InheritanceUtil.isInheritorOrSelf(currentClass, containingClass, true) || containingClass.isInterface()) {
              refExpr.setQualifierExpression(factory.createReferenceExpression(containingClass));
            }
          }
        }
        else {
          final PsiClass realParentClass = refExpr.getCopyableUserData(REF_MEMBER_THIS_CLASS_KEY);
          refExpr.putCopyableUserData(REF_MEMBER_THIS_CLASS_KEY, null);
          if (thisAccessExpr != null && thisClass != null && realParentClass != null &&
              InheritanceUtil.isInheritorOrSelf(thisClass, realParentClass, true)) {
            boolean needQualifier = true;
            PsiElement refElement = refExpr.resolve();
            if (refMember.equals(refElement) ||
                (refElement instanceof PsiMethod && refMember instanceof PsiMethod &&
                 MethodSignatureUtil.isSuperMethod((PsiMethod)refMember, (PsiMethod)refElement))) {
              if (thisAccessExpr instanceof PsiThisExpression thisExpression && thisExpression.getQualifier() == null) {
                //Trivial qualifier
                needQualifier = false;
              }
              else {
                final PsiClass currentClass = findThisClass(refExpr, refMember);
                if (thisAccessExpr instanceof PsiThisExpression thisExpression){
                  PsiJavaCodeReferenceElement thisQualifier = thisExpression.getQualifier();
                  PsiClass thisExprClass = thisQualifier != null
                                           ? (PsiClass)thisQualifier.resolve()
                                           : RefactoringChangeUtil.getThisClass(refExpr);
                  if (thisExprClass != null && (thisExprClass.equals(currentClass) || thisExprClass.isInheritor(realParentClass, true))){ // qualifier is not necessary
                    needQualifier = false;
                  }
                }
              }
            }

            if (needQualifier){
              refExpr.setQualifierExpression(thisAccessExpr);
              maybeRemoveCast(refExpr.getQualifierExpression());
            }
          }
          else if (thisClass != null && realParentClass != null && PsiTreeUtil.isAncestor(realParentClass, thisClass, true)) {
            PsiElement refElement = refExpr.resolve();
            if (refElement != null && !manager.areElementsEquivalent(refMember, refElement)) {
              refExpr = RefactoringChangeUtil.qualifyReference(refExpr, refMember, null);
            }
          }
        }
      }
      else {
        PsiClass refClass = refExpr.getCopyableUserData(REF_CLASS_KEY);
        refExpr.putCopyableUserData(REF_CLASS_KEY, null);
        if (refClass != null && refClass.isValid()){
          refExpr = (PsiReferenceExpression)refExpr.bindToElement(refClass);
        }
      }
    }
    else{
      Boolean couldRemove = refExpr.getCopyableUserData(CAN_REMOVE_QUALIFIER_KEY);
      refExpr.putCopyableUserData(CAN_REMOVE_QUALIFIER_KEY, null);

      if (couldRemove == Boolean.FALSE && canRemoveQualifier(refExpr)){
        PsiReferenceExpression newRefExpr = (PsiReferenceExpression)factory.createExpressionFromText(refExpr.getReferenceName(), null);
        refExpr = (PsiReferenceExpression)refExpr.replace(newRefExpr);
      }
    }
    return refExpr;
  }

  private static PsiClass findThisClass(PsiReferenceExpression refExpr, PsiMember refMember) {
    LOG.assertTrue(refExpr.getQualifierExpression() == null);
    final PsiClass refMemberClass = refMember.getContainingClass();
    if (refMemberClass == null) return null;
    PsiElement parent = refExpr.getContext();
    while(parent != null){
      if (parent instanceof PsiClass){
        if (parent.equals(refMemberClass) || ((PsiClass)parent).isInheritor(refMemberClass, true)){
          return (PsiClass)parent;
        }
      }
      parent = parent.getContext();
    }

    return refMemberClass;
  }

  public static boolean canRemoveQualifier(PsiReferenceExpression refExpr) {
    try{
      PsiExpression qualifier = refExpr.getQualifierExpression();
      if (!(qualifier instanceof PsiReferenceExpression)) return false;
      if (refExpr.getTypeParameters().length > 0) return false;
      PsiElement qualifierRefElement = ((PsiReferenceExpression)qualifier).resolve();
      if (!(qualifierRefElement instanceof PsiClass)) return false;
      PsiElement refElement = refExpr.resolve();
      if (refElement == null) return false;
      PsiElementFactory factory = JavaPsiFacade.getElementFactory(refExpr.getProject());
      if (refExpr.getParent() instanceof PsiMethodCallExpression methodCall){
        PsiMethodCallExpression newMethodCall = (PsiMethodCallExpression)factory.createExpressionFromText(
          refExpr.getReferenceName() + "()", refExpr);
        newMethodCall.getArgumentList().replace(methodCall.getArgumentList());
        PsiElement newRefElement = newMethodCall.getMethodExpression().resolve();
        return refElement.equals(newRefElement);
      }
      else if (refExpr instanceof PsiMethodReferenceExpression) {
        return false;
      }
      else {
        PsiReferenceExpression newRefExpr = (PsiReferenceExpression)factory.createExpressionFromText(refExpr.getReferenceName(), refExpr);
        PsiElement newRefElement = newRefExpr.resolve();
        return refElement.equals(newRefElement);
      }
    }
    catch(IncorrectOperationException e){
      LOG.error(e);
      return false;
    }
  }

  private static PsiElement qualifyThis(PsiElement scope, PsiClass thisClass) throws IncorrectOperationException {
    if (scope instanceof PsiThisExpression thisExpr){
      if (thisExpr.getQualifier() == null){
        if (thisClass instanceof PsiAnonymousClass) return null;
        return RefactoringChangeUtil.createThisExpression(thisClass.getManager(), thisClass);
      }
    }
    else if (!(scope instanceof PsiClass)){
      for(PsiElement child = scope.getFirstChild(); child != null; child = child.getNextSibling()){
        if (qualifyThis(child, thisClass) == null) return null;
      }
    }
    return scope;
  }

  public static void clearContextInfo(PsiElement scope) {
    if (scope instanceof StubBasedPsiElement) {
      scope.putUserData(HARD_REF_TO_AST, null);
    }
    scope.putCopyableUserData(ENCODED_KEY, null);
    scope.putCopyableUserData(THIS_QUALIFIER_CLASS_KEY, null);
    scope.putCopyableUserData(REF_MEMBER_KEY, null);
    scope.putCopyableUserData(CAN_REMOVE_QUALIFIER_KEY, null);
    scope.putCopyableUserData(REF_CLASS_KEY, null);
    scope.putCopyableUserData(REF_MEMBER_THIS_CLASS_KEY, null);
    for(PsiElement child = scope.getFirstChild(); child != null; child = child.getNextSibling()){
      clearContextInfo(child);
    }
  }
}
