#' Chat with a model hosted on Databricks
#'
#' @description
#' Databricks provides out-of-the-box access to a number of [foundation
#' models](https://docs.databricks.com/en/machine-learning/model-serving/score-foundation-models.html)
#' and can also serve as a gateway for external models hosted by a third party.
#'
#' Built on top of [chat_openai_compatible()].
#'
#' ## Authentication
#'
#' `chat_databricks()` picks up on ambient Databricks credentials for a subset
#' of the [Databricks client unified
#' authentication](https://docs.databricks.com/en/dev-tools/auth/unified-auth.html)
#' model. Specifically, it supports:
#'
#' - Personal access tokens
#' - Service principals via OAuth (OAuth M2M)
#' - User account via OAuth (OAuth U2M)
#' - Authentication via the Databricks CLI
#' - Posit Workbench-managed credentials
#' - Viewer-based credentials on Posit Connect. Requires the \pkg{connectcreds}
#'   package.
#'
#' @family chatbots
#' @param workspace The URL of a Databricks workspace, e.g.
#'   `"https://example.cloud.databricks.com"`. Will use the value of the
#'   environment variable `DATABRICKS_HOST`, if set.
#' @param model `r param_model("databricks-claude-3-7-sonnet")`
#'
#'   Available foundational models include:
#'
#'   - `databricks-claude-3-7-sonnet` (the default)
#'   - `databricks-mixtral-8x7b-instruct`
#'   - `databricks-meta-llama-3-1-70b-instruct`
#'   - `databricks-meta-llama-3-1-405b-instruct`
#' @param token An authentication token for the Databricks workspace, or
#'   `NULL` to use ambient credentials.
#' @param params Common model parameters, usually created by [params()].
#' @inheritParams chat_openai
#' @inherit chat_openai return
#' @export
#' @examples
#' \dontrun{
#' chat <- chat_databricks()
#' chat$chat("Tell me three jokes about statisticians")
#' }
chat_databricks <- function(
  workspace = databricks_workspace(),
  system_prompt = NULL,
  model = NULL,
  token = NULL,
  params = NULL,
  api_args = list(),
  echo = c("none", "output", "all"),
  api_headers = character()
) {
  check_string(workspace, allow_empty = FALSE)
  check_string(token, allow_empty = FALSE, allow_null = TRUE)
  model <- set_default(model, "databricks-claude-3-7-sonnet")
  echo <- check_echo(echo)
  if (!is.null(token)) {
    credentials <- function() list(Authorization = paste("Bearer", token))
  } else {
    credentials <- default_databricks_credentials(workspace)
  }

  params <- params %||% params()

  provider <- ProviderDatabricks(
    name = "Databricks",
    base_url = workspace,
    model = model,
    params = params,
    extra_args = api_args,
    credentials = credentials,
    extra_headers = api_headers
  )
  Chat$new(provider = provider, system_prompt = system_prompt, echo = echo)
}

ProviderDatabricks <- new_class(
  "ProviderDatabricks",
  parent = ProviderOpenAICompatible
)

method(base_request, ProviderDatabricks) <- function(provider) {
  req <- request(provider@base_url)
  req <- ellmer_req_credentials(req, provider@credentials())
  req <- ellmer_req_robustify(req)
  req <- ellmer_req_user_agent(req, databricks_user_agent())
  req <- base_request_error(provider, req)
  req
}

method(chat_body, ProviderDatabricks) <- function(
  provider,
  stream = TRUE,
  turns = list(),
  tools = list(),
  type = NULL
) {
  body <- chat_body(
    super(provider, ProviderOpenAICompatible),
    stream = stream,
    turns = turns,
    tools = tools,
    type = type
  )

  params <- chat_params(provider, provider@params)
  body <- modify_list(body, params)

  # Databricks doesn't support stream options
  body$stream_options <- NULL

  body
}

method(chat_params, ProviderDatabricks) <- function(provider, params) {
  # https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/api-reference#chat-request
  standardise_params(
    params,
    c(
      temperature = "temperature",
      top_p = "topP",
      top_k = "topK",
      max_tokens = "maxTokens",
      stop_sequences = "stopSequences"
    )
  )
}

method(chat_path, ProviderDatabricks) <- function(provider) {
  # Note: this API endpoint is undocumented and seems to exist primarily for
  # compatibility with the OpenAI Python SDK. The documented endpoint is
  # `/serving-endpoints/<model>/invocations`.
  "/serving-endpoints/chat/completions"
}

