mock_ns_fun <- function(pkg, fun, replacement) {
  old <- get(fun, envir = asNamespace(pkg), inherits = FALSE)
  assignInNamespace(fun, replacement, ns = pkg)
  list(pkg = pkg, fun = fun, old = old)
}

restore_ns_fun <- function(mock_info) {
  assignInNamespace(mock_info$fun, mock_info$old, ns = mock_info$pkg)
}


test_that("internal helper .getDbms resolves dbms or fails", {
  skip_if_not_installed("DatabaseConnector")

  conn <- structure(list(), dbms = "sql server")

  m <- mock_ns_fun("DatabaseConnector", "dbms", function(connection) "postgresql")
  on.exit(restore_ns_fun(m), add = TRUE)
  expect_equal(OdysseusSurvivalModule:::.getDbms(conn), "postgresql")

  m2 <- mock_ns_fun("DatabaseConnector", "dbms", function(connection) stop("boom"))
  on.exit(restore_ns_fun(m2), add = TRUE)
  expect_equal(OdysseusSurvivalModule:::.getDbms(conn), "sql server")

  conn_no_attr <- list()
  expect_error(OdysseusSurvivalModule:::.getDbms(conn_no_attr), "Could not determine DBMS")
})


test_that("internal helpers validate and build supporting objects", {
  expect_invisible(
    OdysseusSurvivalModule:::.validateSurvivalInputs(
      outcomeDateVariable = "cohort_start_date",
      chunkSize = 10,
      outcomeWashout = Inf,
      minDaysToEvent = 0,
      followUpDays = 365
    )
  )

  expect_error(
    OdysseusSurvivalModule:::.validateSurvivalInputs("bad_col", 10, Inf, 0, 10),
    "outcomeDateVariable"
  )
  expect_error(
    OdysseusSurvivalModule:::.validateSurvivalInputs("cohort_start_date", 0, Inf, 0, 10),
    "chunkSize"
  )
  expect_error(
    OdysseusSurvivalModule:::.validateSurvivalInputs("cohort_start_date", 10, -1, 0, 10),
    "outcomeWashout"
  )
  expect_error(
    OdysseusSurvivalModule:::.validateSurvivalInputs("cohort_start_date", 10, 0, -1, 10),
    "minDaysToEvent"
  )
  expect_error(
    OdysseusSurvivalModule:::.validateSurvivalInputs("cohort_start_date", 10, 0, 0, -1),
    "followUpDays"
  )

  tmp <- OdysseusSurvivalModule:::.generateTempNames()
  expect_equal(sort(names(tmp)), sort(c("target", "obs", "coh_obs", "outcome", "washout", "events", "result", "id")))
  expect_match(tmp$target, "^#surv_tgt_")
  expect_match(tmp$result, "^#surv_res_")

  expr <- OdysseusSurvivalModule:::.buildCensorExpression(
    censorOnCohortExit = TRUE,
    hasCensorDate = TRUE,
    hasFollowUpLimit = TRUE,
    followUpDays = 100
  )
  expect_match(expr, "CASE WHEN")
  expect_match(expr, "cohort_end_date")
  expect_match(expr, "@censor_date")

  sql_template <- OdysseusSurvivalModule:::.getSurvivalSqlTemplate()
  expect_true(is.character(sql_template))
  expect_match(sql_template, "@temp_target")
  expect_match(sql_template, "@outcome_database_schema")
})


test_that("addCohortSurvival runs non-chunked path and returns base columns", {
  skip_if_not_installed("DatabaseConnector")
  skip_if_not_installed("SqlRender")

  exec_calls <- 0L

  m_dbms <- mock_ns_fun("DatabaseConnector", "dbms", function(connection) "sql server")
  m_exec <- mock_ns_fun("DatabaseConnector", "executeSql", function(connection, sql, ...) {
    exec_calls <<- exec_calls + 1L
    invisible(NULL)
  })
  m_query <- mock_ns_fun("DatabaseConnector", "querySql", function(connection, sql, ...) {
    data.frame(
      subject_id = 1:2,
      time = c(10, 20),
      status = c(1, 0),
      age_years = c(50, 60),
      gender = c("Male", "Female")
    )
  })
  on.exit(restore_ns_fun(m_query), add = TRUE)
  on.exit(restore_ns_fun(m_exec), add = TRUE)
  on.exit(restore_ns_fun(m_dbms), add = TRUE)

  out <- OdysseusSurvivalModule:::addCohortSurvival(
    connection = structure(list(), dbms = "sql server"),
    cdmDatabaseSchema = "cdm",
    cohortDatabaseSchema = "cohort",
    targetCohortTable = "cohort_table",
    targetCohortId = 1,
    outcomeCohortTable = "outcome_table",
    outcomeCohortId = 2,
    followUpDays = 180,
    includeAge = FALSE,
    includeGender = FALSE,
    chunkSize = NULL
  )

  expect_true(exec_calls > 0)
  expect_equal(names(out), c("subject_id", "time", "status"))
  expect_equal(nrow(out), 2)
})


test_that("addCohortSurvival runs chunked path and includes demographics", {
  skip_if_not_installed("DatabaseConnector")
  skip_if_not_installed("SqlRender")

  query_calls <- 0L

  m_dbms <- mock_ns_fun("DatabaseConnector", "dbms", function(connection) "sql server")
  m_exec <- mock_ns_fun("DatabaseConnector", "executeSql", function(connection, sql, ...) invisible(NULL))
  m_query <- mock_ns_fun("DatabaseConnector", "querySql", function(connection, sql, ...) {
    query_calls <<- query_calls + 1L
    if (query_calls == 1L) {
      return(data.frame(
        subject_id = 1:2,
        time = c(12, 30),
        status = c(1, 0),
        age_years = c(42, 57),
        gender = c("Female", "Male")
      ))
    }
    data.frame(
      subject_id = integer(0),
      time = numeric(0),
      status = integer(0),
      age_years = numeric(0),
      gender = character(0)
    )
  })
  on.exit(restore_ns_fun(m_query), add = TRUE)
  on.exit(restore_ns_fun(m_exec), add = TRUE)
  on.exit(restore_ns_fun(m_dbms), add = TRUE)

  out <- OdysseusSurvivalModule:::addCohortSurvival(
    connection = structure(list(), dbms = "sql server"),
    cdmDatabaseSchema = "cdm",
    cohortDatabaseSchema = "cohort",
    targetCohortTable = "cohort_table",
    targetCohortId = 1,
    outcomeCohortTable = "outcome_table",
    outcomeCohortId = 2,
    includeAge = TRUE,
    includeGender = TRUE,
    chunkSize = 2,
    outcomeWashout = 30,
    minDaysToEvent = 2,
    censorOnCohortExit = TRUE,
    censorOnDate = as.Date("2024-12-31"),
    followUpDays = 365,
    addDay = TRUE
  )

  expect_true(query_calls >= 2)
  expect_equal(names(out), c("subject_id", "time", "status", "age_years", "gender"))
  expect_equal(nrow(out), 2)
})