test_that("select quotes correctly", {
  out <- local_memdb_frame(x = 1, y = 1) |>
    select(x) |>
    collect()
  expect_equal(out, tibble(x = 1))
})

test_that("select can rename", {
  out <- local_memdb_frame(x = 1, y = 2) |>
    select(y = x) |>
    collect()
  expect_equal(out, tibble(y = 1))
})

test_that("two selects equivalent to one", {
  mf <- local_memdb_frame(a = 1, b = 1, c = 1, d = 2)

  out <- mf |>
    select(a:c) |>
    select(b:c) |>
    collect()

  expect_named(out, c("b", "c"))
})

test_that("correctly inlines across all verbs", {
  lf <- lazy_frame(x = 1, y = 2)

  # single table verbs
  expect_selects(lf |> arrange(x) |> select(x), 1)
  expect_selects(lf |> distinct() |> select(x), 2)
  expect_selects(lf |> filter(x == 1) |> select(x), 1)
  expect_selects(lf |> head(1) |> select(x), 1)
  expect_selects(lf |> mutate(z = x + 1) |> select(x), 1)
  expect_selects(lf |> select(y = x) |> select(y), 1)
  expect_selects(lf |> summarise(y = mean(x)) |> select(y), 1)

  # two table verbs
  lf2 <- lazy_frame(x = 1)
  expect_selects(lf |> left_join(lf2, by = "x") |> select(x), 1)
  expect_selects(lf |> right_join(lf2, by = "x") |> select(x), 1)
  expect_selects(lf |> semi_join(lf2, by = "x") |> select(x), 2)
  expect_selects(lf |> union(lf2) |> select(x), 3)
})

test_that("select after distinct produces subquery", {
  lf <- lazy_frame(x = 1, y = 1:2)
  expect_snapshot(
    lf |> distinct() |> select(x)
  )

  out <- lf |> distinct() |> select(x)
  lq <- out$lazy_query
  expect_false(lq$distinct)
  expect_true(lq$x$distinct)
})

test_that("select after arrange produces subquery, if needed", {
  lf <- lazy_frame(x = 1)

  # shouldn't inline
  out <- lf |> mutate(z = 2) |> arrange(x, z) |> select(x)
  # should inline
  out2 <- lf |> mutate(z = 2) |> arrange(x, z) |> select(x, z)

  inner_query <- out$lazy_query$x
  expect_s3_class(inner_query, "lazy_select_query")
  expect_equal(inner_query$order_by, list(expr(x), expr(z)))
  expect_equal(
    op_vars(inner_query),
    c("x", "z")
  )
  expect_equal(op_vars(out$lazy_query), "x")

  # order vars in a subquery are dropped
  expect_equal(
    inner_query[setdiff(names(inner_query), "order_vars")],
    out2$lazy_query[setdiff(names(out2$lazy_query), "order_vars")]
  )
})


test_that("rename/relocate after distinct is inlined #1141", {
  lf <- lazy_frame(x = 1, y = 1:2)
  expect_snapshot({
    lf |> distinct() |> rename(z = y)
    lf |> distinct() |> relocate(y)
  })

  out <- lf |> distinct() |> rename(z = y)
  lq <- out$lazy_query
  expect_true(lq$distinct)

  out <- lf |> distinct() |> relocate(y)
  lq <- out$lazy_query
  expect_true(lq$distinct)
})

test_that("select operates on mutated vars", {
  mf <- local_memdb_frame(x = c(1, 2, 3), y = c(3, 2, 1))
  expect_equal(mf |> mutate(x, z = x + y) |> select(z) |> pull(), c(4, 4, 4))
})

test_that("select renames variables (#317)", {
  mf <- local_memdb_frame(x = 1, y = 2)
  expect_equal(mf |> select(A = x) |> collect(), tibble(A = 1))
})

test_that("rename renames variables", {
  mf <- local_memdb_frame(x = 1, y = 2)
  expect_equal(mf |> rename(A = x) |> collect(), tibble(A = 1, y = 2))
})

test_that("can rename multiple vars", {
  mf <- local_memdb_frame(a = 1, b = 2)
  exp <- tibble(c = 1, d = 2)

  expect_equal(mf |> rename(c = a, d = b) |> collect(), exp)
  expect_equal(
    mf |> group_by(a) |> rename(c = a, d = b) |> collect(),
    exp |> group_by(c)
  )
})

test_that("can rename with a function", {
  mf <- local_memdb_frame(a = 1, b = 2)

  expect_named(mf |> rename_with(toupper) |> collect(), c("A", "B"))
  expect_named(mf |> rename_with(toupper, 1) |> collect(), c("A", "b"))
})