method(base_request_error, ProviderDatabricks) <- function(provider, req) {
  req_error(req, body = function(resp) {
    if (resp_content_type(resp) == "application/json") {
      # Databrick's "OpenAI-compatible" API has a slightly incompatible error
      # response format, which we account for here.
      resp_body_json(resp)$message
    }
  })
}

# See: https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/api-reference#functionobject
method(as_json, list(ProviderDatabricks, ToolDef)) <- function(
  provider,
  x,
  ...
) {
  # Note: It seems that Databricks doesn't support the "strict" field, despite
  # what their documentation says. It *is* supported for structured outputs,
  # though. I suspect a copy & paste error in their docs.
  compact(list(
    type = "function",
    "function" = compact(list(
      name = x@name,
      description = x@description,
      # Use the same parameter encoding as the OpenAI provider, but only if
      # there actually are parameters.
      parameters = if (length(x@arguments@properties) != 0) {
        as_json(provider, x@arguments, ...)
      }
    ))
  ))
}

# https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/api-reference#toolcall
method(as_json, list(ProviderDatabricks, ContentToolRequest)) <- function(
  provider,
  x,
  ...
) {
  # Databricks seems to require encoding empty arguments as an empty
  # dictionary, rather than an empty array.
  json_args <- to_json(set_names(x@arguments))
  list(
    id = x@id,
    `function` = list(name = x@name, arguments = json_args),
    type = "function"
  )
}

databricks_workspace <- function(error_call = caller_env()) {
  # Default to the well-known DATABRICKS_HOST environment variable.
  host <- Sys.getenv("DATABRICKS_HOST")
  if (nchar(host) != 0) {
    if (!grepl("^https?://", host)) {
      host <- paste0("https://", host)
    }
    return(host)
  }
  # If we have no environment variable, check for a Databricks config file
  # (likely generated via the Databricks CLI or their VS Code extension).
  cfg_file <- Sys.getenv("DATABRICKS_CONFIG_FILE", "~/.databrickscfg")
  if (file.exists(cfg_file)) {
    profile <- Sys.getenv("DATABRICKS_CONFIG_PROFILE", "DEFAULT")
    host <- databricks_config_field(cfg_file, profile, "host")
    if (!is.null(host)) {
      return(host)
    }
  }
  if (is_testing()) {
    testthat::skip("no Databricks workspace configured")
  }
  cli::cli_abort(
    "No env var {.code DATABRICKS_HOST} set or valid {.file {cfg_file}} found.",
    call = error_call
  )
}

databricks_user_agent <- function() {
  user_agent <- paste0("r-ellmer/", utils::packageVersion("ellmer"))
  if (nchar(Sys.getenv("SPARK_CONNECT_USER_AGENT")) != 0) {
    user_agent <- paste(Sys.getenv("SPARK_CONNECT_USER_AGENT"), user_agent)
  }
  user_agent
}

