
operator_tokens <- function() {
  c(
    "'-'", "'+'", "'!'", "'~'", "'?'", "':'", "'*'", "'/'", "'^'",
    "SPECIAL", "LT", "GT", "EQ", "GE", "LE", "AND", "AND2", "OR", "OR2",
    "LEFT_ASSIGN", "RIGHT_ASSIGN", "'$'", "'@'", "EQ_ASSIGN", "PIPE"
  )
}

reserved_words <- function() {
  c("FUNCTION", "'\\\\'", "IF", "ELSE",
    "REPEAT", "WHILE", "FOR", "IN", "NEXT", "BREAK")
}


#' Syntax highlight R code
#'
#' @details
#' See [code_theme_list()] for the default syntax highlighting theme and
#' how to change it.
#'
#' If `code` does not parse, then it is returned unchanged and a
#' `cli_parse_failure` condition is thrown. Note that this is not an error,
#' and the condition is ignored, unless explicitly caught.
#'
#' @param code Character vector, each element is one line of code.
#' @param code_theme Theme see [code_theme_list()].
#' @return Character vector, the highlighted code.
#'
#' @family syntax highlighting
#' @importFrom utils getSrcref getParseData
#' @export
#' @examples
#' code_highlight(deparse(ls))
#' cat(code_highlight(deparse(ls)), sep = "\n")

code_highlight <- function(code, code_theme = NULL) {

  code_theme <- code_theme %||% code_theme_default()

  parsed <- tryCatch(
    parse(text = code, keep.source = TRUE),
    error = function(e) e
  )

  if (inherits(parsed, "error")) {
    cnd <- structure(
      list(message = conditionMessage(parsed), code = code),
      class = c("cli_parse_failure", "condition")
    )
    signalCondition(cnd)
    return(code)
  }

  theme <- code_theme_make(code_theme)
  data <- getParseData(parsed, includeText = NA)

  hitext <- data$text

  cnv <- function(x) do.call(combine_ansi_styles, as.list(x))

  brackettheme <- lapply(theme$bracket, cnv)
  theme <- theme[names(theme) != "bracket"]
  theme <- structure(lapply(theme, cnv), names = names(theme))

  ## Reserved words if else repeat while function for in next break
  if (!is.null(theme$reserved)) {
    reserved <- data$token %in% reserved_words()
    hitext[reserved] <- theme$reserved(data$text[reserved])
  }

  ## Numeric constants, including NAs, NaN and Inf
  if (!is.null(theme$number)) {
    num_const <- data$token == "NUM_CONST"
    hitext[num_const] <- theme$number(data$text[num_const])
  }

  ## NULL
  if (!is.null(theme$null)) {
    null <- data$token == "NULL_CONST"
    hitext[null] <- theme$null(data$text[null])
  }

  ## Operators
  if (!is.null(theme$operator)) {
    operator <- data$token %in% operator_tokens()
    hitext[operator] <- theme$operator(data$text[operator])
  }

  ## Function calls
  if (!is.null(theme$call)) {
    fun_call <- data$token == "SYMBOL_FUNCTION_CALL"
    hitext[fun_call] <- theme$call(data$text[fun_call])
  }

  ## Strings
  if (!is.null(theme$string)) {
    string <- data$token == "STR_CONST"
    reserved <- theme$reserved %||% function(x) x
    raw <- substr(data$text[string], 1, 1) == "r"
    hitext[string][raw] <- paste0(
      rep(reserved("r"), sum(raw)),
      theme$string(substr(data$text[string][raw], 2, nchar(data$text[string][raw])))
    )
    hitext[string][!raw] <- theme$string(data$text[string][!raw])
  }

  ## Comments
  if (!is.null(theme$comment)) {
    comment <- data$token == "COMMENT"
    hitext[comment] <- theme$comment(data$text[comment])
  }

  ## Brackets
  if (length(brackettheme)) {
    bracket <- data$token %in% bracket_tokens()
    hitext[bracket] <- color_brackets(data$text[bracket], brackettheme)
  }

  do_subst(code, data, hitext)
}

do_subst <- function(code, pdata, hitext) {

  pdata$hitext <- hitext

  ## Need to do this line by line. TODO: multiline stuff might be broken
  vapply(seq_along(code), FUN.VALUE = character(1), function(no) {
    my <- pdata[pdata$line1 == no & pdata$line2 == no,, drop = FALSE]
    replace_in_place(code[no], my$col1, my$col2, my$hitext)
  })
}

open_brackets <- function() {
  c("(", "{", "[")
}

close_brackets <- function(){
  c(")", "}", "]")
}

bracket_tokens <- function() {
  s <- c(open_brackets(), close_brackets())
  c(paste0("'", s, "'"), "LBB")
}

apply_color <- function(x, lvl, l){
  k <- (lvl - 1) %% length(l) + 1
  l[[k]](x)
}

