#' Add footnote (new)
#'
#' @description `footnote` provides a more flexible way to add footnote. You
#' can add mutiple sets of footnote using differeny notation system. It is
#' also possible to specify footnote section header one by one and print
#' footnotes as a chunk of texts.
#'
#' @param kable_input HTML or LaTeX table generated by `knitr::kable`
#' @param general Text for general footnote comments. Footnotes in this section
#' won't be labeled with any notations
#' @param number A vector of footnote texts. Footnotes here will be numbered.
#' There is no upper cap for the number of footnotes here
#' @param alphabet A vector of footnote texts, Footnotes here will be labeled
#' with abc. The vector here should not have more than 26 elements.
#' @param symbol A vector of footnote texts, Footnotes here will be labeled
#' with special symbols. The vector here should not have more than 20 elements.
#' @param footnote_order The order of how to arrange `general`, `number`,
#' `alphabet` and `symbol`.
#' @param footnote_as_chunk T/F value. Default is FALSE. It controls whether
#' the footnotes should be printed in a chunk (without line break).
#' @param escape T/F value. It controls whether the contents and titles should
#' be escaped against HTML or LaTeX. Default is TRUE.
#' @param threeparttable T/F value for whether to use LaTeX package
#' threeparttable. Threeparttable will force the width of caption and
#' footnotes be the width of the original table. It's useful when you have
#' long paragraph of footnotes.
#' @param general_title Section header for general footnotes. Default is
#' "Note: ".
#' @param number_title Section header for number footnotes. Default is "".
#' @param alphabet_title Section header for alphabet footnotes. Default is "".
#' @param symbol_title Section header for symbol footnotes. Default is "".
#'
#' @examples dt <- mtcars[1:5, 1:5]
#' footnote(knitr::kable(dt, "html"), alphabet = c("Note a", "Note b"))
#'
#' @export
footnote <- function(kable_input,
                     general = NULL,
                     number = NULL,
                     alphabet = NULL,
                     symbol = NULL,
                     footnote_order = c("general", "number",
                                        "alphabet", "symbol"),
                     footnote_as_chunk = FALSE,
                     escape = TRUE,
                     threeparttable = FALSE,
                     general_title = "Note: ",
                     number_title = "",
                     alphabet_title = "",
                     symbol_title = ""
) {
  kable_format <- attr(kable_input, "format")
  if (!kable_format %in% c("html", "latex")) {
    warning("Please specify format in kable. kableExtra can customize either ",
            "HTML or LaTeX outputs. See https://haozhu233.github.io/kableExtra/ ",
            "for details.")
    return(kable_input)
  }
  if (length(alphabet) > 26) {
    alphabet <- alphabet[1:26]
    warning("Please don't use more than 26 footnotes in table_footnote ",
            "alphabet. Use number instead.")
  }
  if (length(symbol) > 20) {
    symbol <- symbol[1:20]
    warning("Please don't use more than 20 footnotes in table_footnote ",
            "symbol. Use number instead.")
  }
  footnote_titles <- list(
    general = general_title, number = number_title,
    alphabet = alphabet_title, symbol = symbol_title
  )
  footnote_contents <- list(
    general = general, number = number, alphabet = alphabet, symbol = symbol
  )
  notnull <- names(footnote_contents)[!sapply(footnote_contents, is.null)]
  if (length(notnull) == 0) {return(kable_input)}
  footnote_order <- footnote_order[footnote_order %in% notnull]
  footnote_titles <- footnote_titles[footnote_order]
  footnote_contents <- footnote_contents[footnote_order]
  if (escape) {
    if (kable_format == "html") {
      footnote_contents <- lapply(footnote_contents, escape_html)
      footnote_titles <- lapply(footnote_titles, escape_html)
    } else {
      footnote_contents <- lapply(footnote_contents, escape_latex2)
      footnote_titles <- lapply(footnote_titles, escape_latex2)
    }
  }
  footnote_table <- footnote_table_maker(
    kable_format, footnote_titles, footnote_contents
  )
  if (kable_format == "html") {
    return(footnote_html(kable_input, footnote_table, footnote_as_chunk))
  }
  if (kable_format == "latex") {
    return(footnote_latex(kable_input, footnote_table, footnote_as_chunk,
                          threeparttable))
  }
}

