// 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.psi.impl.source;

import com.intellij.lang.LighterAST;
import com.intellij.lang.LighterASTNode;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.psi.JavaTokenType;
import com.intellij.psi.impl.source.tree.ElementType;
import com.intellij.psi.impl.source.tree.LightTreeUtil;
import com.intellij.psi.tree.IElementType;
import com.intellij.util.containers.ContainerUtil;
import com.intellij.util.containers.JBIterable;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Collections;
import java.util.List;

import static com.intellij.psi.impl.source.tree.JavaElementType.*;

/**
 * Resolves unambiguous Java identifiers inside a file (if it can), using {@link LighterAST}. Can be used during indexing.
 */
public class FileLocalResolver {
  private final LighterAST myTree;

  public FileLocalResolver(@NotNull LighterAST tree) {
    myTree = tree;
  }

  public @NotNull LighterAST getLightTree() {
    return myTree;
  }

  /**
   * @param ref reference node
   * @return a resolve result corresponding to a local variable, parameter or field that the given reference resolves to.
   */
  public @NotNull LightResolveResult resolveLocally(@NotNull LighterASTNode ref) {
    String refName = JavaLightTreeUtil.getNameIdentifierText(myTree, ref);
    if (refName == null) return LightResolveResult.UNKNOWN;
    if (!canResolveToLocalVariable(myTree, ref)) return LightResolveResult.NON_LOCAL;

    boolean passedClass = false;
    LighterASTNode lastParent = ref;
    while (true) {
      ProgressManager.checkCanceled();
      LighterASTNode scope = myTree.getParent(lastParent);
      if (scope == null) return LightResolveResult.NON_LOCAL;

      for (LighterASTNode var : getDeclarations(scope, lastParent)) {
        if (refName.equals(JavaLightTreeUtil.getNameIdentifierText(myTree, var))) {
          if (passedClass) {
            IElementType type = var.getTokenType();
            return type == FIELD || type == RECORD_COMPONENT ? LightResolveResult.NON_LOCAL : LightResolveResult.UNKNOWN;
          }
          return LightResolveResult.resolved(var);
        }
      }

      lastParent = scope;
      passedClass |= scope.getTokenType() == CLASS || scope.getTokenType() == ANONYMOUS_CLASS;
    }
  }

  private boolean canResolveToLocalVariable(@NotNull LighterAST tree, @NotNull LighterASTNode ref) {
    LighterASTNode parent = tree.getParent(ref);
    return parent != null && parent.getTokenType() != METHOD_CALL_EXPRESSION && !hasQualifier(ref);
  }

  private boolean hasQualifier(@NotNull LighterASTNode ref) {
    return LightTreeUtil.firstChildOfType(myTree, ref, ElementType.EXPRESSION_BIT_SET) != null;
  }

  private @NotNull Iterable<LighterASTNode> getDeclarations(LighterASTNode scope, @Nullable LighterASTNode lastParent) {
    IElementType type = scope.getTokenType();
    if (type == CODE_BLOCK) {
      return walkChildrenScopes(before(LightTreeUtil.getChildrenOfType(myTree, scope, DECLARATION_STATEMENT), lastParent));
    }
    if (type == DECLARATION_STATEMENT) {
      return before(LightTreeUtil.getChildrenOfType(myTree, scope, LOCAL_VARIABLE), lastParent);
    }
    if (type == FOR_STATEMENT) {
      return walkChildrenScopes(before(LightTreeUtil.getChildrenOfType(myTree, scope, ElementType.JAVA_STATEMENT_BIT_SET), lastParent));
    }
    if (type == FOREACH_STATEMENT || type == CATCH_SECTION) {
      return JBIterable.of(LightTreeUtil.firstChildOfType(myTree, scope, PARAMETER));
    }
    if (type == TRY_STATEMENT) {
      return walkChildrenScopes(before(LightTreeUtil.getChildrenOfType(myTree, scope, RESOURCE_LIST), lastParent));
    }
    if (type == RESOURCE_LIST) {
      return before(LightTreeUtil.getChildrenOfType(myTree, scope, RESOURCE_VARIABLE), lastParent);
    }
    if (type == CLASS) {
      List<LighterASTNode> fields = LightTreeUtil.getChildrenOfType(myTree, scope, FIELD);
      LighterASTNode recordHeader = LightTreeUtil.firstChildOfType(myTree, scope, RECORD_HEADER);
      if (recordHeader == null) return fields;
      return JBIterable.from(fields).append(LightTreeUtil.getChildrenOfType(myTree, recordHeader, RECORD_COMPONENT));
    }
    if (type == LAMBDA_EXPRESSION || type == METHOD) {
      LighterASTNode paramList = LightTreeUtil.firstChildOfType(myTree, scope, PARAMETER_LIST);
      return paramList == null ? ContainerUtil.emptyList() : LightTreeUtil.getChildrenOfType(myTree, paramList, PARAMETER);
    }
    return Collections.emptyList();
  }

