# Constructor -------------------------------------------------------------

#' Guide primitive: brackets
#'
#' This function constructs a brackets [guide primitive][guide-primitives].
#'
#' @param key A [range key][key_range] specification. See more information
#'   in the linked topic.
#' @param bracket A [bracket][bracket_options] by providing one of the
#'  following:
#'  * A bracket `<function>`, such as `bracket_square`.
#'  * A `<character[1]>` naming a bracket function without the
#'  '`bracket_`'-prefix, e.g. `"square"`.
#'  * A two-column `<matrix[n, 2]>` giving line coordinates for a bracket,
#'  like those created by bracket functions, such as `bracket_round()`.
#' @param oob A method for dealing with out-of-bounds (oob) ranges. Can be one
#'  of `"squish"`, `"censor"` or `"none"`.
#' @param drop_zero A `<logical[1]>` whether to drop near-zero width ranges
#'   (`TRUE`, default) or preserve them (`FALSE`).
#' @param pad_discrete A `<numeric[1]>` giving the amount ranges should be
#'   extended when given as a discrete variable. This is applied after
#'   the `drop_zero` setting.
#' @param levels_brackets A list of `<element_line>` objects to customise how
#'   brackets appear at every level.
#' @param levels_text A list of `<element_text>` objects to customise how
#'   text appears at every level.
#' @inheritParams common_parameters
#'
#' @return A `<PrimitiveBracket>` primitive guide that can be used inside other
#'   guides.
#' @family primitives
#' @export
#'
#' @details
#' ## Styling options
#'
#' Below are the [theme][ggplot2::theme] options that determine the styling of
#' this guide, which may differ depending on whether the guide is used in
#' an axis or in a legend context.
#'
#' The possible `{position}` suffixes mentioned below are `x`, `x.top`,
#' `x.bottom`, `y`, `y.left`, `y.right`. The `theta` and `r` position suffixes
#' in \pkg{ggplot2} are *not* obeyed in \pkg{legendry}.
#'
#' | **Theme setting** | **Context** | **Type** | **Description** |
#' | ----------------- | ----------- | -------- | --------------- |
#' | `legendry.bracket` | Both | [`element_line()`] | The bracket lines themselves. |
#' | `legendry.bracket.size` | Both | [`unit()`] | The space (in the orthogonal direction) afforded to a bracket. |
#' | `axis.text.{position}` | Axis | [`element_text()`] | The text over brackets. |
#' | `legend.text` | Legend | [`element_text()`] | The text over brackets. |
#'
#' Styling options *per level* can be set in the `levels_brackets` and
#' `levels_text` arguments. These override theme settings.
#'
#' Styling options *per range* can be set in the [range key][key_range].
#' The `line` and `text` prefixed properties are prioritised for the brackets
#' and text respectively. These override theme settings and 'per level'
#' settings.
#'
#' The context-agnostic alternative to using `theme()` is to use
#' [`theme_guide()`]:
#'
#' ```r
#' primitive_bracket(theme = theme_guide(
#'   bracket = element_line(),
#'   bracket.size = unit(5, "mm")
#' ))
#' ```
#'
#' @examples
#' # A standard plot
#' p <- ggplot(mpg, aes(interaction(drv, year), displ)) +
#'  geom_point()
#'
#' key <- key_range_manual(c(2, 4), c(5, 6), c("A", "B"))
#'
#' # Adding as secondary guides
#' p + guides(
#'   x.sec = primitive_bracket(),
#'   y.sec = primitive_bracket(key = key)
#' )
primitive_bracket <- function(
  key = "range_auto",
  bracket = "line",
  angle = waiver(),
  oob = "squish",
  drop_zero = TRUE,
  pad_discrete = 0.4,
  levels_brackets = NULL,
  levels_text = NULL,
  theme = NULL,
  position = waiver()
) {

  key <- resolve_key(key)
  oob <- arg_match0(oob, c("squish", "censor", "none"))
  check_bool(drop_zero)
  check_number_decimal(pad_discrete, allow_infinite = FALSE)
  check_list_of(
    levels_brackets, element_classes("line", "blank"),
    allow_null = TRUE
  )
  check_list_of(
    levels_text, element_classes("text", "blank"),
    allow_null = TRUE
  )
  bracket <- resolve_bracket(bracket)

  new_guide(
    key = key,
    oob = oob,
    angle = angle,
    drop_zero = drop_zero,
    pad_discrete = pad_discrete,
    bracket = bracket,
    levels_brackets = levels_brackets,
    levels_text = levels_text,
    theme = theme,
    position = position,
    available_aes = c("any", "x", "y", "r", "theta"),
    super = PrimitiveBracket
  )
}

# Class -------------------------------------------------------------------

