package com.jetbrains.python.psi.impl;

import com.intellij.lang.ASTNode;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiPolyVariantReference;
import com.intellij.psi.PsiReference;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.resolve.PyResolveContext;
import com.jetbrains.python.psi.types.PyCollectionType;
import com.jetbrains.python.psi.types.PySubscriptableType;
import com.jetbrains.python.psi.types.PyType;
import com.jetbrains.python.psi.types.TypeEvalContext;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/**
 * @author yole
 */
public class PySubscriptionExpressionImpl extends PyElementImpl implements PySubscriptionExpression {
  public PySubscriptionExpressionImpl(ASTNode astNode) {
    super(astNode);
  }

  public PyExpression getOperand() {
    return childToPsiNotNull(PyElementTypes.EXPRESSIONS, 0);
  }

  @Nullable
  public PyExpression getIndexExpression() {
    return childToPsi(PyElementTypes.EXPRESSIONS, 1);
  }

  @Override
  protected void acceptPyVisitor(final PyElementVisitor pyVisitor) {
    pyVisitor.visitPySubscriptionExpression(this);
  }

  @Nullable
  @Override
  public PyType getType(@NotNull TypeEvalContext context) {
    PyType res = null;
    final PsiReference ref = getReference(PyResolveContext.noImplicits().withTypeEvalContext(context));
    if (ref != null) {
      final PsiElement resolved = ref.resolve();
      if (resolved instanceof Callable) {
        res = ((Callable)resolved).getReturnType(context, null);
      }
    }
    if (res == null) {
      final PyExpression indexExpression = getIndexExpression();
      if (indexExpression != null) {
        final PyType type = context.getType(getOperand());
        if (type instanceof PySubscriptableType) {
          res = ((PySubscriptableType)type).getElementType(indexExpression, context);
        }
        else if (type instanceof PyCollectionType) {
          res = ((PyCollectionType) type).getElementType(context);
        }
      }
    }
    return res;
  }

  @Override
  public PsiReference getReference() {
    return getReference(PyResolveContext.noImplicits());
  }

  @Override
  public PsiPolyVariantReference getReference(PyResolveContext context) {
    return new PyOperatorReferenceImpl(this, context);
  }

  @Override
  public PyExpression getQualifier() {
    return getOperand();
  }

  @Override
  public String getReferencedName() {
    String res = PyNames.GETITEM;
    switch (AccessDirection.of(this)) {
      case READ:
        res = PyNames.GETITEM;
        break;
      case WRITE:
        res = PyNames.SETITEM;
        break;
      case DELETE:
        res = PyNames.DELITEM;
        break;
    }
    return res;
  }

  @Override
  public ASTNode getNameElement() {
    return getNode().findChildByType(PyTokenTypes.LBRACKET);
  }
}

@@ -448,184 +448,9 @@
 pes.
-PyCollectionType;%0Aimport com.jetbrains.python.psi.types.PySubscriptableType;%0Aimport com.jetbrains.python.psi.types.PyType;%0Aimport com.jetbrains.python.psi.types.TypeEvalContext
+*
 ;%0Aim
@@ -1516,24 +1516,53 @@
 (res == null
+ %7C%7C res instanceof PyNoneType
 ) %7B%0A      fi

Return type of elements in sequence if __getattr__ returns None, required for Python stdlib type database