#' @include provider.R
#' @include content.R
#' @include turns.R
#' @include tools-def.R
NULL

#' Chat with an AWS bedrock model
#'
#' @description
#' `r support_badge("official")`
#'
#' [AWS Bedrock](https://aws.amazon.com/bedrock/) provides a number of
#' language models, including those from Anthropic's
#' [Claude](https://aws.amazon.com/bedrock/claude/), using the Bedrock
#' [Converse API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html).
#'
#' ## Authentication
#'
#' `chat_aws_bedrock()` uses \{paws.common\} to resolve credentials,
#' trying the following strategies in order:
#'
#' - A bearer token set in the `AWS_BEARER_TOKEN_BEDROCK` or
#'   `AWS_BEARER_TOKEN` environment variable. This is used by enterprise
#'   API gateways that issue API keys instead of IAM credentials. See the
#'   [AWS documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html)
#'   for details.
#' - Standard IAM credentials resolved from environment variables, AWS
#'   config files, SSO, or instance metadata. See
#'   <https://www.paws-r-sdk.com/#credentials> for details. If your org
#'   uses AWS SSO, you'll need to run `aws sso login` at the terminal.
#'
#' ## Prompt caching
#'
#' Bedrock supports
#' [prompt caching](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html)
#' via cache checkpoints. When caching is enabled, ellmer places cache
#' checkpoints on the system prompt and the last turn, so that the
#' conversation history is cached across turns.
#'
#' By default (`cache = "auto"`), caching is enabled for models known to
#' support it (Anthropic Claude and Amazon Nova) and disabled for all other
#' models. You can also set `cache` to `"5m"` or `"1h"` to force a specific
#' TTL, or `"none"` to disable caching entirely. Note that individual models
#' may have minimum input token thresholds before caching takes effect.
#'
#' Note that [token_usage()] does not currently reflect the cost of writing
#' to the cache, which is priced at a premium over regular input tokens.
#' Cache read savings are reported correctly.
#'
#' @param profile AWS profile to use.
#' @param cache How long to cache inputs? The default, `"auto"`, enables
#'   caching with a 5-minute TTL for models known to support it (Anthropic
#'   Claude and Amazon Nova) and disables caching for all other models.
#'   Set to `"5m"` or `"1h"` to force caching on, or `"none"` to disable it.
#'
#'   See details below.
#' @param model `r param_model("us.anthropic.claude-sonnet-5", "models_aws_bedrock")`.
#'
#'   While ellmer provides a default model, there's no guarantee that you'll
#'   have access to it, so you'll need to specify a model that you can.
#'   If you're using [cross-region inference](https://aws.amazon.com/blogs/machine-learning/getting-started-with-cross-region-inference-in-amazon-bedrock/),
#'   you'll need to use the inference profile ID, e.g.
#'   `model="us.anthropic.claude-sonnet-5"`.
#' @param params Common model parameters, usually created by [params()].
#' @param api_args Named list of arbitrary extra arguments appended to the body
#'   of every chat API call. Use `params` for common parameters. Model-specific
#'   inference parameters can be provided using the
#'   `additionalModelRequestFields` field, for example to enable thinking effort
#'   in Anthropic Claude models:
#'
#'   ```R
#'   api_args = list(
#'     additionalModelRequestFields = list(
#'       thinking = list(type = "enabled", budget_tokens = 4000)
#'     )
#'   )
#'   ```
#'
#'   See <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html>
#'   for more details.
#' @inheritParams chat_openai
#' @inherit chat_openai return
#' @family chatbots
#' @export
#' @examples
#' \dontrun{
#' # Basic usage
#' chat <- chat_aws_bedrock()
#' chat$chat("Tell me three jokes about statisticians")
#' }
chat_aws_bedrock <- function(
  system_prompt = NULL,
  base_url = NULL,
  model = NULL,
  profile = NULL,
  cache = c("auto", "5m", "1h", "none"),
  params = NULL,
  api_args = list(),
  api_headers = character(),
  echo = NULL
) {
  check_installed("paws.common", "AWS authentication")
  check_string(base_url, allow_null = TRUE)
  base_url <- base_url %||%
    \(x) sprintf("https://bedrock-runtime.%s.amazonaws.com", x)
  echo <- check_echo(echo)

  params <- params %||% params()

  provider <- provider_aws_bedrock(
    base_url = base_url,
    model = model,
    profile = profile,
    cache_point = cache,
    params = params,
    extra_args = api_args,
    extra_headers = api_headers
  )
  Chat$new(provider = provider, system_prompt = system_prompt, echo = echo)
}