test_that("select preserves grouping vars", {
  mf <- local_memdb_frame(a = 1, b = 2) |> group_by(b)
  expect_snapshot(out <- mf |> select(a) |> collect())

  expect_named(out, c("b", "a"))
})

test_that("select handles order vars", {
  lf <- lazy_frame(x = 1, y = 1, z = 1)
  # can drop order vars
  expect_equal(lf |> window_order(y) |> select(-y) |> op_sort(), list())
  expect_equal(
    lf |> window_order(desc(y)) |> select(-y) |> op_sort(),
    list()
  )
  # can rename order vars
  expect_equal(
    lf |> window_order(y) |> select(y2 = y) |> op_sort(),
    list(expr(y2))
  )
  expect_equal(
    lf |> window_order(desc(y)) |> select(y2 = y) |> op_sort(),
    list(expr(desc(y2)))
  )
  # keeps sort order
  expect_equal(
    lf |> window_order(x, y) |> select(y2 = y, x) |> op_sort(),
    list(expr(x), expr(y2))
  )
})

test_that("select doesn't relocate grouping vars to the front", {
  mf <- local_memdb_frame(a = 1, b = 2) |> group_by(b)
  expect_equal(mf |> select(a, b) |> op_vars(), c("a", "b"))
})

test_that("relocate works", {
  mf <- local_memdb_frame(a = 1, b = 2, c = 1) |> group_by(b)

  out1 <- mf |> relocate(c) |> collect()
  expect_named(out1, c("c", "a", "b"))
  out2 <- mf |> relocate(a, .after = c) |> collect()
  expect_named(out2, c("b", "c", "a"))
})

test_that("relocate can rename variables", {
  mf <- local_memdb_frame(a = 1, b = 2, c = 1) |> group_by(b)

  out1 <- mf |> relocate(d = b) |> collect()
  expect_named(out1, c("d", "a", "c"))
  expect_equal(group_vars(out1), "d")
})

test_that("only add step if necessary", {
  lf <- lazy_frame(x = 1:3, y = 1:3)
  expect_equal(lf |> select(everything()), lf)
  expect_equal(lf |> select(x, y), lf)

  expect_equal(lf |> rename(x = x), lf)
  expect_equal(lf |> rename(), lf)

  expect_equal(lf |> relocate(x, y), lf)
  expect_equal(lf |> relocate(), lf)
})

test_that("select() after left_join() is inlined", {
  lf1 <- lazy_frame(x = 1, a = 1, .name = "lf1")
  lf2 <- lazy_frame(x = 1, b = 2, .name = "lf2")

  expect_snapshot(
    (out <- left_join(lf1, lf2, by = "x") |>
      select(b, x))
  )
  expect_equal(op_vars(out), c("b", "x"))

  expect_snapshot(
    (out <- left_join(lf1, lf2, by = "x") |>
      relocate(b))
  )
  expect_equal(op_vars(out), c("b", "x", "a"))
  expect_equal(out$lazy_query$vars$var, c("b", "x", "a"))
  expect_equal(out$lazy_query$vars$table, c(2L, 1L, 1L))

  out <- left_join(lf1, lf2, by = "x") |>
    transmute(b, x = x + 1)
  expect_s3_class(out$lazy_query, "lazy_select_query")
})

test_that("select() after filter() on join preserves qualified where", {
  lf1 <- lazy_frame(x = 1, y = 2, .name = "t1")
  lf2 <- lazy_frame(x = 1, z = 3, .name = "t2")

  # drop a column referenced in $where: must fall through to lazy_select_query
  # so the join's WHERE still qualifies z as t2.z.
  out <- lf1 |>
    left_join(lf2, by = "x") |>
    filter(z > 0) |>
    select(-z)
  expect_s3_class(out$lazy_query, "lazy_select_query")
  expect_snapshot(out)

  # rename a column referenced in $where: same handling.
  out <- lf1 |>
    left_join(lf2, by = "x") |>
    filter(z > 0) |>
    rename(zz = z)
  expect_s3_class(out$lazy_query, "lazy_select_query")
  expect_snapshot(out)

  # dropping a column not in $where stays inline.
  out <- lf1 |>
    left_join(lf2, by = "x") |>
    filter(z > 0) |>
    select(-y)
  expect_s3_class(out$lazy_query, "lazy_multi_join_query")
})

