#' @include provider-openai-compatible.R
#' @include content.R
NULL

#' Chat with a model hosted on Snowflake
#'
#' @description
#' The Snowflake provider allows you to interact with LLM models available
#' through the [Cortex LLM REST API](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api).
#'
#' ## Authentication
#'
#' `chat_snowflake()` picks up the following ambient Snowflake credentials:
#'
#' - A static OAuth token defined via the `SNOWFLAKE_TOKEN` environment
#'   variable.
#' - Key-pair authentication credentials defined via the `SNOWFLAKE_USER` and
#'   `SNOWFLAKE_PRIVATE_KEY` (which can be a PEM-encoded private key or a path
#'   to one) environment variables.
#' - Posit Workbench-managed Snowflake credentials for the corresponding
#'   `account`.
#' - Viewer-based credentials on Posit Connect. Requires the \pkg{connectcreds}
#'   package.
#'
#' ## Known limitations
#' Note that Snowflake-hosted models do not support images.
#'
#' @inheritParams chat_openai
#' @param account A Snowflake [account identifier](https://docs.snowflake.com/en/user-guide/admin-account-identifier),
#'   e.g. `"testorg-test_account"`. Defaults to the value of the
#'   `SNOWFLAKE_ACCOUNT` environment variable.
#' @param credentials A list of authentication headers to pass into
#'   [`httr2::req_headers()`], a function that returns them when called, or
#'   `NULL`, the default, to use ambient credentials.
#' @param model `r param_model("claude-3-7-sonnet")`
#' @inherit chat_openai return
#' @examplesIf has_credentials("snowflake")
#' chat <- chat_snowflake()
#' chat$chat("Tell me a joke in the form of a SQL query.")
#' @export
chat_snowflake <- function(
  system_prompt = NULL,
  account = snowflake_account(),
  credentials = NULL,
  model = NULL,
  params = NULL,
  api_args = list(),
  echo = c("none", "output", "all"),
  api_headers = character()
) {
  check_string(account, allow_empty = FALSE)
  model <- set_default(model, "claude-3-7-sonnet")
  params <- params %||% params()
  echo <- check_echo(echo)

  credentials <- credentials %||% default_snowflake_credentials(account)
  check_credentials(credentials)

  provider <- ProviderSnowflakeCortex(
    name = "Snowflake/Cortex",
    base_url = snowflake_url(account),
    account = account,
    credentials = credentials,
    model = model,
    params = params,
    extra_args = api_args,
    extra_headers = api_headers
  )

  Chat$new(provider = provider, system_prompt = system_prompt, echo = echo)
}

ProviderSnowflakeCortex <- new_class(
  "ProviderSnowflakeCortex",
  parent = ProviderOpenAICompatible,
  properties = list(
    account = prop_string()
  )
)

method(base_request, ProviderSnowflakeCortex) <- function(provider) {
  req <- request(provider@base_url)
  req <- ellmer_req_credentials(req, provider@credentials())
  req <- ellmer_req_robustify(req)
  # Snowflake uses the User Agent header to identify "parter applications", so
  # identify requests as coming from "r_ellmer" (unless an explicit partner
  # application is set via the ambient SF_PARTNER environment variable).
  req <- ellmer_req_user_agent(req, Sys.getenv("SF_PARTNER"))

  # Snowflake-specific error response handling:
  req <- req_error(req, body = function(resp) resp_body_json(resp)$message)

  req
}

method(chat_path, ProviderSnowflakeCortex) <- function(provider) {
  "/api/v2/cortex/inference:complete"
}

# See: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#api-reference
method(chat_body, ProviderSnowflakeCortex) <- function(
  provider,
  stream = TRUE,
  turns = list(),
  tools = list(),
  type = NULL
) {
  messages <- as_json(provider, turns)
  tools <- as_json(provider, unname(tools))

  if (!is.null(type)) {
    # Note: Snowflake uses a slightly different format than OpenAI.
    response_format <- list(type = "json", schema = as_json(provider, type))
  } else {
    response_format <- NULL
  }

  params <- chat_params(provider, provider@params)
  compact(
    list2(
      messages = messages,
      model = provider@model,
      !!!params,
      stream = stream,
      tools = tools,
      response_format = response_format
    )
  )
}