footnote_table_maker <- function(format, footnote_titles, footnote_contents) {
  number_index <- read.csv(system.file("symbol_index.csv",
                                       package = "kableExtra"))
  if (format == "latex") {
    symbol_index <- number_index$symbol.latex
  } else {
    symbol_index <- number_index$symbol.html
  }

  if (!is.null(footnote_contents$general)) {
    footnote_contents$general <- data.frame(
      index = "",
      footnote = footnote_contents$general
    )
  }
  if (!is.null(footnote_contents$number)) {
    footnote_contents$number <- data.frame(
      index = as.character(1:length(footnote_contents$number)),
      footnote = footnote_contents$number
    )
  }
  if (!is.null(footnote_contents$alphabet)) {
    footnote_contents$alphabet <- data.frame(
      index = letters[1:length(footnote_contents$alphabet)],
      footnote = footnote_contents$alphabet
    )
  }
  if (!is.null(footnote_contents$symbol)) {
    footnote_contents$symbol <- data.frame(
      index = symbol_index[1:length(footnote_contents$symbol)],
      footnote = footnote_contents$symbol
    )
  }

  out <- list()
  out$contents <- footnote_contents
  out$titles <- footnote_titles
  return(out)
}

# HTML
footnote_html <- function(kable_input, footnote_table, footnote_as_chunk) {
  kable_attrs <- attributes(kable_input)
  kable_xml <- read_kable_as_xml(kable_input)

  new_html_footnote <- html_tfoot_maker(footnote_table, footnote_as_chunk)
  xml_add_child(kable_xml, new_html_footnote)

  out <- as_kable_xml(kable_xml)
  attributes(out) <- kable_attrs
  if (!"kableExtra" %in% class(out)) class(out) <- c("kableExtra", class(out))
  return(out)
}

html_tfoot_maker <- function(footnote_table, footnote_as_chunk) {
  footnote_types <- names(footnote_table$contents)
  footnote_text <- c()
  for (i in footnote_types) {
    footnote_text <- c(footnote_text, html_tfoot_maker_(
      footnote_table$contents[[i]], footnote_table$titles[[i]], i,
      footnote_as_chunk))
  }
  footnote_text <- paste0(
    "<tfoot>", paste0(footnote_text, collapse = ""), "</tfoot>"
  )
  footnote_node <- read_html(footnote_text, options = c("RECOVER", "NOERROR"))
  return(xml_child(xml_child(footnote_node, 1), 1))
}

html_tfoot_maker_ <- function(ft_contents, ft_title, ft_type, ft_chunk) {
  footnote_text <- apply(ft_contents, 1, function(x) {
    paste0('<sup>', x[1], '</sup> ', x[2])
  })
  if (ft_title != "") {
    title_text <- paste0('<strong>', ft_title, '</strong>')
    footnote_text <- c(title_text, footnote_text)
  }
  if (!ft_chunk) {
    footnote_text <- paste0(
      '<tr><td style="padding: 0; border: 0;" colspan="100%">',
      footnote_text, '</td></tr>'
    )
  } else {
    footnote_text <- paste0(
      '<tr><td style="padding: 0; border: 0;" colspan="100%">',
      paste0(footnote_text, collapse = " "),
      '</td></tr>'
    )
  }
  return(footnote_text)
}