test_that("select() after semi_join() is inlined", {
  lf1 <- lazy_frame(x = 1, a = 1, .name = "lf1")
  lf2 <- lazy_frame(x = 1, b = 2, .name = "lf2")

  expect_snapshot(
    (out <- semi_join(lf1, lf2, by = "x") |>
      select(x, a2 = a))
  )
  expect_equal(op_vars(out), c("x", "a2"))

  expect_snapshot(
    (out <- anti_join(lf1, lf2, by = "x") |>
      relocate(a))
  )
  expect_equal(op_vars(out), c("a", "x"))

  out <- semi_join(lf1, lf2, by = "x") |>
    transmute(a, x = x + 1)
  expect_s3_class(out$lazy_query, "lazy_select_query")
})

test_that("select() after join handles previous select", {
  lf <- lazy_frame(x = 1, y = 1, z = 1) |>
    group_by(x, y, z) |>
    select(x, y2 = y, z) |>
    semi_join(
      lazy_frame(x = 1),
      by = "x"
    ) |>
    select(x2 = x, y3 = y2, z)

  expect_equal(op_vars(lf), c("x2", "y3", "z"))
  expect_equal(
    lf$lazy_query$vars,
    tibble(
      name = c("x2", "y3", "z"),
      var = c("x", "y", "z")
    )
  )
  expect_equal(op_grps(lf), c("x2", "y3", "z"))
  expect_snapshot(print(lf))

  lf2 <- lazy_frame(x = 1, y = 1, z = 1) |>
    group_by(x, y, z) |>
    select(x, y2 = y, z) |>
    left_join(
      lazy_frame(x = 1, y = 1),
      by = "x"
    ) |>
    select(x2 = x, y3 = y2, z)

  expect_equal(op_vars(lf2), c("x2", "y3", "z"))
  vars2 <- lf2$lazy_query$vars
  expect_equal(vars2$var, c("x", "y", "z"))
  expect_equal(vars2$table, c(1L, 1L, 1L))

  expect_equal(op_grps(lf2), c("x2", "y3", "z"))
  expect_snapshot(print(lf2))
})

test_that("select() afer join keeps grouping", {
  lf1 <- lazy_frame(x = 1, y = 1) |> group_by(y)
  lf2 <- lazy_frame(x = 1, z = 1) |> group_by(z)

  # just to be sure check without select/renaming
  expect_equal(left_join(lf1, lf2, by = "x") |> op_grps(), "y")

  # rename grouping variable
  expect_equal(
    left_join(lf1, lf2, by = "x") |>
      select(y2 = y) |>
      op_grps(),
    "y2"
  )
})

test_that("select() produces nice error messages", {
  lf <- lazy_frame(x = 1)

  expect_snapshot(error = TRUE, {
    lf |> select(non_existent)
    lf |> select(non_existent + 1)
  })

  expect_snapshot(error = TRUE, {
    lf |> relocate(non_existent)
    lf |> relocate(non_existent + 1)
  })

  expect_snapshot(error = TRUE, {
    # no name
    lf |> rename(x)
    # non-existing column
    lf |> rename(y = non_existent)
    lf |> rename(y = non_existent + 1)
  })

  expect_snapshot(error = TRUE, {
    lf |> rename_with(toupper, .cols = non_existent)
    lf |> rename_with(toupper, .cols = non_existent + 1)
  })
})

test_that("where() isn't suppored", {
  lf <- lazy_frame(x = 1)
  expect_snapshot(error = TRUE, {
    lf |> select(where(is.integer))
  })
})

test_that("arranged computed columns are not inlined away", {
  lf <- lazy_frame(x = 1)

  # shouldn't inline
  out <- lf |> mutate(z = 2) |> arrange(x, z) |> select(x)
  # should inline
  out2 <- lf |> mutate(z = 2) |> arrange(x, z) |> select(x, z)

  inner_query <- out$lazy_query$x
  expect_snapshot({
    lf |> mutate(z = 1) |> arrange(x, z) |> select(x)
  })
  expect_s3_class(inner_query, "lazy_select_query")
  expect_equal(inner_query$order_by, list(expr(x), expr(z)))
  expect_equal(op_vars(inner_query), c("x", "z"))
  expect_equal(op_vars(out$lazy_query), "x")
  expect_equal(
    # order vars in a subquery are dropped
    inner_query[setdiff(names(inner_query), "order_vars")],
    out2$lazy_query[setdiff(names(out2$lazy_query), "order_vars")]
  )
})

# sql_render --------------------------------------------------------------

test_that("multiple selects are collapsed", {
  lf <- lazy_frame(x = 1, y = 2)

  expect_snapshot(lf |> select(2:1) |> select(2:1))
  expect_snapshot(lf |> select(2:1) |> select(2:1) |> select(2:1))
  expect_snapshot(lf |> select(x1 = x) |> select(x2 = x1))
})