# Try various ways to get Databricks credentials. This implements a subset of
# the "Databricks client unified authentication" model.
default_databricks_credentials <- function(workspace = databricks_workspace()) {
  host <- gsub("https://|/$", "", workspace)

  # Detect viewer-based credentials from Posit Connect.
  if (has_connect_viewer_token(resource = workspace)) {
    return(function() {
      token <- connectcreds::connect_viewer_token(workspace)
      list(Authorization = paste("Bearer", token$access_token))
    })
  }

  # An explicit PAT takes precedence over everything else.
  token <- Sys.getenv("DATABRICKS_TOKEN")
  if (nchar(token)) {
    return(function() list(Authorization = paste("Bearer", token)))
  }

  # Next up are explicit OAuth2 M2M credentials.
  client_id <- Sys.getenv("DATABRICKS_CLIENT_ID")
  client_secret <- Sys.getenv("DATABRICKS_CLIENT_SECRET")
  if (nchar(client_id) && nchar(client_secret)) {
    # M2M credentials use an OAuth client credentials flow. We cache the token
    # so we don't need to perform this flow before each turn.
    client <- oauth_client(
      client_id,
      paste0("https://", host, "/oidc/v1/token"),
      secret = client_secret,
      auth = "header",
      name = "ellmer-databricks-m2m"
    )
    return(function() {
      token <- oauth_token_cached(
        client,
        oauth_flow_client_credentials,
        # The "all-apis" scope translates to "everything this service principal
        # has access to", not "all Databricks APIs".
        flow_params = list(scope = "all-apis"),
        # Don't use the cached token when testing.
        reauth = is_testing()
      )
      list(Authorization = paste("Bearer", token$access_token))
    })
  }

  # Check for Workbench-provided credentials.
  cfg_file <- Sys.getenv("DATABRICKS_CONFIG_FILE")
  if (grepl("posit-workbench", cfg_file, fixed = TRUE)) {
    wb_token <- workbench_databricks_token(host, cfg_file)
    if (!is.null(wb_token)) {
      return(function() {
        # Ensure we get an up-to-date token.
        token <- workbench_databricks_token(host, cfg_file)
        list(Authorization = paste("Bearer", token))
      })
    }
  }

  # When on desktop, try using the Databricks CLI for auth.
  cli_path <- Sys.getenv("DATABRICKS_CLI_PATH", "databricks")
  if (!is_hosted_session() && nchar(Sys.which(cli_path)) != 0) {
    token <- databricks_cli_token(cli_path, workspace)
    if (!is.null(token)) {
      return(function() {
        # Ensure we get an up-to-date token.
        token <- databricks_cli_token(cli_path, workspace)
        list(Authorization = paste("Bearer", token))
      })
    }
  }

  # Fall back to OAuth U2M, masquerading as the Databricks CLI. Again, this
  # only works on desktop.
  if (is_interactive() && !is_hosted_session()) {
    # U2M credentials use an OAuth authorization code flow. We cache the token
    # so we don't need to perform this flow before each turn.
    client <- oauth_client(
      "databricks-cli",
      paste0("https://", host, "/oidc/v1/token"),
      auth = "body",
      name = "ellmer-databricks-u2m"
    )
    return(function() {
      token <- oauth_token_cached(
        client,
        oauth_flow_auth_code,
        flow_params = list(
          auth_url = paste0("https://", host, "/oidc/v1/authorize"),
          # The "all-apis" scope translates to "everything this user has access
          # to", not "all Databricks APIs".
          scope = "all-apis offline_access",
          # This is the registered redirect URI for the Databricks CLI.
          redirect_uri = "http://localhost:8020"
        )
      )
      list(Authorization = paste("Bearer", token$access_token))
    })
  }

  if (is_testing()) {
    testthat::skip("no Databricks credentials available")
  }

  cli::cli_abort("No Databricks credentials are available.")
}

# Try to determine whether we can redirect the user's browser to a server on
# localhost, which isn't possible if we are running on a hosted platform.
#
# This is based on the strategy pioneered by the {gargle} package and {httr2}.
is_hosted_session <- function() {
  # If RStudio Server or Posit Workbench is running locally (which is possible,
  # though unusual), it's not acting as a hosted environment.
  Sys.getenv("RSTUDIO_PROGRAM_MODE") == "server" &&
    !grepl("localhost", Sys.getenv("RSTUDIO_HTTP_REFERER"), fixed = TRUE)
}

databricks_cli_token <- function(cli_path, host) {
  output <- suppressWarnings(
    system2(
      cli_path,
      c("auth", "token", "--host", host),
      stdout = TRUE,
      stderr = TRUE
    )
  )
  output <- paste(output, collapse = "\n")
  # If we don't get an error message, try to extract the token from the JSON-
  # formatted output.
  if (grepl("access_token", output, fixed = TRUE)) {
    token <- gsub(".*access_token\":\\s?\"([^\"]+).*", "\\1", output)
    return(token)
  }
  NULL
}

# Reads Posit Workbench-managed Databricks credentials from a
# $DATABRICKS_CONFIG_FILE. The generated file will look as follows:
#
# [workbench]
# host = some-host
# token = some-token
workbench_databricks_token <- function(host, cfg_file) {
  databricks_config_field(cfg_file, "workbench", "token")
}

databricks_config_field <- function(cfg_file, profile, field) {
  cfg <- readLines(cfg_file)
  # We don't attempt a full parse of the INI syntax supported by Databricks
  # config files and use the following approach to cheat, instead:
  #
  # 1. Find the indices of the profile header and the field. If none appear in
  #    the file, we're done.
  # 2. Find the index of the next profile header, if any.
  # 3. Find the index -- if any -- of the field between these headers.
  # 4. Extract the field's value.
  profile <- which(grepl(sprintf("[%s]", profile), cfg, fixed = TRUE))
  if (length(profile) != 1) {
    return(NULL)
  }
  entry <- which(grepl(field, cfg, fixed = TRUE))
  entry <- entry[entry > profile]
  if (!any(entry)) {
    return(NULL)
  }
  headers <- which(startsWith(cfg, "["))
  next_profile <- headers[headers > profile]
  if (!any(next_profile)) {
    next_profile <- length(cfg) + 1
  } else {
    next_profile <- next_profile[1]
  }
  entry <- entry[entry < next_profile]
  if (length(entry) != 1) {
    return(NULL)
  }
  gsub(sprintf("%s\\s*=\\s*(.*)", field), "\\1", cfg[entry])
}
