#' @include get_predict.R
#' @rdname get_predict
#' @keywords internal
#' @export
get_predict.glmmTMB <- function(model,
                                newdata = insight::get_data(model),
                                type = "response",
                                newparams = NULL,
                                ...) {
  if (inherits(vcov, "vcov.glmmTMB")) {
    vcov <- vcov[[1]]
  }

  # hack to avoid re-optimization
  # see https://github.com/vincentarelbundock/marginaleffects/issues/1064
  b_vec <- model$obj$env$parList()$b
  if (length(b_vec) > 0) {
    model$modelInfo$map$b <- factor(rep(NA, length(b_vec)))
  }

  np <- model$fit$par
  if (!is.null(newparams)) {
    np[seq_along(newparams)] <- newparams
  }

  out <- get_predict.default(
    model = model,
    newdata = newdata,
    type = type,
    allow.new.levels = TRUE, # otherwise we get errors in marginal_means()
    newparams = np,
    ...)

  return(out)
}



#' @include get_vcov.R
#' @rdname get_vcov
#' @export
get_vcov.glmmTMB <- function(model, ...) {
  # Extract the full covariance matrix
  out <- stats::vcov(model, full = TRUE)

  # Extract the fixed-effect coefficient names from get_coef
  coef_names <- names(get_coef.glmmTMB(model))

  # Handle dispersion and conditional terms
  cleaned_coef_names <- gsub("^cond~", "", coef_names) # Remove cond~ for conditional terms
  cleaned_coef_names <- gsub("^disp~", "d~", cleaned_coef_names) # Map disp~ to d~ for dispersion terms

  # The 'upper cutoff' and 'lower cutoff' will remain in both, so no removal

  # Get the current row and column names from the covariance matrix
  current_names <- rownames(out)

  # Match cleaned coef_names with current names in the covariance matrix
  matched_indices <- match(current_names, cleaned_coef_names)

  # Replace row/column names only where there is a valid match
  valid_indices <- which(!is.na(matched_indices))

  if (length(valid_indices) > 0) {
    # Apply the correct names from coef_names to matched rows/columns in the covariance matrix
    rownames(out)[valid_indices] <- coef_names[matched_indices[valid_indices]]
    colnames(out)[valid_indices] <- coef_names[matched_indices[valid_indices]]
  } else {
    warning("No matching terms found between the covariance matrix and fixed-effect coefficients.")
  }

  return(out)
}


#' @include get_coef.R
#' @rdname get_coef
#' @export
get_coef.glmmTMB <- function(model, ...) {
  # Extract the fixed-effect coefficients
  out <- unlist(glmmTMB::fixef(model))

  # Apply the gsub logic to rename terms (cond~, disp~, etc.)
  names(out) <- gsub("^(cond|zi|disp)\\.", "\\1~", names(out))

  # No removal of "lower cutoff" and "upper cutoff" - they remain in place
  return(out)
}


#' @include set_coef.R
#' @rdname set_coef
#' @export
set_coef.glmmTMB <- function(model, coefs, ...) {
  # use predict(`newparams`) for this kind of model
  return(model)
}

#' @rdname sanitize_model_specific
sanitize_model_specific.glmmTMB <- function(model, vcov = TRUE, re.form, ...) {
  # re.form=NA
  if (!isTRUE(checkmate::check_flag(vcov))) {
    msg <- "For this model type, `vcov` must be `TRUE` or `FALSE`."
    insight::format_error(msg)
  }

  unsafe <- isFALSE(getOption("marginaleffects_safe", default = TRUE))
  if (unsafe && (missing(re.form) || (!isTRUE(is.na(re.form))))) {
    msg <- "For this model type, `marginaleffects` only takes into account the uncertainty in fixed-effect parameters. This is often appropriate when `re.form=NA`, but may be surprising to users who set `re.form=NULL` (default) or to some other value. Call `options(marginaleffects_safe = FALSE)` to silence this warning."
    insight::format_warning(msg)
  }

  return(model)
}