method(as_json, list(ProviderSnowflakeCortex, TypeObject)) <- function(
  provider,
  x,
  ...
) {
  # Unlike OpenAI, Snowflake does not support the "additionalProperties" field.
  names <- names2(x@properties)
  required <- map_lgl(x@properties, function(prop) prop@required)
  properties <- as_json(provider, x@properties, ...)
  names(properties) <- names
  list(
    type = "object",
    description = x@description %||% "",
    properties = properties,
    required = as.list(names[required])
  )
}

# See: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#optional-json-arguments
method(chat_params, ProviderSnowflakeCortex) <- function(provider, params) {
  standardise_params(
    params,
    c(
      temperature = "temperature",
      top_p = "top_p",
      max_tokens = "max_tokens"
    )
  )
}

# Snowflake -> ellmer --------------------------------------------------------

method(stream_content, ProviderSnowflakeCortex) <- function(provider, event) {
  if (length(event$choices) == 0) {
    return(NULL)
  }
  delta <- event$choices[[1]]$delta
  if (is.null(delta) || !identical(delta$type, "text")) {
    return(NULL)
  }
  text <- delta[["content"]] %||% delta[["text"]]
  if (is.null(text) || !nzchar(text)) {
    return(NULL)
  }
  ContentText(text)
}

method(stream_merge_chunks, ProviderSnowflakeCortex) <- function(
  provider,
  result,
  chunk
) {
  delta <- chunk$choices[[1]]$delta

  if (is.null(result)) {
    if (delta$type == "tool_use") {
      content_list <- list(
        list(
          type = "tool_use",
          tool_use_id = delta$tool_use_id,
          name = delta$name,
          input = delta$input %||% ""
        )
      )
    } else {
      content_list <- list(
        list(
          type = "text",
          text = delta$text %||% ""
        )
      )
    }
    # Non-streaming responses use "message" instead of "delta".
    chunk$choices[[1]]$message <- list(content_list = content_list)
    chunk$choices[[1]]$delta <- NULL
    return(chunk)
  }

  content_list <- result$choices[[1]]$message$content_list

  if (delta$type == "text") {
    text_idx <- NULL
    for (i in rev(seq_along(content_list))) {
      if (identical(content_list[[i]]$type, "text")) {
        text_idx <- i
        break
      }
    }
    if (is.null(text_idx)) {
      content_list[[length(content_list) + 1L]] <- list(
        type = "text",
        text = delta$text %||% ""
      )
    } else {
      paste(content_list[[text_idx]]$text) <- delta$text %||% ""
    }
  } else if (delta$type == "tool_use") {
    if (!is.null(delta$tool_use_id)) {
      content_list[[length(content_list) + 1L]] <- list(
        type = "tool_use",
        tool_use_id = delta$tool_use_id,
        name = delta$name,
        input = delta$input %||% ""
      )
    } else if (!is.null(delta$input)) {
      for (i in rev(seq_along(content_list))) {
        if (identical(content_list[[i]]$type, "tool_use")) {
          paste(content_list[[i]]$input) <- delta$input
          break
        }
      }
    }
  } else {
    cli::cli_abort(
      "Unsupported content type {.str {delta$type}}.",
      .internal = TRUE
    )
  }

  result$choices[[1]]$message$content_list <- content_list
  result$usage <- chunk$usage
  result
}

method(value_tokens, ProviderSnowflakeCortex) <- function(provider, json) {
  usage <- json$usage
  tokens(
    input = usage$prompt_tokens %||% 0,
    output = usage$completion_tokens %||% 0
  )
}

