package de.plushnikov.intellij.plugin.extension;

import com.intellij.ide.structureView.StructureViewExtension;
import com.intellij.ide.structureView.StructureViewTreeElement;
import com.intellij.ide.structureView.impl.java.JavaClassTreeElement;
import com.intellij.ide.structureView.impl.java.PsiFieldTreeElement;
import com.intellij.ide.structureView.impl.java.PsiMethodTreeElement;
import com.intellij.openapi.editor.Editor;
import com.intellij.psi.PsiClass;
import com.intellij.psi.PsiElement;
import de.plushnikov.intellij.plugin.psi.LombokLightClassBuilder;
import de.plushnikov.intellij.plugin.psi.LombokLightFieldBuilder;
import de.plushnikov.intellij.plugin.psi.LombokLightMethodBuilder;
import org.jetbrains.annotations.Nullable;

import java.util.Arrays;
import java.util.stream.Stream;

/**
 * Extension to populate StructureView with all fields, methods and inner classes generated by lombok
 */
public final class LombokStructureViewExtension implements StructureViewExtension {
  @Override
  public Class<? extends PsiElement> getType() {
    return PsiClass.class;
  }

  @Override
  public StructureViewTreeElement[] getChildren(PsiElement parent) {
    final PsiClass parentClass = (PsiClass) parent;

    final Stream<PsiFieldTreeElement> lombokFields = Arrays.stream(parentClass.getFields())
      .filter(LombokLightFieldBuilder.class::isInstance)
      .map(psiField -> new PsiFieldTreeElement(psiField, false));

    final Stream<PsiMethodTreeElement> lombokMethods = Arrays.stream(parentClass.getMethods())
      .filter(LombokLightMethodBuilder.class::isInstance)
      .map(psiMethod -> new PsiMethodTreeElement(psiMethod, false));

    final Stream<JavaClassTreeElement> lombokInnerClasses = Arrays.stream(parentClass.getInnerClasses())
      .filter(LombokLightClassBuilder.class::isInstance)
      .map(psiClass -> new JavaClassTreeElement(psiClass, false));

    return Stream.concat(Stream.concat(lombokFields, lombokMethods), lombokInnerClasses)
      .toArray(StructureViewTreeElement[]::new);
  }

  @Override
  public @Nullable Object getCurrentEditorElement(Editor editor, PsiElement parent) {
    return null;
  }
}
