#' Plot Moderation Effect
#'
#' Produces a plot of the moderation effect.
#'
#' @param object A fitted model
#'
#' @return A ggplot object
#'
#' @examples
#' set.seed(123)
#'
#' n <- 100
#' x <- rnorm(n)
#' z <- 0.5 * x + sqrt(1 - 0.5^2) * rnorm(n)
#'
#' y <- 0.3 * x + 0.3 * z + 0.8 * x * z + rnorm(n)
#'
#' data <- data.frame(x, z, y)
#'
#' result <- moderated_regression(data, iv = "x", moderator = "z", dv = "y")
#'
#' plot_moderation(result)
#'
#' @export

plot_moderation <- function(object) {

  if (!inherits(object, "modlr")) {
    stop("object must be of class modlr")
  }

  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("ggplot2 package required")
  }

  data <- object$data
  iv   <- object$iv
  mod  <- object$moderator
  dv   <- object$dv

  # ✅ Use RAW model (consistent with JN + slopes)
  mod_raw <- lm(
    as.formula(paste(dv, "~", iv, "*", mod)),
    data = data
  )

  # ✅ Generate x sequence
  x_seq <- seq(
    min(data[[iv]], na.rm = TRUE),
    max(data[[iv]], na.rm = TRUE),
    length.out = 100
  )

  # ===============================
  # ✅ Define moderator values
  # ===============================
  if (object$type_moderator == "binary") {

    z_vals <- sort(unique(data[[mod]]))
    labels <- paste(mod, "=", z_vals)

  } else {

    z_mean <- mean(data[[mod]], na.rm = TRUE)
    z_sd   <- sd(data[[mod]], na.rm = TRUE)

    z_vals <- c(z_mean - z_sd, z_mean, z_mean + z_sd)
    labels <- c("Low (-1 SD)", "Mean", "High (+1 SD)")
  }

  # ===============================
  # ✅ Build prediction data
  # ===============================
  pred_data <- data.frame()

  for (i in seq_along(z_vals)) {

    temp <- data.frame(x_seq, rep(z_vals[i], length(x_seq)))
    names(temp) <- c(iv, mod)

    # ✅ Predict from model (handles covariates automatically)
    temp$.pred <- predict(mod_raw, newdata = temp)

    temp$group <- labels[i]

    pred_data <- rbind(pred_data, temp)
  }

  # ===============================
  # ✅ Plot
  # ===============================
  p <- ggplot2::ggplot(
    pred_data,
    ggplot2::aes(x = .data[[iv]], y = .data[[".pred"]], color = .data[["group"]])
    ) +
    ggplot2::geom_line(linewidth = 1.2) +
    ggplot2::labs(
      title = "Moderation Effect",
      subtitle = paste("Interaction p =", round(object$interaction_p, 4)),
      x = iv,
      y = paste("Predicted", dv),
      color = "Moderator"
    ) +
    ggplot2::theme_bw() +
    ggplot2::theme(legend.position = "bottom")


  # -------------------------------
  # Covariate note
  # -------------------------------
  if (!is.null(object$covariates) && length(object$covariates) > 0) {
    message("Note: Results are conditional on the included covariates (held constant).\n")
  }

  return(p)


}