#' Colored brackets
#'
#' Add color to brackets. Brackets will be coloured consecutively with the
#' colors provided in \code{color_seq} by scope.
#'
#' @param x a character vector of brackets consisting of a valid sequence of any
#'   of the following: \code{'[[', '[', ']', '(', ')', '{', '}'}
#' @param color_seq a list of functions that take and return a character scalar. The
#' ordering defines the sequence of color functions to apply to a given scope level.
#' Color functions are recycled when the scope level exceeds the length of \code{color_seq}
#'
#' @details Meant for coloring brackets encountered within \code{highlight}.
#'   Note that occurrences of 'orphan' brackets are not taken into account
#'   mainly due to the fact that cases such as
#'
#'   \code{foo <- function(x){ `[[`(x, 1) }}
#'
#'   will either be converted to
#'
#'   \code{foo <- function(x){ x[[1]] }}
#'
#'   before the brackets are coloured if passed in as
#'   \code{highlight(deparse(foo))} or will be identified as a
#'   'SYMBOL_FUNCTION_CALL' token instead of 'LBB' if passed in as
#'
#'   \code{highlight("foo <- function(x){ `[[`(x, 1) }")}
#'
#'   Similarly, invalid code that would lead to orphaned brackets is not taken
#'   into account as this would be caught before hand in \code{highlight}.
#'
#' @noRd

color_brackets <- function(x, color_seq = list(col_yellow, col_blue, col_cyan)) {
  stopifnot(vapply(color_seq, is.function, logical(1)))
  open <- c(open_brackets(), "[[")
  o <- character()
  lvl <- 0
  i <- 1
  while (i <= length(x)) {

    if (x[i] %in% open) {
      o[length(o) + 1] <- x[i]
      lvl <- lvl + 1
      x[i] <- apply_color(x[i], lvl, color_seq)
      i <- i + 1
      next
    }

    j <- nchar(o[length(o)])
    x[i:(i + j - 1)] <-
      apply_color(x[i:(i + j - 1)], lvl, color_seq)
    i <- i + j
    lvl <- lvl - 1
    o <- o[-length(o)]
  }
  x
}

replace_in_place <- function(str, start, end, replacement) {

  stopifnot(
    length(str) == 1,
    length(start) == length(end),
    length(end) == length(replacement)
  )

  keep <- substring(str, c(1, end + 1), c(start - 1, nchar(str)))

  pieces <- character(length(replacement) * 2 + 1)

  even <- seq_along(replacement) * 2
  odd <- c(1, even + 1)
  pieces[even] <- replacement
  pieces[odd] <- keep

  paste0(pieces, collapse = "")
}

code_theme_default <- function() {
  opt <- code_theme_opt("cli.code_theme")
  if (!is.null(opt)) return(opt)

  rs <- rstudio_detect()
  if (rs$type %in% c("rstudio_console", "rstudio_console_starting")) {
    opt <- code_theme_opt("cli.code_theme_rstudio")
    if (!is.null(opt)) return(opt)
    code_theme_default_rstudio()
  } else {
    opt <- code_theme_opt("cli.code_theme_terminal")
    if (!is.null(opt)) return(opt)
    code_theme_default_term()
  }
}

code_theme_opt <- function(option) {
  theme <- getOption(option)
  if (is.null(theme)) return(NULL)

  code_theme_make(theme)
}

code_theme_make <- function(theme) {
  if (is.list(theme)) return(theme)
  if (is_string(theme)) {
    if (theme %in% names(rstudio_themes)) return(rstudio_themes[[theme]])
    lcs <- gsub(" ", "_", tolower(names(rstudio_themes)))
    if (theme %in% lcs) return(rstudio_themes[[ match(theme, lcs)[1] ]])
    warning("Unknown cli code theme: `", theme, "`.")
    return(NULL)
  }
  warning("Invalid cli code theme, see documentation")
  NULL
}

code_theme_default_rstudio <- function() {
  theme <- rstudioapi::getThemeInfo()$editor
  if (! theme %in% names(rstudio_themes)) {
    warning("cli does not know this RStudio theme: ", theme)
    return(list())
  }
  rstudio_themes[[theme]]
}

code_theme_default_term <- function() {
  list(
    reserved = "red",
    number   = "blue",
    null     = c("blue", "bold"),
    operator = "green",
    call     = "cyan",
    string   = "yellow",
    comment  = c("#a9a9a9", "italic"),
    bracket  = list("yellow", "blue", "cyan")
  )
}

#' Syntax highlighting themes
#'
#' @description
#' `code_theme_list()` lists the built-in code themes.
#'
#' # Code themes
#' A theme is a list of character vectors, except for `bracket`, see below.
#' Each character vector must contain RGB colors (e.g. `"#a9a9a9"`),
#' and cli styles, e.g. `"bold"`. Entries in the list:
#' * `reserved`: reserved words
#' * `number`: numeric literals
#' * `null`: the `NULL` constant
#' * `operator`: operators, including assignment
#' * `call`: function calls
#' * `string`: character literals
#' * `comment`: comments
#' * `bracket`: brackets: \code{(){}[]} This is a list of character vectors,
#'   to create "rainbow" brackets. It is recycled for deeply nested lists.
#'
#' # The default code theme
#'
#' In RStudio, it matches the current theme of the IDE.
#'
#' You can use three options to customize the code theme:
#' * If `cli.code_theme` is set, it is used.
#' * Otherwise if R is running in RStudio and `cli.code_theme_rstudio` is
#'   set, then it is used.
#' * Otherwise if T is not running in RStudio and `cli.code_theme_terminal`
#'   is set, then it is used.
#'
#' You can set these options to the name of a built-in theme, or to list
#' that specifies a custom theme. See [code_theme_list()] for the list
#' of the built-in themes.
#'
#' @return Character vector of the built-in code theme names.
#'
#' @family syntax highlighting
#' @export
#' @examples
#' code_theme_list()
#' code_highlight(deparse(get), code_theme = "Solarized Dark")

code_theme_list <- function() {
  names(rstudio_themes)
}
