#' Schedule a grid
#'
#' @param grid A tibble containing the parameter grid.
#' @param wflow The workflow object for which we schedule the grid.
#'
#' @returns A schedule object, inheriting from either 'single_schedule',
#' 'grid_schedule', or 'resample_schedule'.
#'
#' @keywords internal
#' @export
schedule_grid <- function(grid, wflow) {
  if (!tibble::is_tibble(grid)) {
    cli::cli_abort(
      "Argument {.arg grid} must be a tibble, not {.obj_type_friendly {grid}}."
    )
  }
  if (!inherits(wflow, "workflow")) {
    cli::cli_abort(
      "Argument {.arg wflow} must be a workflow object, not {.obj_type_friendly {wflow}}."
    )
  }

  schedule <- schedule_stages(grid, wflow)

  param_info <- tune_args(wflow)

  og_cls <- class(schedule)
  if (nrow(param_info) == 0) {
    cls <- "resample_schedule"
  } else {
    cls <- "grid_schedule"
  }

  if (nrow(grid) == 1) {
    cls <- c("single_schedule", cls)
  }

  class(schedule) <- c(cls, "schedule", og_cls)

  schedule
}

schedule_stages <- function(grid, wflow) {
  # Which parameter belongs to which stage and which is a submodel parameter?
  param_info <- get_param_info(wflow)

  model_source <- get_model_source(wflow)

  # schedule preprocessing stage and push the rest into a nested tibble
  param_pre_stage <- param_info |>
    dplyr::filter(source == "recipe") |>
    dplyr::pull(id)
  schedule <- grid |>
    tidyr::nest(.by = dplyr::all_of(param_pre_stage), .key = "model_stage")

  # schedule next stages nested within `schedule_model_stage_i()`
  schedule |>
    dplyr::mutate(
      model_stage = purrr::map(
        model_stage,
        schedule_model_stage_i,
        param_info = param_info,
        wflow = wflow,
        model_source = model_source
      )
    )
}

schedule_model_stage_i <- function(
  model_stage,
  param_info,
  wflow,
  model_source = "model_spec"
) {
  model_param <- param_info |>
    dplyr::filter(source == model_source) |>
    dplyr::pull(id)
  non_submodel_param <- param_info |>
    dplyr::filter(source == model_source & !has_submodel) |>
    dplyr::pull(id)

  any_non_submodel_param <- length(non_submodel_param) > 0

  # schedule model parameters
  schedule <- min_model_grid(model_stage, model_param, wflow)

  # push remaining parameters into the next stage
  next_stage <- model_stage |>
    tidyr::nest(
      .by = dplyr::all_of(non_submodel_param),
      .key = "predict_stage"
    )

  if (any_non_submodel_param) {
    # min_model_grid() may change the row order, thus use next_stage as the
    # "left" data frame here to preserve the original row order
    schedule <- next_stage |>
      dplyr::left_join(schedule, by = non_submodel_param) |>
      dplyr::relocate(dplyr::all_of(model_param))
  } else {
    schedule <- dplyr::bind_cols(schedule, next_stage)
  }

  # schedule next stages nested within `schedule_predict_stage_i()`
  schedule |>
    dplyr::mutate(
      predict_stage = purrr::map(
        predict_stage,
        schedule_predict_stage_i,
        param_info = param_info
      )
    )
}

min_model_grid <- function(grid, model_param, wflow) {
  # work on only the model parameters
  model_grid <- grid |>
    dplyr::select(dplyr::all_of(model_param)) |>
    dplyr::distinct()

  if (nrow(model_grid) < 1) {
    return(model_grid)
  }

  # But we check to see if there are _any_ model parameters. If not,
  # act like `min_grid.model_spec()` and just add a placeholder for submodels.
  # We want to avoid method dispatch to a more specific method (e.g.
  # `min_grid.logistic_reg()`) because it assumes that _there is_ a model
  # parameter and, at least for glmnet models, it fails. See issue #1119
  if (ncol(model_grid) == 0) {
    res <-
      min_grid.model_spec(extract_spec_parsnip(wflow), model_grid) |>
      dplyr::select(dplyr::all_of(model_param))
  } else {
    res <-
      min_grid(extract_spec_parsnip(wflow), model_grid) |>
      dplyr::select(dplyr::all_of(model_param))
  }
  res
}

schedule_predict_stage_i <- function(predict_stage, param_info) {
  submodel_param <- param_info |>
    dplyr::filter(source == "model_spec" & has_submodel) |>
    dplyr::pull(id)

  predict_stage |>
    tidyr::nest(
      .by = dplyr::all_of(submodel_param),
      .key = "post_stage"
    )
}

get_param_info <- function(wflow) {
  param_info <- tune_args(wflow) |>
    dplyr::select(name, id, source)

  model_spec <- extract_spec_parsnip(wflow)
  model_type <- class(model_spec)[1]
  model_eng <- model_spec$engine

  # NULL for tidyclust workflows since tidyclust models are not in parsnip's
  # registry
  model_param <- parsnip::get_from_env(paste0(model_type, "_args"))

  if (!is.null(model_param)) {
    model_param <- model_param |>
      dplyr::filter(engine == model_spec$engine) |>
      dplyr::select(name = parsnip, has_submodel)
    param_info <- dplyr::left_join(param_info, model_param, by = "name")
  } else {
    # tidyclust models never use submodels
    param_info$has_submodel <- FALSE
  }

  if (!"has_submodel" %in% names(param_info)) {
    param_info$has_submodel <- FALSE
  } else {
    param_info$has_submodel[is.na(param_info$has_submodel)] <- FALSE
  }

  param_info
}

# Determine the source type for model parameters
# Returns "model_spec" for parsnip models, "cluster_spec" for tidyclust models
get_model_source <- function(wflow) {
  spec <- extract_spec_parsnip(wflow)
  if (inherits(spec, "cluster_spec")) {
    "cluster_spec"
  } else {
    "model_spec"
  }
}
