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

#' Chat with a Google Gemini or Vertex AI model
#'
#' @description
#' `r support_badge("official")`
#'
#' Google's AI offering is broken up into two parts: Gemini and Vertex AI.
#' Most enterprises are likely to use Vertex AI, and individuals are likely
#' to use Gemini.
#'
#' Use [google_upload()] to upload files (PDFs, images, video, audio, etc.)
#'
#' ## Authentication
#' These functions try a number of authentication strategies, in this order:
#'
#' * An API key set in the `GOOGLE_API_KEY` or `GEMINI_API_KEY` env var
#'   (Gemini only).
#' * Google's default application credentials, if the \pkg{gargle} package
#'   is installed.
#' * Viewer-based credentials on Posit Connect, if the \pkg{connectcreds}
#'   package.
#' * `r lifecycle::badge("experimental")`. An browser-based OAuth flow, if
#'   you're in an interactive session. This currently uses an unverified
#'   OAuth app (so you will get a scary warning); we plan to verify in the
#'   near future.
#'
#' @param api_key `r lifecycle::badge("deprecated")` Use `credentials` instead.
#' @param credentials A function that returns a list of authentication headers
#'   or `NULL`, the default, to use ambient credentials. See above for details.
#' @param model `r param_model("gemini-3.5-flash", "google_gemini")`
#' @inheritParams chat_openai
#' @inherit chat_openai return
#' @family chatbots
#' @export
#' @examples
#' \dontrun{
#' chat <- chat_google_gemini()
#' chat$chat("Tell me three jokes about statisticians")
#' }
chat_google_gemini <- function(
  system_prompt = NULL,
  base_url = "https://generativelanguage.googleapis.com/v1beta/",
  api_key = NULL,
  credentials = NULL,
  model = NULL,
  params = NULL,
  api_args = list(),
  api_headers = character(),
  echo = NULL
) {
  model <- set_default(model, "gemini-3.5-flash")
  echo <- check_echo(echo)

  credentials <- as_credentials(
    "chat_google_gemini",
    default_google_credentials(variant = "gemini"),
    credentials = credentials,
    api_key = api_key
  )

  provider <- ProviderGoogleGemini(
    name = "Google/Gemini",
    base_url = base_url,
    model = model,
    params = params %||% params(),
    extra_args = api_args,
    extra_headers = api_headers,
    credentials = credentials
  )
  Chat$new(provider = provider, system_prompt = system_prompt, echo = echo)
}