method(value_turn, ProviderSnowflakeCortex) <- function(
  provider,
  result,
  has_type = FALSE
) {
  raw_content <- result$choices[[1]]$message$content_list
  contents <- compact(
    lapply(raw_content, function(content) {
      if (identical(content$type, "text")) {
        if (!nzchar(content$text %||% "")) {
          return(NULL)
        }
        if (has_type) {
          ContentJson(string = content$text)
        } else {
          ContentText(content$text)
        }
      } else if (identical(content$type, "tool_use")) {
        # Streaming produces flat format; non-streaming has nested tool_use
        if (!is.null(content[["tool_use"]])) {
          content <- content[["tool_use"]]
        }
        input <- content[["input"]] %||% ""
        id <- content[["tool_use_id"]]
        name <- content[["name"]]
        if (is_string(input)) {
          input <- jsonlite::parse_json(input)
        }
        ContentToolRequest(id, name, input %||% list())
      } else {
        cli::cli_abort(
          "Unknown content type {.str {content$type}}.",
          .internal = TRUE
        )
      }
    })
  )
  tokens <- value_tokens(provider, result)
  cost <- get_token_cost(provider, tokens)
  AssistantTurn(contents, json = result, tokens = unlist(tokens), cost = cost)
}

# ellmer -> Snowflake --------------------------------------------------------

method(as_json, list(ProviderSnowflakeCortex, Turn)) <- function(
  provider,
  x,
  ...
) {
  # Attempting to omit the `content` field and use `content_list` instead
  # yields:
  #
  # > messages[0].content cannot be empty string (was '')
  #
  # So we emulate what Snowflake do and put the text content in both.
  if (S7_inherits(x@contents[[1]], ContentText)) {
    content <- x@contents[[1]]@text
    if (nchar(content) == 0) {
      content <- "<empty>"
    }
  } else if (S7_inherits(x@contents[[1]], ContentJson)) {
    # Match the existing as_json() implementation.
    content <- "<structured data/>"
  } else if (S7_inherits(x@contents[[1]], ContentToolResult)) {
    # Completely undocumented, but: it seems like the model is expecting the
    # tool result in textual format here, too -- otherwise it gets confused,
    # like it can't see the output.
    content <- tool_string(x@contents[[1]])
  } else if (S7_inherits(x@contents[[1]], ContentToolRequest)) {
    # Tool-only response (no preceding text).
    content <- "<empty>"
  } else {
    cli::cli_abort("Unsupported content type: {.cls {class(x@contents[[1]])}}.")
  }
  x <- turn_contents_expand(x)
  list(
    role = x@role,
    content_list = as_json(provider, x@contents, ...)
  )
}

# See: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#tools-configuration
method(as_json, list(ProviderSnowflakeCortex, ToolDef)) <- function(
  provider,
  x,
  ...
) {
  list(
    tool_spec = compact(
      list(
        type = "generic",
        name = x@name,
        description = x@description,
        input_schema = as_json(provider, x@arguments, ...)
      )
    )
  )
}

# See: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#tools-configuration
method(as_json, list(ProviderSnowflakeCortex, ContentToolRequest)) <- function(
  provider,
  x,
  ...
) {
  input <- x@arguments
  if (length(input) == 0) {
    # Snowflake requires an empty object, rather than an empty array.
    input <- set_names(list())
  }
  list(
    type = "tool_use",
    tool_use = list(
      tool_use_id = x@id,
      name = x@name,
      input = input
    )
  )
}

# See: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api#tool-results
method(as_json, list(ProviderSnowflakeCortex, ContentToolResult)) <- function(
  provider,
  x,
  ...
) {
  list(
    type = "tool_results",
    tool_results = compact(
      list(
        tool_use_id = x@request@id,
        name = x@request@name,
        content = list(
          list(type = "text", text = tool_string(x))
        ),
        # TODO: Is this the correct format?
        status = if (tool_errored(x)) "error"
      )
    )
  )
}

# Utilities ------------------------------------------------------------------

snowflake_account <- function() {
  key_get("SNOWFLAKE_ACCOUNT")
}

snowflake_url <- function(account) {
  paste0("https://", account, ".snowflakecomputing.com")
}

