#' Kaplan-Meier plot by provider risk group
#'
#' @param surv_profile Data frame from vowr_survival_merge()
#' @return A survminer Kaplan-Meier plot object
#' @export
vowr_km_by_risk <- function(surv_profile) {

  required <- c("time_days", "status", "risk_flag")
  missing_cols <- setdiff(required, names(surv_profile))

  if (length(missing_cols) > 0) {
    stop("Missing required columns: ", paste(missing_cols, collapse = ", "))
  }

  plot_data <- surv_profile[!is.na(surv_profile$risk_flag), ]

  plot_data$risk_group <- dplyr::case_when(
    plot_data$risk_flag == "Critical (Breach & Unstable)" ~ "Critical",
    plot_data$risk_flag == "High (Breach)" ~ "High",
    plot_data$risk_flag == "Warning (Unstable)" ~ "Warning",
    plot_data$risk_flag == "Stable" ~ "Stable",
    TRUE ~ as.character(plot_data$risk_flag)
  )

  plot_data$risk_group <- factor(
    plot_data$risk_group,
    levels = c("Stable", "Warning", "High", "Critical")
  )

  fit <- survival::survfit(
    survival::Surv(time_days, status) ~ risk_group,
    data = plot_data
  )

  survminer::ggsurvplot(
    fit,
    data = plot_data,
    risk.table = FALSE,
    conf.int = FALSE,
    pval = TRUE,
    pval.coord = c(45, 0.18),
    legend.title = "Risk group",
    legend.labs = c("Stable", "Warning", "High", "Critical"),
    xlab = "Days until threshold breach",
    ylab = "Probability of remaining below threshold",
    break.time.by = 200,
    censor.shape = "|",
    censor.size = 3,
    linewidth = 1.1,
    ggtheme = ggplot2::theme_minimal(base_size = 14) +
      ggplot2::theme(
        legend.position = "top",
        legend.title = ggplot2::element_text(face = "bold"),
        plot.title = ggplot2::element_text(face = "bold", hjust = 0.5),
        axis.title = ggplot2::element_text(face = "bold"),
        panel.grid.minor = ggplot2::element_blank()
      )
  )
}