chat_google_gemini_test <- function(
  ...,
  model = "gemini-3.5-flash",
  params = NULL,
  echo = "none"
) {
  params <- params %||% params()
  params$temperature <- params$temperature %||% 0
  params$seed <- 1014

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

#' @export
#' @rdname chat_google_gemini
#' @param location Location, e.g. `us-east1`, `me-central1`, `africa-south1` or
#'   `global`.
#' @param project_id Project ID.
chat_google_vertex <- function(
  location = Sys.getenv("GOOGLE_CLOUD_LOCATION"),
  project_id = Sys.getenv("GOOGLE_CLOUD_PROJECT"),
  system_prompt = NULL,
  model = NULL,
  params = NULL,
  api_args = list(),
  api_headers = character(),
  echo = NULL
) {
  check_string(location, allow_empty = FALSE)
  check_string(project_id, allow_empty = FALSE)

  model <- set_default(model, "gemini-3.5-flash")
  echo <- check_echo(echo)
  credentials <- default_google_credentials(variant = "vertex")

  provider <- ProviderGoogleGemini(
    name = "Google/Vertex",
    base_url = vertex_url(location, project_id),
    model = model,
    params = params %||% params(),
    extra_args = api_args,
    extra_headers = api_headers,
    credentials = credentials,
    project_id = project_id
  )
  Chat$new(provider = provider, system_prompt = system_prompt, echo = echo)
}

# https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent
vertex_url <- function(location, project_id) {
  paste_c(
    c("https://", google_location(location), "aiplatform.googleapis.com"),
    "/v1",
    c("/projects/", project_id),
    c("/locations/", location),
    "/publishers/google/"
  )
}

ProviderGoogleGemini <- new_class(
  "ProviderGoogleGemini",
  parent = Provider,
  properties = list(
    model = prop_string(),
    project_id = prop_string(allow_null = TRUE)
  )
)

# Base request -----------------------------------------------------------------

method(base_request, ProviderGoogleGemini) <- function(provider) {
  req <- request(provider@base_url)
  req <- ellmer_req_credentials(req, provider@credentials(), "x-goog-api-key")
  req <- ellmer_req_robustify(req)
  req <- ellmer_req_user_agent(req)
  req <- req_error(req, body = function(resp) {
    json <- resp_body_json(resp, check_type = FALSE)
    json$error$message
  })
  req
}

# Chat -------------------------------------------------------------------------

method(chat_request, ProviderGoogleGemini) <- function(
  provider,
  stream = TRUE,
  turns = list(),
  tools = list(),
  type = NULL
) {
  req <- base_request(provider)

  # Can't use chat_path() because it varies based on stream
  req <- req_url_path_append(req, "models")
  if (stream) {
    # https://ai.google.dev/api/generate-content#method:-models.streamgeneratecontent
    req <- req_url_path_append(
      req,
      paste0(provider@model, ":", "streamGenerateContent")
    )
    req <- req_url_query(req, alt = "sse")
  } else {
    # https://ai.google.dev/api/generate-content#method:-models.generatecontent
    req <- req_url_path_append(
      req,
      paste0(provider@model, ":", "generateContent")
    )
  }

  body <- chat_body(
    provider = provider,
    stream = stream,
    turns = turns,
    tools = tools,
    type = type
  )
  body <- modify_list(body, provider@extra_args)

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

method(chat_body, ProviderGoogleGemini) <- function(
  provider,
  stream = TRUE,
  turns = list(),
  tools = list(),
  type = NULL
) {
  if (length(turns) >= 1 && is_system_turn(turns[[1]])) {
    system <- list(parts = list(text = turns[[1]]@text))
  } else {
    system <- list(parts = list(text = ""))
  }

  generation_config <- chat_params(provider, provider@params)
  if (!is.null(type)) {
    generation_config$response_mime_type <- "application/json"
    generation_config$response_schema <- as_json(provider, type)
  }

  if (
    has_name(generation_config, "thinkingBudget") ||
      has_name(generation_config, "thinkingLevel")
  ) {
    generation_config$thinkingConfig <- compact(list(
      thinkingBudget = generation_config$thinkingBudget,
      thinkingLevel = generation_config$thinkingLevel,
      includeThoughts = TRUE
    ))
    generation_config$thinkingBudget <- NULL
    generation_config$thinkingLevel <- NULL
  }

  contents <- as_json(provider, turns)

  # https://ai.google.dev/api/caching#Tool
  tool_config <- NULL
  if (length(tools) > 0) {
    is_builtin <- map_lgl(tools, \(tool) S7_inherits(tool, ToolBuiltIn))
    # Mixing built-in and custom tools on Gemini 3+ requires setting
    # `includeServerSideToolInvocations`. This is only supported by the
    # Gemini Developer API; Vertex AI rejects the field.
    if (
      any(is_builtin) &&
        !all(is_builtin) &&
        has_google_mixed_tool_support(provider)
    ) {
      tool_config <- list(includeServerSideToolInvocations = TRUE)
    }
  }
  tools <- chat_body_tools(provider, tools)

  compact(list(
    contents = contents,
    tools = tools,
    toolConfig = tool_config,
    systemInstruction = system,
    generationConfig = generation_config
  ))
}

method(chat_params, ProviderGoogleGemini) <- function(provider, params) {
  standardise_params(
    params,
    c(
      temperature = "temperature",
      topP = "top_p",
      topK = "top_k",
      frequencyPenalty = "frequency_penalty",
      presencePenalty = "presence_penalty",
      seed = "seed",
      maxOutputTokens = "max_tokens",
      responseLogprobs = "log_probs",
      stopSequences = "stop_sequences",
      thinkingBudget = "reasoning_tokens",
      thinkingLevel = "reasoning_effort"
    )
  )
}

# Gemini -> ellmer --------------------------------------------------------------

method(stream_parse, ProviderGoogleGemini) <- function(provider, event) {
  if (is.null(event)) {
    NULL
  } else {
    jsonlite::parse_json(event$data)
  }
}
method(stream_content, ProviderGoogleGemini) <- function(provider, event) {
  parts <- event$candidates[[1]]$content$parts
  if (is.null(parts) || length(parts) == 0) {
    return(NULL)
  }

  part <- parts[[1]]
  if (isTRUE(part$thought) && !is.null(part$text)) {
    ContentThinking(part$text)
  } else if (!is.null(part$text)) {
    ContentText(part$text)
  }
}
method(stream_merge_chunks, ProviderGoogleGemini) <- function(
  provider,
  result,
  chunk
) {
  if (is.null(result)) {
    chunk
  } else {
    merge_gemini_chunks(result, chunk)
  }
}

method(value_tokens, ProviderGoogleGemini) <- function(provider, json) {
  # https://ai.google.dev/api/generate-content#UsageMetadata
  usage <- json$usageMetadata

  # Total token count for the generation request (prompt + response candidates).
  # Not documented, but appears to include thinking and tool use, i.e.
  # usage$promptTokenCount + usage$candidatesTokenCount +
  #  usage$toolUsePromptTokenCount + usage$thoughtsTokenCount ==
  #  usage$totalTokenCount
  total <- usage$totalTokenCount %||% 0

  # Number of tokens in the prompt. When cachedContent is set, this is
  # still the total effective prompt size meaning this includes the number
  # of tokens in the cached content.
  input <- usage$promptTokenCount %||% 0

  cached <- usage$cachedContentTokenCount %||% 0

  tokens(
    input = input - cached,
    output = total - input,
    cached_input = cached
  )
}

# https://ai.google.dev/api/generate-content
method(value_finish_reason, ProviderGoogleGemini) <- function(
  provider,
  result
) {
  reason <- result$candidates[[1]]$finishReason
  if (is.null(reason)) {
    return(NA_character_)
  }
  switch(
    reason,
    STOP = "success",
    MAX_TOKENS = "max_tokens",
    SAFETY = ,
    RECITATION = ,
    BLOCKLIST = ,
    PROHIBITED_CONTENT = ,
    SPII = "content_filter",
    I(reason)
  )
}

method(value_turn, ProviderGoogleGemini) <- function(
  provider,
  result,
  has_type = FALSE
) {
  message <- result$candidates[[1]]$content

  contents <- lapply(message$parts, function(content) {
    if (isTRUE(content$thought) && has_name(content, "text")) {
      ContentThinking(content$text)
    } else if (has_name(content, "text")) {
      if (has_type) {
        ContentJson(string = content$text)
      } else {
        ContentText(content$text)
      }
    } else if (has_name(content, "functionCall")) {
      extra <- compact(list(
        thoughtSignature = content$thoughtSignature
      ))
      ContentToolRequest(
        content$functionCall$name,
        content$functionCall$name,
        content$functionCall$args,
        extra = extra
      )
    } else if (has_name(content, "inlineData")) {
      ContentImageInline(
        type = content$inlineData$mimeType,
        data = content$inlineData$data
      )
    } else if (
      has_name(content, "toolCall") || has_name(content, "toolResponse")
    ) {
      # Server-side built-in tool invocations (e.g. google_search) returned
      # when includeServerSideToolInvocations is TRUE; skip silently.
      NULL
    } else {
      cli::cli_abort(
        "Unknown content type with names {.str {names(content)}}.",
        .internal = TRUE
      )
    }
  })
  contents <- compact(contents)
  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 -> Gemini --------------------------------------------------------------

# https://ai.google.dev/api/caching#Content
method(as_json, list(ProviderGoogleGemini, Turn)) <- function(
  provider,
  x,
  ...
) {
  if (is_system_turn(x)) {
    # System messages go in the top-level API parameter
  } else if (is_user_turn(x)) {
    x <- turn_contents_expand(x)
    list(
      role = x@role,
      parts = as_json(provider, x@contents, ...)
    )
  } else if (is_assistant_turn(x)) {
    list(role = "model", parts = as_json(provider, x@contents, ...))
  } else {
    cli::cli_abort("Unknown role {x@role}", .internal = TRUE)
  }
}


method(as_json, list(ProviderGoogleGemini, ToolDef)) <- function(
  provider,
  x,
  ...
) {
  compact(list(
    name = x@name,
    description = x@description,
    parameters = as_json(provider, x@arguments, ...)
  ))
}

method(as_json, list(ProviderGoogleGemini, ContentText)) <- function(
  provider,
  x,
  ...
) {
  if (identical(x@text, "")) {
    # Gemini tool call requests can include a Content with empty text,
    # but it doesn't like it if you send this back
    NULL
  } else {
    list(text = x@text)
  }
}

method(as_json, list(ProviderGoogleGemini, ContentThinking)) <- function(
  provider,
  x,
  ...
) {
  # https://ai.google.dev/gemini-api/docs/thinking
  list(thought = TRUE, text = x@thinking)
}

method(as_json, list(ProviderGoogleGemini, ContentPDF)) <- function(
  provider,
  x,
  ...
) {
  list(
    inlineData = list(
      mimeType = x@type,
      data = x@data
    )
  )
}

# https://ai.google.dev/api/caching#FileData
method(as_json, list(ProviderGoogleGemini, ContentUploaded)) <- function(
  provider,
  x,
  ...
) {
  list(
    fileData = list(
      mimeType = x@mime_type,
      fileUri = x@uri
    )
  )
}

# https://ai.google.dev/api/caching#FileData
method(as_json, list(ProviderGoogleGemini, ContentImageRemote)) <- function(
  provider,
  x,
  ...
) {
  cli::cli_abort("Gemini doesn't support remote images")
}

# https://ai.google.dev/api/caching#Blob
method(as_json, list(ProviderGoogleGemini, ContentImageInline)) <- function(
  provider,
  x,
  ...
) {
  list(
    inlineData = list(
      mimeType = x@type,
      data = x@data
    )
  )
}

# https://ai.google.dev/api/caching#FunctionCall
method(as_json, list(ProviderGoogleGemini, ContentToolRequest)) <- function(
  provider,
  x,
  ...
) {
  compact(list(
    functionCall = list(
      name = x@id,
      args = x@arguments
    ),
    thoughtSignature = x@extra$thoughtSignature
  ))
}

# https://ai.google.dev/api/caching#FunctionResponse
method(as_json, list(ProviderGoogleGemini, ContentToolResult)) <- function(
  provider,
  x,
  ...
) {
  list(
    functionResponse = list(
      name = x@request@id,
      response = list(value = tool_string(x))
    )
  )
}

method(as_json, list(ProviderGoogleGemini, TypeObject)) <- function(
  provider,
  x,
  ...
) {
  if (x@additional_properties) {
    cli::cli_abort("{.arg .additional_properties} not supported for Gemini.")
  }

  if (length(x@properties) == 0) {
    return(list())
  }

  required <- map_lgl(x@properties, function(prop) prop@required)

  compact(list(
    type = "object",
    description = x@description,
    properties = as_json(provider, x@properties, ...),
    required = as.list(names2(x@properties)[required])
  ))
}

# Gemini-specific merge logic --------------------------------------------------

merge_last <- function() {
  function(left, right, path = NULL) {
    right
  }
}

merge_identical <- function() {
  function(left, right, path = NULL) {
    if (!identical(left, right)) {
      stop(
        "Expected identical values, but got ",
        deparse(left),
        " and ",
        deparse(right)
      )
    }
    left
  }
}

merge_any_or_empty <- function() {
  function(left, right, path = NULL) {
    if (!is.null(left) && nzchar(left)) {
      left
    } else if (!is.null(right) && nzchar(right)) {
      right
    } else {
      ""
    }
  }
}

merge_optional <- function(merge_func) {
  function(left, right, path = NULL) {
    if (is.null(left) && is.null(right)) {
      NULL
    } else {
      merge_func(left, right, path)
    }
  }
}

merge_objects <- function(...) {
  spec <- list(...)
  function(left, right, path = NULL) {
    if (is.null(left)) {
      return(right)
    } else if (is.null(right)) {
      return(left)
    }

    # cat(paste(collapse = "", path), "\n")
    stopifnot(is.list(left), is.list(right), all(nzchar(names(spec))))
    mapply(
      names(spec),
      spec,
      FUN = function(key, value) {
        value(left[[key]], right[[key]], c(path, ".", key))
      },
      USE.NAMES = TRUE,
      SIMPLIFY = FALSE
    )
  }
}

merge_candidate_lists <- function(...) {
  merge_unindexed <- merge_objects(...)
  merge_indexed <- merge_objects(index = merge_identical(), ...)

  function(left, right, path = NULL) {
    if (length(left) == 1 && length(right) == 1) {
      list(merge_unindexed(left[[1]], right[[1]], c(path, "[]")))
    } else {
      # left and right are lists of objects with [["index"]]
      # We need to find the elements that have matching indices and merge them
      left_indices <- vapply(left, `[[`, integer(1), "index")
      right_indices <- vapply(right, `[[`, integer(1), "index")
      # I know this seems weird, but according to Google's Go SDK, we should
      # only retain indices on the right that *already* appear on the left.
      # Citations:
      # https://github.com/google/generative-ai-go/blob/3d14f4039eaef321b15bcbf70839389d7f000233/genai/client_test.go#L655
      # https://github.com/google/generative-ai-go/blob/3d14f4039eaef321b15bcbf70839389d7f000233/genai/client.go#L396
      lapply(left_indices, function(index) {
        left_item <- left[[which(left_indices == index)]]
        right_item <- right[[which(right_indices == index)]]
        if (is.null(right_item)) {
          left_item
        } else {
          merge_indexed(left_item, right_item, c(path, "[", index, "]"))
        }
      })
    }
  }
}

merge_append <- function() {
  function(left, right, path = NULL) {
    c(left, right)
  }
}

merge_parts <- function() {
  function(left, right, path = NULL) {
    joined <- c(left, right)

    # Identify text parts
    is_text <- map_lgl(joined, ~ is.list(.x) && identical(names(.x), "text"))

    # Create groups for contiguous sections
    groups <- cumsum(c(TRUE, diff(is_text) != 0))

    # Split into groups and process each
    split_parts <- split(joined, groups)
    merged_split_parts <- map2(
      split_parts,
      split(is_text, groups),
      function(parts, is_text_group) {
        if (!is_text_group[[1]]) {
          # Non-text group: return parts unchanged
          return(parts)
        } else {
          # Text group: merge text values
          text_values <- map_chr(parts, ~ .x[["text"]])
          list(list(text = paste0(text_values, collapse = "")))
        }
      }
    )
    unlist(merged_split_parts, recursive = FALSE, use.names = FALSE)
  }
}

# Put it all together...
# https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse
merge_gemini_chunks <- merge_objects(
  candidates = merge_candidate_lists(
    content = merge_objects(
      role = merge_any_or_empty(),
      parts = merge_parts()
    ),
    finishReason = merge_last(),
    safetyRatings = merge_last(),
    citationMetadata = merge_optional(
      merge_objects(citationSources = merge_append())
    ),
    tokenCount = merge_last()
  ),
  promptFeedback = merge_last(),
  usageMetadata = merge_last()
)

default_google_credentials <- function(
  error_call = caller_env(),
  variant = c("gemini", "vertex")
) {
  variant <- arg_match(variant)

  if (variant == "gemini") {
    api_key <- Sys.getenv("GOOGLE_API_KEY")
    if (!nzchar(api_key)) {
      api_key <- Sys.getenv("GEMINI_API_KEY")
    }
    if (nzchar(api_key)) {
      return(\() api_key)
    }
  }

  gemini_scope <- switch(
    variant,
    gemini = "https://www.googleapis.com/auth/generative-language.retriever",
    # https://github.com/googleapis/python-genai/blob/cc9e470326e0c1b84ec3ce9891c9f96f6c74688e/google/genai/_api_client.py#L184
    vertex = "https://www.googleapis.com/auth/cloud-platform"
  )

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

  if (is_testing()) {
    testthat::skip_if_not_installed("gargle")
  }

  check_installed("gargle", "for Google authentication")
  gargle::with_cred_funs(
    funs = list(
      # We don't want to use *all* of gargle's default credential functions --
      # in particular, we don't want to try and authenticate using the bundled
      # OAuth client -- so winnow down the list.
      credentials_app_default = gargle::credentials_app_default
    ),
    {
      token <- gargle::token_fetch(scopes = gemini_scope)
    },
    action = "replace"
  )

  if (is.null(token) && is_testing()) {
    testthat::skip("no Google credentials available")
  }

  if (is.null(token) && is_interactive()) {
    return(function() {
      function(req) {
        req_oauth_auth_code(
          req,
          client = gemini_client(),
          auth_url = "https://accounts.google.com/o/oauth2/auth",
          scope = "https://www.googleapis.com/auth/generative-language.retriever"
        )
      }
    })
  }

  if (is.null(token)) {
    cli::cli_abort(
      c(
        "No Google credentials are available.",
        "i" = "Try suppling an API key or configuring Google's application default credentials."
      ),
      call = error_call
    )
  }

  # A token can exist but have a NULL access_token if it has expired.
  if (is.null(token$credentials$access_token)) {
    cli::cli_abort(
      c(
        "Google credentials were found but are not valid.",
        "i" = "Try refreshing your credentials or configuring new ones."
      ),
      call = error_call
    )
  }

  # gargle emits an httr-style token, which we awkwardly shim into something
  # httr2 can work with.
  if (!token$can_refresh()) {
    # TODO: Not really sure what to do in this case when the token expires.
    return(function() {
      list(Authorization = paste("Bearer", token$credentials$access_token))
    })
  }

  # gargle tokens don't track the expiry time, so we do it ourselves (with a
  # grace period).
  expiry <- Sys.time() + token$credentials$expires_in - 5
  return(function() {
    if (expiry < Sys.time()) {
      token$refresh()
    }
    list(Authorization = paste("Bearer", token$credentials$access_token))
  })
}

google_oauth_reset <- function() {
  httr2::oauth_cache_clear(gemini_client())
}

gemini_client <- function() {
  httr2::oauth_client(
    id = "148439353047-kit3pok9u920mhmqbc3c0pdr50bvb7pt.apps.googleusercontent.com",
    secret = httr2::obfuscated(
      "o2yDPr_4BNgZvhLT9kIZS6jAYp43sAzAjMrmW60FUC-N4btRmTwOQ1650vS2pDRSvbKK"
    ),
    token_url = "https://oauth2.googleapis.com/token",
    name = "gemini-r-client"
  )
}

method(chat_body_tools, ProviderGoogleGemini) <- function(provider, tools) {
  if (length(tools) == 0) {
    return(NULL)
  }
  is_builtin <- map_lgl(tools, \(tool) S7_inherits(tool, ToolBuiltIn))
  funs <- as_json(provider, unname(tools))
  user_tools <- funs[!is_builtin]
  c(
    if (length(user_tools) > 0) list(list(functionDeclarations = user_tools)),
    funs[is_builtin]
  )
}

# Token counting ----------------------------------------------------------

# https://ai.google.dev/api/tokens
method(count_tokens, ProviderGoogleGemini) <- function(
  provider,
  ...,
  system_prompt = NULL,
  tools = list(),
  type = NULL
) {
  req <- base_request(provider)
  req <- req_url_path_append(
    req,
    "models",
    paste0(provider@model, ":", "countTokens")
  )

  if (!is.null(system_prompt)) {
    system <- list(parts = list(text = system_prompt))
  } else {
    system <- list(parts = list(text = ""))
  }

  contents <- as_json(provider, list(user_turn(...)))
  tools <- chat_body_tools(provider, tools)

  if (!is.null(type)) {
    generation_config <- list(
      response_mime_type = "application/json",
      response_schema = as_json(provider, type)
    )
  } else {
    generation_config <- NULL
  }

  token_body <- compact(list(
    contents = contents,
    tools = tools,
    systemInstruction = system,
    generationConfig = generation_config
  ))

  if (identical(provider@name, "Google/Gemini")) {
    token_body$model <- paste0("models/", provider@model)
    token_body <- list(generateContentRequest = token_body)
  }

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

  resp <- req_perform(req)
  resp_body_json(resp)$totalTokens
}

# Pricing ----------------------------------------------------------------------

# Models -----------------------------------------------------------------------

#' @export
#' @rdname chat_google_gemini
models_google_gemini <- function(
  base_url = "https://generativelanguage.googleapis.com/v1beta/",
  api_key = NULL,
  credentials = NULL
) {
  check_string(base_url)

  credentials <- as_credentials(
    "models_google_gemini",
    default_google_credentials(variant = "gemini"),
    credentials = credentials,
    api_key = api_key
  )

  provider <- ProviderGoogleGemini(
    name = "Google/Gemini",
    model = "",
    base_url = base_url,
    credentials = credentials
  )
  models_list(provider)
}

#' @rdname chat_google_gemini
#' @export
models_google_vertex <- function(
  location = Sys.getenv("GOOGLE_CLOUD_LOCATION"),
  project_id = Sys.getenv("GOOGLE_CLOUD_PROJECT"),
  credentials = NULL
) {
  check_string(location, allow_empty = FALSE)
  check_string(project_id, allow_empty = FALSE)

  credentials <- credentials %||% default_google_credentials(variant = "vertex")
  check_credentials(credentials)

  base_url <- paste_c(
    c("https://", google_location(location), "aiplatform.googleapis.com"),
    "/v1beta1",
    "/publishers/google/"
  )

  provider <- ProviderGoogleGemini(
    name = "Google/Vertex",
    model = "",
    base_url = base_url,
    credentials = credentials,
    project_id = project_id
  )
  models_list(provider)
}

method(models_list, ProviderGoogleGemini) <- function(provider) {
  is_vertex <- grepl(
    "aiplatform.googleapis.com",
    provider@base_url,
    fixed = TRUE
  )

  req <- base_request(provider)
  if (is_vertex) {
    req <- req_headers(req, `x-goog-user-project` = provider@project_id)
  }
  req <- req_headers(req, !!!provider@extra_headers)
  req <- req_url_path_append(req, "/models")
  resp <- req_perform(req)

  json <- resp_body_json(resp)

  if (is_vertex) {
    name <- map_chr(json$publisherModels, "[[", "name")
    name <- gsub("^publishers/google/models/", "", name)
    # this is the closest to "generateContent" in "supportedGenerationMethods" for Gemini
    # https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/publishers.models
    can_generate <- json$publisherModels |>
      map_lgl(\(x) "openGenerationAiStudio" %in% names(x$supportedActions))
  } else {
    name <- map_chr(json$models, "[[", "name")
    name <- gsub("^models/", "", name)

    methods <- map(json$models, \(x) unlist(x$supportedGenerationMethods))
    can_generate <- map_lgl(methods, \(x) "generateContent" %in% x)
  }

  df <- data.frame(id = name)
  df <- cbind(df, match_prices(provider@name, df$id))
  df <- df[can_generate, ]
  unrowname(df[order(df$id), ])
}

# for location "global", there is no location in the final base URL
# https://github.com/googleapis/python-genai/blob/cc9e470326e0c1b84ec3ce9891c9f96f6c74688e/google/genai/_api_client.py#L646-L654
google_location <- function(location) {
  if (location == "global") "" else paste0(location, "-")
}

# Batched requests -------------------------------------------------------------

# https://ai.google.dev/gemini-api/docs/batch-api
# Only the Gemini Developer API is supported; Vertex AI's batch API has a
# different request shape (GCS bucket URIs instead of file uploads).
method(has_batch_support, ProviderGoogleGemini) <- function(provider) {
  identical(provider@name, "Google/Gemini")
}

method(batch_submit, ProviderGoogleGemini) <- function(
  provider,
  conversations,
  type = NULL
) {
  path <- withr::local_tempfile(fileext = ".jsonl")

  requests <- map(seq_along(conversations), function(i) {
    body <- chat_body(
      provider,
      stream = FALSE,
      turns = conversations[[i]],
      type = type
    )

    list(
      key = paste0("chat-", i),
      request = gemini_prepare_batch_body(body)
    )
  })

  json_lines <- map_chr(requests, to_json)
  writeLines(json_lines, path)

  uploaded <- gemini_upload_file(provider, path)
  if (is.null(uploaded$name) || !nzchar(uploaded$name)) {
    cli::cli_abort(
      "Gemini upload did not return a file resource name.",
      .internal = TRUE
    )
  }

  req <- base_request(provider)
  req <- req_url_path_append(
    req,
    "models",
    paste0(provider@model, ":batchGenerateContent")
  )
  req <- req_body_json(
    req,
    list(
      batch = list(
        displayName = paste0("ellmer-", as.integer(Sys.time())),
        model = paste0("models/", provider@model),
        inputConfig = list(fileName = uploaded$name)
      )
    )
  )

  resp <- req_perform(req)
  resp_body_json(resp)
}

method(batch_poll, ProviderGoogleGemini) <- function(provider, batch) {
  req <- base_request(provider)
  req <- req_url_path_append(req, batch$name)
  resp <- req_perform(req)
  resp_body_json(resp)
}

method(batch_status, ProviderGoogleGemini) <- function(provider, batch) {
  metadata <- batch$metadata %||% list()
  stats <- metadata$batchStats %||% list()
  state <- metadata$state %||% "BATCH_STATE_UNSPECIFIED"

  total <- as.integer(stats$requestCount %||% 0L)
  pending <- as.integer(stats$pendingRequestCount %||% 0L)
  succeeded <- as.integer(stats$successfulRequestCount %||% 0L)
  failed <- as.integer(stats$failedRequestCount %||% 0L)

  if (!is.null(batch$error) && total > 0 && failed == 0L) {
    failed <- total
  }

  terminal_states <- c(
    "BATCH_STATE_SUCCEEDED",
    "BATCH_STATE_FAILED",
    "BATCH_STATE_CANCELLED",
    "BATCH_STATE_EXPIRED"
  )

  is_done <- state %in% terminal_states

  # The API can report BATCH_STATE_SUCCEEDED before responsesFile is populated.
  if (state == "BATCH_STATE_SUCCEEDED") {
    responses_file <- batch$response$responsesFile %||% ""
    if (!nzchar(responses_file)) {
      is_done <- FALSE
    }
  }

  n_processing <- max(pending, total - succeeded - failed, 0L)

  list(
    working = !is_done,
    n_processing = n_processing,
    n_succeeded = max(succeeded, 0L),
    n_failed = max(failed, 0L)
  )
}

method(batch_retrieve, ProviderGoogleGemini) <- function(provider, batch) {
  metadata <- batch$metadata %||% list()
  stats <- metadata$batchStats %||% list()
  request_count <- as.integer(stats$requestCount %||% 0L)

  if (!is.null(batch$error)) {
    code <- as.integer(batch$error$code %||% 500L)
    return(rep(
      list(list(status_code = code, body = NULL)),
      max(0L, request_count)
    ))
  }

  responses_file <- batch$response$responsesFile %||% ""

  if (!nzchar(responses_file)) {
    cli::cli_abort(
      "Gemini batch completed but no output file was returned.",
      .internal = TRUE
    )
  }

  path_output <- withr::local_tempfile(fileext = ".jsonl")
  gemini_download_file(provider, responses_file, path_output)

  parsed <- read_ndjson(path_output)

  normalized <- imap(parsed, function(x, i) {
    gemini_normalize_result(x, index_default = as.integer(i))
  })

  ids <- vapply(normalized, function(x) x$index, integer(1))
  results <- lapply(normalized, function(x) x$result)
  results[order(ids)]
}

method(batch_result_turn, ProviderGoogleGemini) <- function(
  provider,
  result,
  has_type = FALSE
) {
  if (!is.null(result) && result$status_code == 200L && !is.null(result$body)) {
    value_turn(provider, result$body, has_type = has_type)
  } else {
    NULL
  }
}

# Gemini batch helpers ---------------------------------------------------------

# The Gemini REST API accepts both camelCase and snake_case, but the batch
# JSONL file parser requires protobuf field names which are always snake_case.
# Without this conversion, the batch API rejects requests with HTTP 400.
gemini_to_snake_case <- function(x) {
  if (is.list(x)) {
    if (!is.null(names(x))) {
      names(x) <- gsub("([a-z])([A-Z])", "\\1_\\2", names(x), perl = TRUE) |>
        tolower()
    }
    lapply(x, gemini_to_snake_case)
  } else {
    x
  }
}

gemini_prepare_batch_body <- function(body) {
  # Drop empty system instructions ; the batch JSONL parser rejects empty text.
  text <- body$systemInstruction$parts$text
  if (!is.null(text) && identical(text, "")) {
    body$systemInstruction <- NULL
  }

  # Save schema before snake_case conversion so user property names like
  # "firstName" survive. Look up both spellings since chat_body mixes cases.
  gc <- body$generationConfig
  saved_schema <- gc$response_schema %||% gc$responseSchema

  body <- gemini_to_snake_case(body)

  # The batch JSONL parser uses response_json_schema, not response_schema.
  if (!is.null(saved_schema)) {
    body$generation_config$response_json_schema <- saved_schema
    body$generation_config$response_schema <- NULL
  }

  body
}

gemini_extract_index <- function(x, default = NA_integer_) {
  key <- x$key %||% ""
  if (grepl("^chat-[0-9]+$", key)) {
    return(as.integer(sub("^chat-([0-9]+)$", "\\1", key)))
  }
  as.integer(default)
}

gemini_normalize_result <- function(x, index_default) {
  index <- gemini_extract_index(x, default = index_default)

  if (!is.null(x$response) && is.null(x$error) && is.null(x$status)) {
    return(list(
      index = index,
      result = list(status_code = 200L, body = x$response)
    ))
  }

  if (!is.null(x$error) || !is.null(x$status)) {
    code <- x$error$code %||% 500L
    return(list(
      index = index,
      result = list(status_code = as.integer(code), body = NULL)
    ))
  }

  list(index = index, result = list(status_code = 500L, body = NULL))
}

has_google_mixed_tool_support <- function(provider) {
  identical(provider@name, "Google/Gemini") &&
    grepl("^gemini-([3-9]|[0-9]{2,})", provider@model)
}
