package fuchsia.developer.plugin.fidl;

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.openapi.util.TextRange;
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 org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class ContextAwareHighlighter implements Annotator {

  private static final TextAttributesKey KEYWORD_ATTRIBUTE;
  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];

    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;
  }

  /**
   * Determines whether a sequence of characters is a legal escape sequence. The offset passed is
   * the character *after* the \.
   *
   * @return the length of the escape if true, -1 if unknown escape sequence, -2 if malformed.
   */
  private static int isEscape(char[] chars, int offset) {
    // Is it a single-character escape sequence?
    if ("nrt\\\"".indexOf(chars[offset]) >= 0) {
      return 1;
    }

    // Is it a unicode escape sequence?
    if (chars[offset] == 'u') {
      if (offset + 2 >= chars.length - 1 || chars[offset + 1] != '{') {
        return -1;
      }

      int i = offset + 2;
      StringBuilder number = new StringBuilder();
      while (Character.digit(chars[i], 16) != -1) {
        if (i - offset == 7) {
          // too long!
          return -1;
        }
        number.append(chars[i]);
        i++;
      }
      if (chars[i] != '}' || number.length() == 0) {
        return -1;
      }
      int value = Integer.parseInt(number.toString(), 16);
      if (value < 0 || value >= 0x10fff) {
        return -1; // illegal unicode value
      }
      return i - offset;
    }

    // Unknown escape sequence
    return -1;
  }

  /**
   * Checks the node in question, and returns its type, or null if it's not a layout.
   *
   * @param element PsiElement of the node in question.
   * @return The IElementType describing this layout (or null if this is not a layout).
   */
  @Nullable
  private static IElementType getLayoutKind(PsiElement element) {
    if (element != null) {
      ASTNode layoutKindNode = element.getNode().findChildByType(Types.LAYOUT_KIND);
      if (layoutKindNode != null) {
        return layoutKindNode.getFirstChildNode().getElementType();
      }
    }
    return null;
  }

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

    IElementType thisType = element.getNode().getElementType();
    IElementType maybeThisLayoutKind = getLayoutKind(element);
    IElementType parentType = parent.getNode().getElementType();
    IElementType grandParentType = null;
    IElementType maybeGrandParentLayoutKind = null;
    IElementType maybeGreatGrandParentLayoutKind = null;

    PsiElement grandParent = parent.getParent();
    PsiElement greatGrandParent = null;
    if (grandParent != null && grandParent.getNode() != null) {
      grandParentType = grandParent.getNode().getElementType();
      maybeGrandParentLayoutKind = getLayoutKind(grandParent);

      greatGrandParent = grandParent.getParent();
      if (greatGrandParent != null && greatGrandParent.getNode() != null) {
        maybeGreatGrandParentLayoutKind = getLayoutKind(greatGrandParent);
      }
    }

    // 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) {
      if (grandParentType != null && grandParentType == Types.TYPE_CONSTRUCTOR) {
        if (FidlLexer.ALL_KEYWORDS.contains(element.getText())) {
          Annotation annotation = holder.createInfoAnnotation(element, null);
          annotation.setTextAttributes(KEYWORD_ATTRIBUTE);
        }
      }
    }

    // The enum layout allows the more liberal type-constructor in the grammar, but the compiler
    // limits this to signed or unsigned integer types.  The bits layout is even more restrictive,
    // and only allows unsigned integer types.
    if (thisType == Types.TYPE_CONSTRUCTOR
        && (maybeGrandParentLayoutKind == Types.BITS || maybeGrandParentLayoutKind == Types.ENUM)) {
      // LAYOUT -> COMPOUND_IDENTIFIER-> IDENTIFIER_TOKEN -> INT
      ASTNode layout = element.getNode().findChildByType(Types.LAYOUT);
      if (layout != null) {
        ASTNode compoundIdentifier = layout.findChildByType(Types.COMPOUND_IDENTIFIER);
        if (compoundIdentifier != null) {
          ASTNode identifierToken = compoundIdentifier.findChildByType(Types.IDENTIFIER_TOKEN);
          if (maybeGrandParentLayoutKind == Types.BITS) {
            // Only unsigned integers allowed for bits.
            if (identifierToken != null
                && identifierToken.findChildByType(UNSIGNED_INTEGRAL_TYPES) == null) {
              holder.createErrorAnnotation(
                  element, "Expected integral type, found " + identifierToken.getText());
            }
          } else {
            // Only integers allowed for enums.
            if (identifierToken != null
                && identifierToken.findChildByType(INTEGRAL_TYPES) == null) {
              holder.createErrorAnnotation(
                  element, "Expected integral type, found " + identifierToken.getText());
            }
          }
        }
      }
    }

    // The rule from the grammar:
    // -----
    // The VALUE_LAYOUT_MEMBER 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.VALUE_LAYOUT_MEMBER) {
      // INLINE_LAYOUT -> LAYOUT_BODY -> VALUE_LAYOUT -> VALUE_LAYOUT_MEMBER
      if (maybeGreatGrandParentLayoutKind == Types.BITS
          || maybeGreatGrandParentLayoutKind == Types.ENUM) {
        // 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 {
              // INLINE_LAYOUT -> LAYOUT_BODY -> VALUE_LAYOUT -> VALUE_LAYOUT_MEMBER
              ASTNode subtype = greatGrandParent.getNode().findChildByType(Types.LAYOUT_SUBTYPE);
              if (subtype != null) {
                String value =
                    correctTypeOrError(
                        subtype.findChildByType(Types.TYPE_CONSTRUCTOR), integralLiteral);
                if (value != null) {
                  holder.createErrorAnnotation(element, value);
                }
              }
              if (maybeGreatGrandParentLayoutKind == Types.BITS) {
                String value = unsignedLongPowerOfTwoOrError(integralLiteral);
                if (value != null) {
                  holder.createErrorAnnotation(element, value);
                }
              }
            }
          }
        }
      }
    }

    // Attributes cannot be placed on a reserved member.
    if (thisType == Types.RESERVED) {
      // Parent is ORDINAL_LAYOUT_MEMBER
      if (parent.getNode().getElementType() == Types.ORDINAL_LAYOUT_MEMBER) {
        ASTNode maybeAttributeList = parent.getNode().findChildByType(Types.ATTRIBUTE_BLOCK);
        if (maybeAttributeList != null && !maybeAttributeList.getText().equals("")) {
          holder.createErrorAnnotation(
              maybeAttributeList, "Attributes are not allowed on reserved table members");
        }
      }
    }

    // The grammar allows `( declaration-modifiers )*` on all declarations, but the
    // compiler limits this as follows:
    // * A modifier cannot occur twice on the same declaration.
    //    * The `flexible` and `strict` modifiers cannot be used together.
    //    * The `flexible` and `strict` modifiers can only be used on `bits`, `enum`, and `union`.
    // * The `resource` modifier can only be used on `struct`, `table`, and `union`.
    if (element.getNode().findChildByType(Types.DECLARATION_MODIFIERS) != null) {
      ASTNode flexible = null;
      ASTNode strict = null;
      ASTNode resource = null;
      for (PsiElement sibling : element.getChildren()) {
        ASTNode siblingNode = sibling.getNode();
        if (siblingNode.getElementType() == Types.DECLARATION_MODIFIERS) {
          ASTNode found;
          found = siblingNode.findChildByType(Types.FLEXIBLE);
          if (found != null) {
            if (flexible == null) {
              flexible = found;
            } else {
              holder.createErrorAnnotation(
                  found, "The `flexible` modifier cannot be used twice in the same declaration");
            }
          }
          found = siblingNode.findChildByType(Types.STRICT);
          if (found != null) {
            if (strict == null) {
              strict = found;
            } else {
              holder.createErrorAnnotation(
                  found, "The `strict` modifier cannot be used twice in the same declaration");
            }
          }
          found = siblingNode.findChildByType(Types.RESOURCE);
          if (found != null) {
            if (resource == null) {
              resource = found;
            } else {
              holder.createErrorAnnotation(
                  found, "The `resource` modifier cannot be used twice in the same declaration");
            }
          }
        }
      }

      if (flexible != null || strict != null) {
        if (maybeThisLayoutKind != Types.BITS
            && maybeThisLayoutKind != Types.ENUM
            && maybeThisLayoutKind != Types.UNION) {
          ASTNode theNode = flexible != null ? flexible : strict;
          holder.createErrorAnnotation(
              theNode,
              "The `flexible` and `strict` modifiers can only be used on bits, enum, and union.");
        }
        if (flexible != null && strict != null) {
          holder.createErrorAnnotation(
              flexible, "The flexible and strict modifiers cannot be used together");
        }
      }
      if (resource != null) {
        if (maybeThisLayoutKind != Types.STRUCT
            && maybeThisLayoutKind != Types.TABLE
            && maybeThisLayoutKind != Types.UNION) {
          holder.createErrorAnnotation(
              resource, "The `resource` modifier can only be used on struct, table, and union.");
        }
      }
    }

    // The grammar for `STRING-LITERAL` is as follows:
    //
    // ```
    // STRING-LITERAL       = "\"" ( unicode-value | byte-value )* "\"" ;
    // unicode-value        = limited-unicode-char | little-u-value |
    //                        big-u-value | escaped-char ;
    // limited-unicode-char = any unicode character except CR, LF, "\" or "\"" ;
    // byte-value           = octal-byte-value | hex-byte-value ;
    // octal-byte-value     = "\" octal-digit octal-digit octal-digit ;
    // hex-byte-value       = "\x" hex-digit hex-digit ;
    // little-u-value       = "\\u" hex-digit hex-digit hex-digit hex-digit ;
    // big-u-value          = "\U" hex-digit hex-digit hex-digit hex-digit
    //                             hex-digit hex-digit hex-digit hex-digit ;
    // escaped-char         = "\" ( "a" | "b" | "f" | "n" | "r" | "t" | "v" | "\" | "\"" ) ;
    if (thisType == Types.STRING_LITERAL) {
      String literal = element.getText();
      char[] chars = new char[literal.length()];
      literal.getChars(0, literal.length(), chars, 0);
      for (int i = 1; i < chars.length; i++) {
        if (chars[i] == '\n' || chars[i] == '\r') {
          holder.createErrorAnnotation(element, "CR and LF not allowed in string literal");
        }
        if (chars[i] == '\\') {
          if (i < chars.length - 2) {
            int jump = isEscape(chars, i + 1);
            int startOffset = element.getTextOffset() + i;
            if (jump == -1) {
              holder.createErrorAnnotation(
                  new TextRange(startOffset, startOffset + 2),
                  "Unknown escape sequence \\" + chars[i + 1]);
            } else if (jump == -2) {
              holder.createErrorAnnotation(
                  new TextRange(startOffset, startOffset + 2),
                  "Malformed \\" + chars[i + 1] + " escape sequence");
            } else {
              i += jump;
            }
          } else {
            holder.createErrorAnnotation(element, "Illegal string termination");
          }
        }
      }
    }

    // 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.
  }
}