# LaTeX
footnote_latex <- function(kable_input, footnote_table, footnote_as_chunk,
                           threeparttable) {
  table_info <- magic_mirror(kable_input)
  out <- enc2utf8(as.character(kable_input))

  footnote_text <- latex_tfoot_maker(footnote_table, footnote_as_chunk,
                                     table_info$ncol, threeparttable)
  if (threeparttable) {
    if (table_info$tabular %in% c("longtable", "longtabu") ) {
      out <- sub(paste0("\\\\begin\\{", table_info$tabular, "\\}"),
                 paste0("\\\\begin{ThreePartTable}\n\\\\begin{TableNotes}",
                        ifelse(footnote_as_chunk, "[para]", ""),
                        "\n\\\\small\n", footnote_text,
                        "\n\\\\end{TableNotes}\n\\\\begin{",
                        table_info$tabular, "}"),
                 out)
      out <- sub(paste0("\\\\end\\{",table_info$tabular, "\\}"),
                 paste0("\\\\end{", table_info$tabular,
                        "}\n\\\\end{ThreePartTable}"),
                 out)
      if (table_info$booktabs) {
        out <- sub("\\\\bottomrule", "\\\\bottomrule\n\\\\insertTableNotes", out)
      } else {
        out <- sub("\\\\hline\n\\\\end\\{longtable\\}",
                   "\\\\hline\n\\\\insertTableNotes\n\\\\end\\{longtable\\}",
                   out)
      }
    } else {
      if (table_info$tabular == "tabu") {
        stop("Please use `longtable = T` in your kable function. ",
             "Full width threeparttable only works with longtable.")
      }
      out <- sub(paste0("\\\\begin\\{", table_info$tabular, "\\}"),
                 paste0("\\\\begin{threeparttable}\n\\\\begin{",
                        table_info$tabular, "}"),
                 out)
      out <- sub(table_info$end_tabular,
                 paste0("\\\\end{", table_info$tabular,
                        "}\n\\\\begin{tablenotes}",
                        ifelse(footnote_as_chunk, "[para]", ""),
                        "\n\\\\small\n", footnote_text,
                        "\n\\\\end{tablenotes}\n\\\\end{threeparttable}"),
                 out)
    }
  } else {
    if (table_info$booktabs) {
      out <- sub("\\\\bottomrule",
                 paste0("\\\\bottomrule\n", footnote_text), out)
    } else {
      out <- sub(table_info$end_tabular,
                 paste0(footnote_text, "\n\\\\end{", table_info$tabular, "}"),
                 out)
    }
  }

  out <- structure(out, format = "latex", class = "knitr_kable")
  attr(out, "kable_meta") <- table_info
  return(out)
}

latex_tfoot_maker <- function(footnote_table, footnote_as_chunk, ncol,
                              threeparttable) {
  footnote_types <- names(footnote_table$contents)
  footnote_text <- c()
  if (threeparttable) {
    for (i in footnote_types) {
      footnote_text <- c(footnote_text, latex_tfoot_maker_tpt_(
        footnote_table$contents[[i]], footnote_table$titles[[i]],
        footnote_as_chunk, ncol))
    }
  } else {
    for (i in footnote_types) {
      footnote_text <- c(footnote_text, latex_tfoot_maker_(
        footnote_table$contents[[i]], footnote_table$titles[[i]],
        footnote_as_chunk, ncol))
    }
  }
  footnote_text <- paste0(footnote_text, collapse = "\n")
  return(footnote_text)
}

latex_tfoot_maker_ <- function(ft_contents, ft_title, ft_chunk, ncol) {
  footnote_text <- apply(ft_contents, 1, function(x) {
    if (x[1] == "") {
      x[2]
    } else {
      paste0('\\\\textsuperscript{', x[1], '} ', x[2])
    }
  })
  if (ft_title != "") {
    title_text <- paste0('\\\\textbf{', ft_title, '} ')
    footnote_text <- c(title_text, footnote_text)
  }
  if (!ft_chunk) {
    footnote_text <- paste0(
      '\\\\multicolumn{', ncol, '}{l}{', footnote_text, '}\\\\\\\\'
    )
  } else {
    footnote_text <- paste0(
      '\\\\multicolumn{', ncol, '}{l}{',
      paste0(footnote_text, collapse = " "),
      '}\\\\\\\\'
    )
  }
  return(footnote_text)
}

latex_tfoot_maker_tpt_ <- function(ft_contents, ft_title, ft_chunk, ncol) {
  footnote_text <- apply(ft_contents, 1, function(x) {
    if (x[1] == "") {
      paste0('\\\\item ', x[2])
    } else {
      paste0('\\\\item[', x[1], '] ', x[2])
    }
  })
  if (ft_title != "") {
    title_text <- paste0('\\\\item \\\\textbf{', ft_title, '} ')
    footnote_text <- c(title_text, footnote_text)
  }
  footnote_text <- paste0(footnote_text, collapse = "\n")
  # if (!ft_chunk) {
  #   footnote_text <- paste0(footnote_text, collapse = "\n")
  # } else {
  #   footnote_text <- paste0(footnote_text, collapse = " ")
  # }
  return(footnote_text)
}