///|
struct CompletionItem {
  label : String
  description : String
}

pub fn CompletionItem::new(label : String, description : String) -> CompletionItem {
  { label, description }
}

///|
struct PromptServices {
  completer_hook : (StringView) -> Array[CompletionItem]
}

///|
pub fn PromptServices::new(
  completer_hook? : (StringView) -> Array[CompletionItem] = _ => [],
) -> PromptServices {
  { completer_hook, }
}

///|
priv enum TokenCat {
  Keyword
  CtrlKeyword
  Constant
  UIdent
  LIdent
  Operator
  PackageName
  Others
  Comment
  Attribute
}

///|
fn TokenCat::from_token(tok : @tokens.Token) -> TokenCat {
  match tok {
    LEXMATCH_QUESTION => CtrlKeyword
    LEXMATCH => CtrlKeyword
    TRY_EXCLAMATION => CtrlKeyword
    TRY_QUESTION => CtrlKeyword
    WHERE => CtrlKeyword
    NOBREAK => CtrlKeyword
    NORAISE => CtrlKeyword
    DECLARE => Keyword
    ENUMVIEW => Keyword
    LETREC => Keyword
    AND => Keyword
    SUBERROR => Keyword
    IS => CtrlKeyword
    IN => CtrlKeyword
    FOR => CtrlKeyword
    DEFER => CtrlKeyword
    GUARD => CtrlKeyword
    LOOP => CtrlKeyword
    TEST => Keyword
    ELLIPSIS => Operator
    RANGE_EXCLUSIVE_REV => Operator
    RANGE_INCLUSIVE_REV => Operator
    RANGE_EXCLUSIVE => Operator
    RANGE_LT_INCLUSIVE => Operator
    RANGE_INCLUSIVE => Operator
    DOTDOT => Operator
    RETURN => CtrlKeyword
    WHILE => CtrlKeyword
    THIN_ARROW => Operator
    FAT_ARROW => Operator
    TYPE => CtrlKeyword
    MUTABLE => Keyword
    USING => CtrlKeyword
    MATCH => CtrlKeyword
    CONST => Keyword
    LET => Keyword
    IF => CtrlKeyword
    FN => Keyword
    ELSE => CtrlKeyword
    PIPE => Operator
    AS => CtrlKeyword
    PACKAGE_NAME(_) => PackageName
    BARBAR => Operator
    CARET => Operator
    AMPER => Operator
    AMPERAMPER => Operator
    RBRACE => Operator
    LBRACE => Operator
    BAR => Operator
    UNDERSCORE => LIdent
    RBRACKET => Operator
    PLUS => Operator
    LBRACKET => Operator
    SEMI(_) => Operator
    COLON => Operator
    COLONCOLON => Operator
    DOT_LPAREN => Operator
    DOT_INT(_) => Constant
    DOT_UIDENT(_) => UIdent
    DOT_LIDENT(_) => LIdent
    EXCLAMATION => Operator
    QUESTION => Operator
    MINUS => Operator
    COMMA => Operator
    RPAREN => Operator
    LPAREN => Operator
    EQUAL => Operator
    FNALIAS => Keyword
    TRAITALIAS => Keyword
    TYPEALIAS => Keyword
    ASYNC => CtrlKeyword
    CATCH => CtrlKeyword
    TRY => CtrlKeyword
    THROW => CtrlKeyword
    RAISE => CtrlKeyword
    WITH => Keyword
    IMPL => Keyword
    DERIVE => Keyword
    TRAIT => Keyword
    ENUM => Keyword
    STRUCT => Keyword
    CONTINUE => CtrlKeyword
    BREAK => CtrlKeyword
    EXTERN => Keyword
    IMPORT => Keyword
    READONLY => Keyword
    PRIV => Keyword
    PUB => Keyword
    TRUE => Keyword
    FALSE => Keyword
    EOF => Others
    AUGMENTED_ASSIGNMENT(_) => Operator
    INFIX4(_) => Operator
    INFIX3(_) => Operator
    INFIX2(_) => Operator
    INFIX1(_) => Operator
    NEWLINE => Others
    COMMENT(_) => Comment
    POST_LABEL(_) => LIdent
    UIDENT(_) => UIdent
    LIDENT(_) => LIdent
    ATTRIBUTE(_) => Attribute
    REGEX_INTERP(_) => Constant
    REGEX_LITERAL(_) => Constant
    INTERP(_) => Constant
    MULTILINE_INTERP(_) => Constant
    MULTILINE_STRING(_) => Constant
    STRING(_) => Constant
    DOUBLE(_) => Constant
    FLOAT(_) => Constant
    BYTES(_) => Constant
    BYTE(_) => Constant
    INT(_) => Constant
    CHAR(_) => Constant
    PACKAGE => PackageName
  }
}

///|
fn TokenCat::format(self : TokenCat) -> @pp.ConsoleFormat? {
  fn foreground(color) {
    Some(@pp.ConsoleFormat::new(foreground=color))
  }
  match self {
    Comment => foreground(Green)
    Others => None
    PackageName => foreground(BrightGreen)
    Operator => None
    LIdent => foreground(White)
    UIdent => foreground(BrightGreen)
    Constant => foreground(Yellow)
    CtrlKeyword => foreground(BrightMagenta)
    Keyword => foreground(Blue)
    Attribute => foreground(Magenta)
  }
}