#' @export
#' @rdname chat_aws_bedrock
models_aws_bedrock <- function(profile = NULL, base_url = NULL) {
  check_string(base_url, allow_null = TRUE)
  base_url <- base_url %||% \(x) sprintf("https://bedrock.%s.amazonaws.com", x)

  provider <- provider_aws_bedrock(
    base_url = base_url,
    model = "",
    profile = profile,
  )
  models_list(provider)
}

chat_aws_bedrock_test <- function(
  ...,
  model = "us.anthropic.claude-haiku-4-5-20251001-v1:0",
  params = NULL,
  echo = "none"
) {
  params <- params %||% params()
  params$temperature <- params$temperature %||% 0

  chat_aws_bedrock(model = model, params = params, ..., echo = echo)
}

provider_aws_bedrock <- function(
  base_url,
  model = "",
  profile = NULL,
  cache_point = "none",
  params = list(),
  extra_args = list(),
  extra_headers = character()
) {
  cache <- aws_creds_cache(profile)
  credentials <- paws_credentials(profile, cache = cache)

  if (is.function(base_url)) {
    base_url <- base_url(credentials$region)
  }

  model <- set_default(model, "us.anthropic.claude-sonnet-5")

  cache_point <- as_bedrock_cache_point(cache_point, model)

  ProviderAWSBedrock(
    name = "AWS/Bedrock",
    base_url = base_url,
    model = model,
    profile = profile,
    region = credentials$region,
    cache = cache,
    cache_point = cache_point,
    params = params,
    extra_args = extra_args,
    extra_headers = extra_headers
  )
}

ProviderAWSBedrock <- new_class(
  "ProviderAWSBedrock",
  parent = Provider,
  properties = list(
    profile = prop_string(allow_null = TRUE),
    region = prop_string(),
    cache = class_list,
    cache_point = prop_string()
  )
)

method(models_list, ProviderAWSBedrock) <- function(provider) {
  # ListFoundationModels uses the control-plane endpoint (bedrock.*) not the
  # data-plane endpoint (bedrock-runtime.*) used for inference.
  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_ListFoundationModels.html
  provider@base_url <- sub(
    "bedrock-runtime",
    "bedrock",
    provider@base_url,
    fixed = TRUE
  )

  req <- base_request(provider)
  req <- req_url_path_append(req, "foundation-models")
  resp <- req_perform(req)
  json <- resp_body_json(resp)
  models <- json$modelSummaries

  df <- data.frame(
    id = map_chr(models, "[[", "modelId"),
    name = map_chr(models, "[[", "modelName"),
    provider = map_chr(models, "[[", "providerName")
  )
  cbind(df, match_prices("AWS/Bedrock", df$id))
}

method(base_request, ProviderAWSBedrock) <- function(provider) {
  creds <- paws_credentials(provider@profile, provider@cache)

  req <- request(provider@base_url)
  if (nzchar(creds$access_token)) {
    req <- req_auth_bearer_token(req, creds$access_token)
  } else {
    req <- req_auth_aws_v4(
      req,
      aws_access_key_id = creds$access_key_id,
      aws_secret_access_key = creds$secret_access_key,
      aws_session_token = creds$session_token
    )
  }
  req <- ellmer_req_robustify(req)
  req <- ellmer_req_user_agent(req)
  req <- base_request_error(provider, req)
  req
}

method(base_request_error, ProviderAWSBedrock) <- function(provider, req) {
  req_error(req, body = function(resp) {
    body <- resp_body_json(resp)
    body$Message %||% body$message
  })
}