test_that("mutate collapses over nested select", {
  lf <- lazy_frame(g = 0, x = 1, y = 2)

  expect_snapshot(lf |> mutate(a = 1, b = 2) |> select(a))
  expect_snapshot(lf |> mutate(a = 1, b = 2) |> select(x))
})

test_that("output is styled", {
  local_reproducible_output(crayon = TRUE)
  withr::local_options(dbplyr_highlight = cli::combine_ansi_styles("blue"))

  lf <- lazy_frame(x = 1, y = 1, z = 1)
  out <- lf |>
    group_by(x) |>
    mutate(y = mean(y, na.rm = TRUE), z = z + 1) |>
    filter(z == 1) |>
    left_join(lf, by = "x")

  expect_snapshot(show_query(out, sql_options = sql_options(cte = TRUE)))
})

# sql_build -------------------------------------------------------------

test_that("select picks variables", {
  out <- lazy_frame(x1 = 1, x2 = 1, x3 = 2) |>
    select(x1:x2) |>
    sql_build()

  expect_equal(out$select, sql("\"x1\"", "\"x2\""))
})

test_that("select renames variables", {
  out <- lazy_frame(x1 = 1, x2 = 1, x3 = 2) |>
    select(y = x1, z = x2) |>
    sql_build()

  expect_equal(out$select, sql("\"x1\" AS \"y\"", "\"x2\" AS \"z\""))
})

test_that("select can refer to variables in local env", {
  vars <- c("x", "y")
  out <- lazy_frame(x = 1, y = 1, z = 1) |>
    select(dplyr::one_of(vars)) |>
    sql_build()

  expect_equal(out$select, sql("\"x\"", "\"y\""))
})

test_that("rename preserves existing vars", {
  out <- lazy_frame(x = 1, y = 1) |>
    rename(z = y) |>
    sql_build()

  expect_equal(out$select, sql("\"x\"", "\"y\" AS \"z\""))
})


# ops ---------------------------------------------------------------------

test_that("select reduces variables", {
  out <- mtcars |> tbl_lazy() |> select(mpg:disp)
  expect_equal(op_vars(out), c("mpg", "cyl", "disp"))
})

test_that("rename preserves existing", {
  out <- tibble(x = 1, y = 2) |> tbl_lazy() |> rename(z = y)
  expect_equal(op_vars(out), c("x", "z"))
})

test_that("rename renames grouping vars", {
  df <- lazy_frame(a = 1, b = 2)
  expect_equal(df |> group_by(a) |> rename(c = a) |> op_grps(), "c")
})

test_that("mutate preserves grouping vars (#396)", {
  df <- lazy_frame(a = 1, b = 2, c = 3) |> group_by(a, b)
  expect_equal(df |> mutate(a = 1) |> op_grps(), c("a", "b"))
  expect_equal(df |> mutate(b = 1) |> op_grps(), c("a", "b"))
})

test_that("select after arrange(desc()) works (#1206)", {
  out <- lazy_frame(x = 1, y = 1) |>
    arrange(desc(x)) |>
    select(x)
  expect_equal(op_vars(out), "x")
})


# lazy_select_query -------------------------------------------------------

test_that("select, relocate, and rename work", {
  lf <- lazy_frame(x = 1, y = 1)

  expect_equal(
    select(lf, x)$lazy_query$select,
    new_lazy_select(exprs(x = x))
  )

  expect_equal(
    relocate(lf, y)$lazy_query$select,
    new_lazy_select(exprs(y = y, x = x))
  )

  expect_equal(
    rename(lf, b = y, a = x)$lazy_query$select,
    new_lazy_select(exprs(a = x, b = y))
  )
})

test_that("renaming handles groups correctly", {
  lf <- lazy_frame(x = 1, y = 1) |>
    group_by(x) |>
    rename(ax = x)

  result <- lf$lazy_query
  expect_equal(
    result$select,
    new_lazy_select(exprs(ax = x, y = y))
  )

  expect_equal(result$group_vars, "ax")
  expect_equal(op_grps(result), "ax")

  result <- rename(lf, x = ax)$lazy_query
  expect_equal(
    result$select,
    new_lazy_select(exprs(x = x, y = y))
  )

  expect_equal(result$group_vars, "x")
  expect_equal(op_grps(result), "x")

  # https://github.com/tidyverse/dbplyr/issues/928
  result <- lazy_frame(cyl = 1, vs = 1) |>
    rename(vs = cyl, new_vs = vs) |>
    group_by(vs)
  expect_equal(op_grps(result), "vs")
})