///|
fn[A, B] scan(iter : Iter[A], init : B, f : (A, B) -> B) -> Iter[B] {
  let mut acc = init
  Iter::new(() => {
    iter
    .next()
    .map(x => {
      acc = f(x, acc)
      acc
    })
  })
}

///|
priv struct TokenInfo {
  start : Int
  length : Int
  token : @tokens.Token
}

///|
fn tokenize(text : String, name : String) -> Array[TokenInfo] {
  let lines = text.split("\n").map(line => line.trim_end(chars="\r"))
  let lengths = [0, ..scan(lines, 0, (line, acc) => line.length() + acc)]
  let lex_result = @lexer.tokens_from_string(text, comment=true, name~)
  lex_result.tokens.filter_map(tok => {
    if tok.0 is EOF {
      return None
    }
    let start = lengths[tok.1.lnum - 1] + tok.1.column() - 1
    let end = lengths[tok.2.lnum - 1] + tok.2.column() - 1
    Some({ start, length: end - start, token: tok.0 })
  })
}

///|
pub fn highlighter(snapshot : @pp.PromptSnapshot) -> Array[@pp.FormatSpan] {
  let tokens = tokenize(snapshot.text, "stdin")
  let brackets_map = {}
  let parens = []
  let brackets = []
  let braces = []
  for tok in tokens {
    match tok.token {
      LPAREN => parens.push(tok.start)
      LBRACE => braces.push(tok.start)
      LBRACKET => brackets.push(tok.start)
      RPAREN =>
        if parens.pop() is Some(open) {
          brackets_map[open] = tok.start
          brackets_map[tok.start] = open
        }
      RBRACKET =>
        if brackets.pop() is Some(open) {
          brackets_map[open] = tok.start
          brackets_map[tok.start] = open
        }
      RBRACE =>
        if braces.pop() is Some(open) {
          brackets_map[open] = tok.start
          brackets_map[tok.start] = open
        }
      _ => ()
    }
  }
  let caret = if snapshot.caret == snapshot.text.length() {
    snapshot.caret - 1
  } else {
    snapshot.caret
  }
  let brackets_at = if brackets_map.get(caret) is Some(other) {
    [caret, other]
  } else {
    []
  }
  tokens.filter_map(tok => {
    if brackets_at.contains(tok.start) {
      Some(
        @pp.FormatSpan::new(
          tok.start,
          tok.length,
          @pp.ConsoleFormat::new(underline=true),
        ),
      )
    } else {
      TokenCat::from_token(tok.token)
      .format()
      .map(format => @pp.FormatSpan::new(tok.start, tok.length, format))
    }
  })
}

///|
let moonbit_keywords : ReadOnlyArray[String] = [
  "as", "else", "extern", "fn", "fnalias", "if", "let", "const", "match", "using",
  "mut", "type", "typealias", "struct", "enum", "trait", "traitalias", "derive",
  "while", "break", "continue", "import", "return", "throw", "raise", "try", "catch",
  "pub", "priv", "readonly", "true", "false", "_", "test", "loop", "for", "in", "impl",
  "with", "guard", "async", "is", "suberror", "and", "letrec", "enumview", "noraise",
  "defer",
]

///|
pub fn PromptServices::completer(
  self : PromptServices,
  req : @pp.CompletionRequest,
) -> Array[@pp.CompletionItem] {
  let s = req.snapshot.text[:req.snapshot.caret]
  let result = (self.completer_hook)(s).map(item => {
    @pp.CompletionItem::new(item.label, description=item.description)
  })
  for kwd in moonbit_keywords {
    if kwd.has_prefix(s) {
      result.push(@pp.CompletionItem::new(kwd, description="keyword"))
    }
  }
  result
}

///|
pub fn brackets_match(snapshot : @pp.PromptSnapshot) -> Bool {
  let tokens = @lexer.tokens_from_string(
      snapshot.text,
      comment=true,
      name="stdin",
    ).tokens
  let stack : Array[@tokens.TokenKind] = []
  for tok in tokens {
    let kind = tok.0.kind()
    match kind {
      TK_LBRACE => stack.push(TK_RBRACE)
      TK_LBRACKET => stack.push(TK_RBRACKET)
      TK_LPAREN => stack.push(TK_RPAREN)
      TK_RBRACE | TK_RBRACKET | TK_RPAREN =>
        if stack.pop() != Some(kind) {
          return false
        }
      _ => ()
    }
  }
  stack.is_empty()
}

///|
pub fn PromptServices::as_callbacks(
  self : PromptServices,
) -> @pp.PromptCallbacks {
  @pp.PromptCallbacks::new(
    highlight_callback=snapshot => highlighter(snapshot),
    completion_items_provider=request => self.completer(request),
    should_insert_soft_newline=(snapshot, _) => !brackets_match(snapshot),
  )
}