method(chat_params, ProviderAWSBedrock) <- function(provider, params) {
  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html
  standardise_params(
    params,
    c(
      temperature = "temperature",
      topP = "top_p",
      maxTokens = "max_tokens",
      stopSequences = "stop_sequences"
    )
  )
}

method(chat_request, ProviderAWSBedrock) <- function(
  provider,
  stream = TRUE,
  turns = list(),
  tools = list(),
  type = NULL
) {
  req <- base_request(provider)
  suffix <- if (stream) "converse-stream" else "converse"
  req <- req_url_path_append(
    req,
    paste0("model/", curl::curl_escape(provider@model), "/", suffix)
  )

  if (length(turns) >= 1 && is_system_turn(turns[[1]])) {
    system <- c(
      list(list(text = turns[[1]]@text)),
      bedrock_cache_point(provider)
    )
  } else {
    system <- NULL
  }

  is_last <- seq_along(turns) == length(turns)
  messages <- compact(map2(turns, is_last, function(turn, is_last) {
    as_json(provider, turn, is_last = is_last)
  }))

  if (!is.null(type)) {
    tool_def <- ToolDef(
      function(...) {},
      name = "structured_tool_call__",
      description = "Extract structured data",
      arguments = type_object(data = type)
    )
    tools[[tool_def@name]] <- tool_def
    tool_choice <- list(tool = list(name = tool_def@name))
  } else {
    tool_choice <- NULL
  }

  if (length(tools) > 0) {
    tools <- chat_body_tools(provider, tools)
    toolConfig <- compact(list(tools = tools, tool_choice = tool_choice))
  } else {
    toolConfig <- NULL
  }

  # Merge params into inferenceConfig, giving precedence to manual api_args
  params <- chat_params(provider, provider@params)

  extra_args <- provider@extra_args
  extra_args$inferenceConfig <- modify_list(params, extra_args$inferenceConfig)

  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
  body <- compact(list2(
    messages = messages,
    system = system,
    toolConfig = toolConfig,
    !!!extra_args
  ))

  req <- req_body_json(req, body)
  req <- req_headers(req, !!!provider@extra_headers)

  req
}

method(chat_resp_stream, ProviderAWSBedrock) <- function(provider, resp) {
  resp_stream_aws(resp)
}

# Bedrock -> ellmer -------------------------------------------------------------

method(stream_parse, ProviderAWSBedrock) <- function(provider, event) {
  if (is.null(event)) {
    return()
  }

  body <- event$body
  body$event_type <- event$headers$`:event-type`
  body$p <- NULL # padding? Looks like: "p": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJ",

  body
}

method(stream_content, ProviderAWSBedrock) <- function(provider, event) {
  if (event$event_type == "contentBlockDelta") {
    text <- event$delta$text
    if (is.null(text)) {
      return(NULL)
    }
    ContentText(text)
  }
}

method(stream_merge_chunks, ProviderAWSBedrock) <- function(
  provider,
  result,
  chunk
) {
  i <- chunk$contentBlockIndex + 1

  if (chunk$event_type == "messageStart") {
    result <- list(role = chunk$role, content = list())
  } else if (chunk$event_type == "contentBlockStart") {
    result$content[[i]] <- list(toolUse = chunk$start$toolUse)
  } else if (chunk$event_type == "contentBlockDelta") {
    if (i > length(result$content)) {
      result$content[[i]] <- list()
    }
    if (has_name(chunk$delta, "text")) {
      paste(result$content[[i]]$text) <- chunk$delta$text
    } else if (has_name(chunk$delta, "toolUse")) {
      paste(result$content[[i]]$toolUse$input) <- chunk$delta$toolUse$input
    } else if (has_name(chunk$delta, "reasoningContent")) {
      if (is.null(result$content[[i]]$reasoningContent)) {
        result$content[[i]]$reasoningContent <- list(reasoningText = list())
      }
      # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ReasoningContentBlockDelta.html
      delta <- chunk$delta$reasoningContent
      if (has_name(delta, "text")) {
        paste(result$content[[i]]$reasoningContent$reasoningText$text) <-
          delta$text
      } else if (has_name(delta, "signature")) {
        result$content[[i]]$reasoningContent$reasoningText$signature <-
          delta$signature
      }
    } else {
      cli::cli_abort(
        "Unknown chunk type {names(chunk$delta)}",
        .internal = TRUE
      )
    }
  } else if (chunk$event_type == "contentBlockStop") {
    if (has_name(result$content[[i]], "toolUse")) {
      input <- result$content[[i]]$toolUse$input
      if (input == "") {
        result$content[[i]]$toolUse$input <- set_names(list())
      } else {
        result$content[[i]]$toolUse$input <- jsonlite::parse_json(input)
      }
    }
  } else if (chunk$event_type == "messageStop") {
    # match structure of non-streaming
    result <- list(
      output = list(
        message = result
      ),
      stopReason = chunk$stopReason
    )
  } else if (chunk$event_type == "metadata") {
    result$usage <- chunk$usage
    result$metrics <- chunk$metrics
  } else {
    cli::cli_inform(c("!" = "Unknown chunk type {.str {event_type}}."))
  }

  result
}