  private @NotNull JBIterable<LighterASTNode> walkChildrenScopes(JBIterable<? extends LighterASTNode> children) {
    return children.flatMap(child -> getDeclarations(child, null));
  }

  private static @NotNull JBIterable<LighterASTNode> before(List<LighterASTNode> children, @Nullable LighterASTNode lastParent) {
    return JBIterable.from(children).filter(node -> lastParent == null || node.getStartOffset() < lastParent.getStartOffset());
  }

  /**
   * Determine the type name of the given variable. Can be used to later iterate over all classes with this name for lightweight checks
   * before loading AST and fully resolving the type.
   * @param var Variable node
   * @return the short name of the class corresponding to the type of the variable, or null if the variable is not of class type or
   * the type is generic
   */
  public @Nullable String getShortClassTypeName(@NotNull LighterASTNode var) {
    return getShortClassTypeName(var, 0);
  }

  /**
   * Determine the type name of the given variable, unwrapping the expected array type. 
   * Can be used to later iterate over all classes with this name for lightweight checks
   * before loading AST and fully resolving the type.
   * @param var Variable node
   * @param arrayDepth Expected array depth
   * @return the short name of the class corresponding to the type of the variable, or null if the variable is not of class type or
   * the type is generic
   */
  public @Nullable String getShortClassTypeName(@NotNull LighterASTNode var, int arrayDepth) {
    LighterASTNode typeNode = LightTreeUtil.firstChildOfType(myTree, var, TYPE);
    while (arrayDepth > 0) {
      LighterASTNode bracket = LightTreeUtil.firstChildOfType(myTree, typeNode, JavaTokenType.LBRACKET);
      if (bracket == null) return null;
      typeNode = LightTreeUtil.firstChildOfType(myTree, typeNode, TYPE);
      arrayDepth--;
    }
    LighterASTNode typeRef = LightTreeUtil.firstChildOfType(myTree, typeNode, JAVA_CODE_REFERENCE);
    String refName = JavaLightTreeUtil.getNameIdentifierText(myTree, typeRef);
    if (refName == null) return null;

    if (LightTreeUtil.firstChildOfType(myTree, typeRef, JAVA_CODE_REFERENCE) != null) return refName;
    if (isTypeParameter(refName, var)) return null;
    return refName;
  }

  private boolean isTypeParameter(String name, LighterASTNode place) {
    LighterASTNode scope = place;
    while (scope != null) {
      if (scope.getTokenType() == CLASS || scope.getTokenType() == METHOD) {
        if (hasOwnTypeParameter(name, scope)) {
          return true;
        }
        if (isStatic(scope)) {
          break;
        }
      }
      scope = myTree.getParent(scope);
    }
    return false;
  }

  private boolean hasOwnTypeParameter(String name, LighterASTNode member) {
    LighterASTNode typeParams = LightTreeUtil.firstChildOfType(myTree, member, TYPE_PARAMETER_LIST);
    if (typeParams != null) {
      for (LighterASTNode param : LightTreeUtil.getChildrenOfType(myTree, typeParams, TYPE_PARAMETER)) {
        if (name.equals(JavaLightTreeUtil.getNameIdentifierText(myTree, param))) {
          return true;
        }
      }
    }
    return false;
  }

  private boolean isStatic(LighterASTNode scope) {
    LighterASTNode modList = LightTreeUtil.firstChildOfType(myTree, scope, MODIFIER_LIST);
    return modList != null && LightTreeUtil.firstChildOfType(myTree, modList, JavaTokenType.STATIC_KEYWORD) != null;
  }

  public static class LightResolveResult {
    /** We know nothing about the resolve result */
    public static final LightResolveResult UNKNOWN = new LightResolveResult();

    /** The result can't be determined, but it's definitely not a local variable/parameter */
    public static final LightResolveResult NON_LOCAL = new LightResolveResult();

    static @NotNull LightResolveResult resolved(@NotNull LighterASTNode target) {
      return new LightResolveResult() {
        @Override
        public LighterASTNode getTarget() {
          return target;
        }
      };
    }

    public @Nullable LighterASTNode getTarget() {
      return null;
    }
  }
}