package fuchsia.developer.plugin.fidl;

import com.google.common.collect.ImmutableSet;
import com.intellij.lang.ASTNode;
import com.intellij.lang.annotation.Annotation;
import com.intellij.lang.annotation.AnnotationHolder;
import com.intellij.lang.annotation.Annotator;
import com.intellij.openapi.editor.colors.TextAttributesKey;
import com.intellij.psi.PsiElement;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.tree.TokenSet;
import fuchsia.developer.plugin.fidl.psi.Types;
import java.util.Arrays;
import java.util.Set;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class ContextAwareHighlighter implements Annotator {

  private static final TextAttributesKey KEYWORD_ATTRIBUTE;
  private static final Set<String> TYPE_KEYWORDS;
  private static final TextAttributesKey IDENTIFIER_ATTRIBUTE;
  private static final SyntaxHighlighter SYNTAX_HIGHLIGHTER;
  private static final TokenSet INTEGRAL_TYPES =
      TokenSet.create(
          Types.INT8,
          Types.INT16,
          Types.INT32,
          Types.INT64,
          Types.UINT8,
          Types.UINT16,
          Types.UINT32,
          Types.UINT64);
  private static final TokenSet UNSIGNED_INTEGRAL_TYPES =
      TokenSet.create(Types.UINT8, Types.UINT16, Types.UINT32, Types.UINT64);

  static {
    SYNTAX_HIGHLIGHTER = new SyntaxHighlighter();
    TextAttributesKey[] key = SYNTAX_HIGHLIGHTER.getKeywordHighlights();
    KEYWORD_ATTRIBUTE = key[0];

    ImmutableSet.Builder<String> builder = ImmutableSet.builder();
    builder.addAll(
        Arrays.asList(
            "array",
            "vector",
            "string",
            "handle",
            "process",
            "thread",
            "vmo",
            "channel",
            "clock",
            "event",
            "port",
            "interrupt",
            "log",
            "socket",
            "resource",
            "eventpair",
            "exception",
            "job",
            "vmar",
            "fifo",
            "guest",
            "timer",
            "request",
            "iommu",
            "pager",
            "pcidevice",
            "pmt",
            "suspendtoken",
            "vcpu",
            "strict",
            "bool",
            "float32",
            "float64",
            "int8",
            "int16",
            "int32",
            "int64",
            "uint8",
            "uint16",
            "uint32",
            "uint64"));
    TYPE_KEYWORDS = builder.build();

    key = SYNTAX_HIGHLIGHTER.getTokenHighlights(Types.IDENTIFIER);
    IDENTIFIER_ATTRIBUTE = key[0];
  }

  private static class NumberAndRadix {
    String representation;
    int radix;
  }

  private static NumberAndRadix numberAndRadix(ASTNode literalNode) {
    NumberAndRadix out = new NumberAndRadix();
    // Assume decimal
    out.representation = literalNode.getText();
    out.radix = 10;
    if (literalNode.findChildByType(Types.BINARY_INTEGRAL_LITERAL) != null) {
      String[] pieces = out.representation.split("0[bB]");
      out.representation = pieces[0] + pieces[1];
      out.radix = 2;
    } else if (literalNode.findChildByType(Types.HEX_INTEGRAL_LITERAL) != null) {
      String[] pieces = out.representation.split("0[xX]");
      out.representation = pieces[0] + pieces[1];
      out.radix = 16;
    }
    return out;
  }

  @Nullable
  private static Long signedLong(ASTNode literalNode) {
    NumberAndRadix val = numberAndRadix(literalNode);
    Long lval;
    try {
      lval = Long.parseLong(val.representation, val.radix);
    } catch (NumberFormatException e) {
      // This may happen if you say something like "0B1234"
      return null;
    }
    return lval;
  }

  @Nullable
  private static Long unsignedLong(ASTNode literalNode) {
    NumberAndRadix val = numberAndRadix(literalNode);
    Long lval;
    try {
      lval = Long.parseUnsignedLong(val.representation, val.radix);
    } catch (NumberFormatException e) {
      // This may happen if you say something like "0B1234"
      return null;
    }
    return lval;
  }

  /**
   * @param literalNode An ASTNode of type NUMERIC_LITERAL.
   * @return null if the value associated with the node isn't supported or is an unsigned power of
   *     two, an error message otherwise.
   */
  private static String unsignedLongPowerOfTwoOrError(ASTNode literalNode) {
    Long lval = unsignedLong(literalNode);
    if (lval == null || Long.compareUnsigned(lval, 0) < 0 || Long.bitCount(lval) != 1) {
      return "Bit value must be non-negative power of two, is  " + literalNode.getText();
    }
    return null;
  }

  /**
   * Checks the literal is of the (currently numeric) type given in typeConstructor.
   *
   * @param typeConstructor ASTNode containing the type the element is supposed to be.
   * @param literal ASTNode containing the element to check
   * @return A string describing the error, or null if there was no error
   */
  private static String correctTypeOrError(@Nullable ASTNode typeConstructor, ASTNode literal) {
    String type;
    if (typeConstructor == null) {
      // Only happens with enums; this is default value for enum.
      type = "uint32";
    } else {
      type = typeConstructor.getText();
    }
    boolean error = false;
    Long val;
    switch (type) {
      case "int8":
        val = signedLong(literal);
        error = val == null || (val != val.byteValue());
        break;
      case "int16":
        val = signedLong(literal);
        error = val == null || (val != val.shortValue());
        break;
      case "int32":
        val = signedLong(literal);
        error = val == null || (val != val.intValue());
        break;
      case "int64":
        val = signedLong(literal);
        error = val == null;
        break;
      case "uint8":
        val = unsignedLong(literal);
        error = val == null || (val & ~0xFFL) != 0;
        break;
      case "uint16":
        val = unsignedLong(literal);
        error = val == null || (val & ~0xFFFFL) != 0;
        break;
      case "uint32":
        val = unsignedLong(literal);
        error = val == null || (val & ~0xFFFFFFFFL) != 0;
        break;
      case "uint64":
        val = unsignedLong(literal);
        error = val == null;
        break;
    }
    String result = null;
    if (error) {
      result = "Expected value of type " + typeConstructor.getText() + ", got " + literal.getText();
    }
    return result;
  }

  @Override
  public void annotate(@NotNull PsiElement element, @NotNull AnnotationHolder holder) {
    PsiElement parent = element.getParent();
    if (parent == null || parent.getNode() == null) {
      return;
    }
    IElementType parentType = parent.getNode().getElementType();
    IElementType thisType = element.getNode().getElementType();

    // Context-sensitive keywords: if the token is used in the place of an identifier, and the
    // SyntaxHighlighter might highlight it, then make sure it is colored as an identifier
    if (parentType == Types.IDENTIFIER_TOKEN
        && SYNTAX_HIGHLIGHTER.getTokenHighlights(thisType).length != 0) {
      Annotation annotation = holder.createInfoAnnotation(element, null);
      annotation.setTextAttributes(IDENTIFIER_ATTRIBUTE);
    }

    // Context-sensitive keywords, part II: The Syntax Highlighter does not color types.  If we
    // think this identifier is a type, color it.
    if (parentType == Types.COMPOUND_IDENTIFIER) {
      IElementType grandParentType = parent.getParent().getNode().getElementType();
      if (grandParentType == Types.TYPE_CONSTRUCTOR) {
        if (TYPE_KEYWORDS.contains(element.getText())) {
          Annotation annotation = holder.createInfoAnnotation(element, null);
          annotation.setTextAttributes(KEYWORD_ATTRIBUTE);
        }
      }
    }

    // The enum-declaration allows the more liberal type-constructor in the grammar, but the
    // compiler limits this to signed or unsigned integer types
    if (parentType == Types.ENUM_DECLARATION && thisType == Types.TYPE_CONSTRUCTOR) {
      // compound-identifier -> identifier-token -> int
      ASTNode compoundIdentifier = element.getNode().findChildByType(Types.COMPOUND_IDENTIFIER);
      if (compoundIdentifier != null) {
        ASTNode identifierToken = compoundIdentifier.findChildByType(Types.IDENTIFIER_TOKEN);

        if (identifierToken != null && identifierToken.findChildByType(INTEGRAL_TYPES) == null) {
          holder.createErrorAnnotation(
              element, "Expected integral type, found " + identifierToken.getText());
        }
      }
    }

    if (thisType == Types.XUNION) {
      holder.createWarningAnnotation(
          element,
          "Xunions are transitional, and will be renamed to unions"
              + " in a future language revision");
    }

    // The bits-declaration allows the more liberal type-constructor in the grammar, but the
    // compiler limits this to unsigned integer types
    if (parentType == Types.BITS_DECLARATION && thisType == Types.TYPE_CONSTRUCTOR) {
      // type-constructor -> compound-identifier -> identifier-token -> int8-or-whatever
      ASTNode compoundIdentifier = element.getNode().findChildByType(Types.COMPOUND_IDENTIFIER);
      if (compoundIdentifier != null) {
        ASTNode identifierToken = compoundIdentifier.findChildByType(Types.IDENTIFIER_TOKEN);
        if (identifierToken != null
            && identifierToken.findChildByType(UNSIGNED_INTEGRAL_TYPES) == null) {
          holder.createErrorAnnotation(
              element, "Expected unsigned integral type, found " + identifierToken.getText());
        }
      }
    }

    // The rule from the grammar:
    // -----
    // The bits-or-enum-member-value allows the more liberal literal in the grammar, but the
    // compiler limits this to:
    //
    // - A NUMERIC-LITERAL in the context of an enum;
    // - A NUMERIC-LITERAL which must be a power of two, in the context of a bits.
    // -----
    // We take it a bit farther (as does the compiler): we ensure that the types match the size of
    // the bits or enums (they are big enough to fit into uint8 and so on).
    if (thisType == Types.BITS_OR_ENUM_MEMBER_VALUE) {
      ASTNode grandParent = element.getParent().getParent().getNode();
      IElementType grandParentType = grandParent.getElementType();
      if (grandParentType == Types.ENUM_DECLARATION || grandParentType == Types.BITS_DECLARATION) {
        // literal -> numeric-literal -> integral-literal -> {decimal, hex, binary}-integral-literal
        ASTNode literal = element.getNode().findChildByType(Types.LITERAL);
        if (literal != null) {
          ASTNode numericLiteral = literal.findChildByType(Types.NUMERIC_LITERAL);
          if (numericLiteral == null) {
            // Must be a numeric literal in either case:
            holder.createErrorAnnotation(
                element, "Expected integer value for member, found " + literal.getText());
          } else {
            ASTNode integralLiteral = numericLiteral.findChildByType(Types.INTEGRAL_LITERAL);
            if (integralLiteral == null) {
              // Must be an integral literal in either case:
              holder.createErrorAnnotation(
                  element, "Expected integer value for member, found " + literal.getText());
            } else {
              // bits/enum-declaration -> bits-or-enum-member -> bits-or-enum-member-value
              String value =
                  correctTypeOrError(
                      grandParent.findChildByType(Types.TYPE_CONSTRUCTOR), integralLiteral);
              if (value != null) {
                holder.createErrorAnnotation(element, value);
              }
              if (grandParentType == Types.BITS_DECLARATION) {
                value = unsignedLongPowerOfTwoOrError(integralLiteral);
                if (value != null) {
                  holder.createErrorAnnotation(element, value);
                }
              }
            }
          }
        }
      }
    }

    // Attributes cannot be placed on a reserved member.
    if (thisType == Types.RESERVED) {
      // Parent is TABLE_FIELD_DECL, Grandparent is ORDINAL_MEMBER_FIELD
      ASTNode grandParent = element.getParent().getParent().getNode();
      if (grandParent.getElementType() == Types.ORDINAL_MEMBER_FIELD) {
        ASTNode maybeAttributeList = grandParent.findChildByType(Types.ATTRIBUTE_LIST);
        if (maybeAttributeList != null) {
          holder.createErrorAnnotation(
              maybeAttributeList, "Attributes are not allowed on reserved table members");
        }
      }
    }

    // TODO: A service-member allows the more liberal type-constructor in the grammar, but the
    // compiler limits this to protocols.  This requires us to track legal protocol names.
  }
}