method(value_tokens, ProviderAWSBedrock) <- function(provider, json) {
  usage <- json$usage
  tokens(
    input = usage$inputTokens %||% 0,
    output = usage$outputTokens %||% 0,
    cached_input = usage$cacheReadInputTokens %||% 0
  )
}

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
method(value_finish_reason, ProviderAWSBedrock) <- function(provider, result) {
  reason <- result$stopReason
  if (is.null(reason)) {
    return(NA_character_)
  }
  switch(
    reason,
    end_turn = "success",
    tool_use = "tool_use",
    max_tokens = "max_tokens",
    model_context_window_exceeded = "context_window",
    stop_sequence = "stop_sequence",
    guardrail_intervened = ,
    content_filtered = "content_filter",
    I(reason)
  )
}

method(value_turn, ProviderAWSBedrock) <- function(
  provider,
  result,
  has_type = FALSE
) {
  contents <- lapply(result$output$message$content, function(content) {
    if (has_name(content, "text")) {
      ContentText(content$text)
    } else if (has_name(content, "toolUse")) {
      if (has_type) {
        ContentJson(data = content$toolUse$input$data)
      } else {
        ContentToolRequest(
          name = content$toolUse$name,
          arguments = content$toolUse$input,
          id = content$toolUse$toolUseId
        )
      }
    } else if (has_name(content, "reasoningContent")) {
      ContentThinking(
        content$reasoningContent$reasoningText$text,
        extra = list(
          signature = content$reasoningContent$reasoningText$signature
        )
      )
    } else {
      cli::cli_abort(
        "Unknown content type {.str {names(content)}}.",
        .internal = TRUE
      )
    }
  })

  tokens <- value_tokens(provider, result)
  cost <- get_token_cost(provider, tokens)

  AssistantTurn(
    contents,
    json = result,
    tokens = unlist(tokens),
    cost = cost,
    finish_reason = value_finish_reason(provider, result)
  )
}

# ellmer -> Bedrock -------------------------------------------------------------

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html
method(as_json, list(ProviderAWSBedrock, Turn)) <- function(
  provider,
  x,
  ...,
  is_last = FALSE
) {
  if (is_system_turn(x)) {
    NULL
  } else if (is_user_turn(x) || is_assistant_turn(x)) {
    x <- turn_contents_expand(x)
    content <- as_json(provider, x@contents, ...)

    if (is_last) {
      content <- c(content, bedrock_cache_point(provider))
    }

    list(role = x@role, content = content)
  } else {
    cli::cli_abort("Unknown role {x@role}", .internal = TRUE)
  }
}

method(as_json, list(ProviderAWSBedrock, ContentText)) <- function(
  provider,
  x,
  ...
) {
  if (is_whitespace(x@text)) {
    list(text = "[empty string]")
  } else {
    list(text = x@text)
  }
}