default_snowflake_credentials <- function(account = snowflake_account()) {
  # Detect viewer-based credentials from Posit Connect.
  url <- snowflake_url(account)
  if (is_installed("connectcreds") && connectcreds::has_viewer_token(url)) {
    return(function() {
      token <- connectcreds::connect_viewer_token(url)
      list(
        Authorization = paste("Bearer", token$access_token),
        `X-Snowflake-Authorization-Token-Type` = "OAUTH"
      )
    })
  }

  token <- Sys.getenv("SNOWFLAKE_TOKEN")
  if (nchar(token) != 0) {
    return(function() {
      list(
        Authorization = paste("Bearer", token),
        # See: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication#using-oauth
        `X-Snowflake-Authorization-Token-Type` = "OAUTH"
      )
    })
  }

  # Support for Snowflake key-pair authentication.
  # See: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication#generate-a-jwt-token
  user <- Sys.getenv("SNOWFLAKE_USER")
  private_key <- Sys.getenv("SNOWFLAKE_PRIVATE_KEY")
  if (nchar(user) != 0 && nchar(private_key) != 0) {
    check_installed("jose", "for key-pair authentication")
    key <- openssl::read_key(private_key)
    return(function() {
      token <- snowflake_keypair_token(account, user, key)
      list(
        Authorization = paste("Bearer", token),
        `X-Snowflake-Authorization-Token-Type` = "KEYPAIR_JWT"
      )
    })
  }

  # Check for Workbench-managed credentials.
  sf_home <- Sys.getenv("SNOWFLAKE_HOME")
  if (grepl("posit-workbench", sf_home, fixed = TRUE)) {
    token <- workbench_snowflake_token(account, sf_home)
    if (!is.null(token)) {
      return(function() {
        # Ensure we get an up-to-date token.
        token <- workbench_snowflake_token(account, sf_home)
        list(
          Authorization = paste("Bearer", token),
          `X-Snowflake-Authorization-Token-Type` = "OAUTH"
        )
      })
    }
  }

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

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

snowflake_keypair_token <- function(
  account,
  user,
  key,
  cache = snowflake_keypair_cache(account, key),
  lifetime = 600L,
  reauth = FALSE
) {
  # Producing a signed JWT is a fairly expensive operation (in the order of
  # ~10ms), but adding a cache speeds this up approximately 500x.
  creds <- cache$get()
  if (reauth || is.null(creds) || creds$expiry < Sys.time()) {
    cache$clear()
    expiry <- Sys.time() + lifetime
    # We can't use openssl::fingerprint() here because it uses a different
    # algorithm.
    fp <- openssl::base64_encode(
      openssl::sha256(openssl::write_der(key$pubkey))
    )
    if (grepl(".+\\.privatelink$", account)) {
      # account identifier is everything up to the first period
      account <- gsub("^([^.]*).+", "\\1", account)
    }
    sub <- toupper(paste0(account, ".", user))
    iss <- paste0(sub, ".SHA256:", fp)
    # Note: Snowflake employs a malformed issuer claim, so we have to inject it
    # manually after jose's validation phase.
    claim <- jwt_claim("dummy", sub, exp = as.integer(expiry))
    claim$iss <- iss
    creds <- list(expiry = expiry, token = jwt_encode_sig(claim, key))
    cache$set(creds)
  }
  creds$token
}

snowflake_keypair_cache <- function(account, key) {
  credentials_cache(key = hash(c("sf", account, openssl::fingerprint(key))))
}

# Credential handling ----------------------------------------------------------

snowflake_credentials_exist <- function(...) {
  tryCatch(
    is_list(default_snowflake_credentials(...)),
    error = function(e) FALSE
  )
}

# Reads Posit Workbench-managed Snowflake credentials from a
# $SNOWFLAKE_HOME/connections.toml file, as used by the Snowflake Connector for
# Python implementation. The file will look as follows:
#
# [workbench]
# account = "account-id"
# token = "token"
# authenticator = "oauth"
workbench_snowflake_token <- function(account, sf_home) {
  cfg <- readLines(file.path(sf_home, "connections.toml"))
  # We don't attempt a full parse of the TOML syntax, instead relying on the
  # fact that this file will always contain only one section.
  if (!any(grepl(account, cfg, fixed = TRUE))) {
    # The configuration doesn't actually apply to this account.
    return(NULL)
  }
  line <- grepl("token = ", cfg, fixed = TRUE)
  token <- gsub("token = ", "", cfg[line])
  if (nchar(token) == 0) {
    return(NULL)
  }
  # Drop enclosing quotes.
  gsub("\"", "", token)
}