#' @export
#' @rdname legendry_extensions
#' @format NULL
#' @usage NULL
PrimitiveBracket <- ggproto(
  "PrimitiveBracket", Guide,

  params = new_params(
    key = NULL, oob = "squish", drop_zero = TRUE,
    pad_discrete = 0.4, angle = waiver(), bracket = cbind(c(0.0, 1.0), 0.5),
    levels_text = NULL, levels_brackets = NULL
  ),

  hashables = exprs(key, decor, bracket),

  elements = list(
    position = list(
      text = "axis.text", line = I("legendry.bracket"),
      size = I("legendry.bracket.size")
    ),
    legend = list(
      text = "legend.text", line = I("legendry.bracket"),
      size = I("legendry.bracket.size")
    )
  ),

  extract_key = range_extract_key,

  extract_params = extract_range_params,

  extract_decor = function(scale, aesthetic, position, key, bracket, ...) {
    key <- vec_slice(key, key$.draw)
    n_keys <- nrow(key)
    if (n_keys < 1) {
      return(NULL)
    }
    bracket  <- resolve_bracket(bracket)
    n_vertex <- nrow(bracket)
    decor <- vec_rep_each(key, n_vertex)
    brackets <- vec_rep(bracket, n_keys)

    decor[[aesthetic]] <-
      brackets[, 1L] * (decor$end - decor$start) + decor$start
    decor[["offset"]] <- brackets[, 2L]
    decor[["group"]] <- rep(seq_len(n_keys), each = n_vertex)
    decor[setdiff(names(decor), c("start", "end"))]
  },

  transform = function(self, params, coord, panel_params) {
    params$key <-
      transform_key(params$key, params$position, coord, panel_params)
    params$bbox  <- panel_params$bbox %||%
      list(x = c(0.0, 1.0), y = c(0.0, 1.0))
    params$decor <-
      transform_bracket(params$decor, params$position, coord, panel_params)
    params
  },

  setup_params = setup_range_params,

  setup_elements = primitive_setup_elements,

  override_elements = function(params, elements, theme) {
    elements$size <- cm(elements$size)
    elements
  },

  build_bracket = function(key, decor, elements, params) {
    levels   <- unique(c(key$.level, decor$.level))
    nlevels  <- length(levels)
    position <- params$position

    # Recycle custom elements per level to appropriate length
    bracket_levels <- rep0(params$levels_brackets, length.out = nlevels)
    text_levels    <- rep0(params$levels_text,     length.out = nlevels)

    # Justify labels along their ranges
    key <- justify_ranges(key, levels, elements$text, text_levels)

    if (is_theta(position)) {
      add <- if (position == "theta.sec") pi else 0.0
      key <- polar_xy(key, key$r, key$theta + add, params$bbox)
    }

    if (is_blank(elements$line) || is_empty(decor)) {
      decor <- vec_slice(decor, 0L)
    } else if (position %in% .trbl) {
      offset  <- decor$offset
      offset  <- if (position %in% .trbl[1L:2L]) 1.0 - offset else offset
      decor$x <- switch(position, left = , right = offset, decor$x)
      decor$y <- switch(position, top = , bottom = offset, decor$y)
      decor$offset <- 0.0
    }

    offset <- elements$offset
    angle  <- params$angle %|W|% NULL
    size   <- elements$size
    text   <- angle_labels(elements$text, angle, params$position)

    brackets <- vector("list", nlevels)
    labels   <- vector("list", nlevels)

    for (i in seq_len(nlevels)) {

      # Render bracket
      brackets[[i]] <- draw_bracket(
        decor   = vec_slice(decor, decor$.level == levels[[i]]),
        element = combine_elements(bracket_levels[[i]], elements$line),
        size = size, offset = offset, position = position
      )
      offset <- offset + get_size_attr(brackets[[i]])

      # Render text
      labels[[i]] <- draw_labels(
        key     = vec_slice(key, key$.level == levels[[i]]),
        element = combine_elements(text_levels[[i]], text),
        angle = angle, offset = offset, position = position
      )
      offset <- offset + get_size_attr(labels[[i]])
    }
    if (params$position %in% c("top", "left")) {
      brackets <- rev(brackets)
      labels   <- rev(labels)
    }

    list(brackets = brackets, labels = labels)
  },

  measure_grobs = function(grobs, params, elements) {
    labels <- switch(
      params$position,
      top  = , bottom = height_cm(grobs$labels),
      left = , right  =  width_cm(grobs$labels),
      map_dbl(grobs$labels, get_size_attr)
    )
    is_bracket <- as.numeric(!map_lgl(grobs$brackets, is_zero))
    bracket <- is_bracket * elements$size
    list(brackets = bracket, labels = labels)
  },

  draw = function(self, theme, position = NULL, direction = NULL,
                  params = self$params) {

    params <- replace_null(params, position = position, direction = direction)
    params <- self$setup_params(params)

    elems    <- self$setup_elements(params, self$elements, theme)
    elems    <- self$override_elements(params, elems, theme)
    brackets <- self$build_bracket(params$key, params$decor, elems, params)
    size     <- self$measure_grobs(brackets, params, elems)

    if (params$position %in% c("top", "left")) {
      grobs <- vec_interleave(brackets$labels, brackets$brackets)
      size  <- vec_interleave(size$labels, size$brackets)
    } else {
      grobs <- vec_interleave(brackets$brackets, brackets$labels)
      size  <- vec_interleave(size$brackets, size$labels)
    }

    primitive_grob(
      grob = grobs,
      size = unit(size, "cm"),
      position = params$position,
      name = "bracket"
    )
  }
)

# Helpers -----------------------------------------------------------------

draw_bracket <- function(decor, element, size, offset, position) {
  if (nrow(decor) < 2L) {
    return(zeroGrob())
  }
  x <- unit(decor$x, "npc")
  y <- unit(decor$y, "npc")

  if (is_theta(position)) {
    offset <- (1.0 - decor$offset) * size + offset
    x <- x + unit(sin(decor$theta) * offset, "cm")
    y <- y + unit(cos(decor$theta) * offset, "cm")
  }

  id <- new_rle(decor$group)
  props <- element_key_properties(vec_slice(decor, id$start), "line")

  grob <- inject(element_grob(
    element, x = x, y = y,
    id.lengths = id$times,
    !!!props
  ))
  if (!is_blank(element)) {
    attr(grob, "size") <- size
  }
  grob
}