method(as_json, list(ProviderAWSBedrock, ContentImageRemote)) <- function(
  provider,
  x,
  ...
) {
  cli::cli_abort("Bedrock doesn't support remote images")
}

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html
method(as_json, list(ProviderAWSBedrock, ContentImageInline)) <- function(
  provider,
  x,
  ...
) {
  type <- switch(
    x@type,
    "image/png" = "png",
    "image/gif" = "gif",
    "image/jpeg" = "jpeg",
    "image/webp" = "webp",
    cli::cli_abort("Image type {content@type} is not supported by bedrock")
  )

  list(
    image = list(
      format = type,
      source = list(bytes = x@data)
    )
  )
}

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_DocumentBlock.html
method(as_json, list(ProviderAWSBedrock, ContentPDF)) <- function(
  provider,
  x,
  ...
) {
  list(
    document = list(
      #> This field is vulnerable to prompt injections, because the model
      #> might inadvertently interpret it as instructions. Therefore, we
      #> that you specify a neutral name.
      name = bedrock_document_name(),
      format = "pdf",
      source = list(bytes = x@data)
    )
  )
}

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolUseBlock.html
method(as_json, list(ProviderAWSBedrock, ContentToolRequest)) <- function(
  provider,
  x,
  ...
) {
  list(
    toolUse = list(
      toolUseId = x@id,
      name = x@name,
      input = x@arguments
    )
  )
}

# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultBlock.html
method(as_json, list(ProviderAWSBedrock, ContentToolResult)) <- function(
  provider,
  x,
  ...
) {
  list(
    toolResult = list(
      toolUseId = x@request@id,
      content = list(list(text = tool_string(x))),
      status = if (tool_errored(x)) "error" else "success"
    )
  )
}

method(as_json, list(ProviderAWSBedrock, ToolDef)) <- function(
  provider,
  x,
  ...
) {
  list(
    toolSpec = list(
      name = x@name,
      description = x@description,
      inputSchema = list(json = compact(as_json(provider, x@arguments, ...)))
    )
  )
}

method(as_json, list(ProviderAWSBedrock, ContentThinking)) <- function(
  provider,
  x,
  ...
) {
  if (identical(x@thinking, "")) {
    return()
  }

  list(
    reasoningContent = list(
      reasoningText = list(
        text = x@thinking,
        signature = x@extra$signature
      )
    )
  )
}

# Helpers ----------------------------------------------------------------

as_bedrock_cache_point <- function(cache_point, model) {
  cache_point <- arg_match(
    cache_point,
    values = c("auto", "5m", "1h", "none")
  )
  if (cache_point != "auto") {
    return(cache_point)
  }
  supports_caching <-
    grepl("(^|\\.)anthropic\\.", model) || grepl("(^|\\.)amazon\\.nova", model)
  if (supports_caching) "5m" else "none"
}

bedrock_cache_point <- function(provider) {
  if (provider@cache_point == "none") {
    return(list())
  }
  cp <- list(type = "default")
  if (provider@cache_point != "5m") {
    cp$ttl <- provider@cache_point
  }
  list(list(cachePoint = cp))
}

paws_credentials <- function(
  profile,
  cache = aws_creds_cache(profile),
  reauth = FALSE
) {
  creds <- cache$get()
  if (reauth || is.null(creds) || creds$expiration < Sys.time()) {
    cache$clear()
    try_fetch(
      creds <- locate_aws_credentials(profile),
      error = function(cnd) {
        if (is_testing()) {
          testthat::skip("Failed to locate AWS credentials")
        }
        cli::cli_abort("No IAM credentials found.", parent = cnd)
      }
    )
    cache$set(creds)
  }
  creds
}

# Wrapper for paws.common::locate_credentials() so we can mock it in tests.
locate_aws_credentials <- function(profile) {
  paws.common::locate_credentials(profile, signing_name = "bedrock")
}

aws_creds_cache <- function(profile) {
  credentials_cache(key = hash(c("aws", profile)))
}

bedrock_document_name <- local({
  i <- 1
  function() {
    i <<- i + 1
    paste0("document-", i)
  }
})
