defmodule LangChain.ChatModels.ChatVertexAI do
  @moduledoc """
  Parses and validates inputs for making a request for the Google AI  Chat API.

  Converts response into more specialized `LangChain` data structures.

  Example Usage:

  ```elixir
  alias LangChain.Chains.LLMChain
  alias LangChain.Message
  alias LangChain.Message.ContentPart
  alias LangChain.ChatModels.ChatVertexAI

  config = %{
        model: "gemini-2.0-flash",
        api_key: ..., # vertex requires gcloud auth token https://cloud.google.com/vertex-ai/generative-ai/docs/start/quickstarts/quickstart-multimodal#rest
        temperature: 1.0,
        top_p: 0.8,
        receive_timeout: ...
      }
   model = ChatVertexAI.new!(config)

      %{llm: model, verbose: false, stream: false}
      |> LLMChain.new!()
      |> LLMChain.add_message(
        Message.new_user!([
          ContentPart.new!(%{type: :text, content: "Analyse the provided file and share a summary"}),
          ContentPart.new!(%{
            type: :file_url,
            content: ...,
            options: [media: ...]
          })
        ])
      )
      |> LLMChain.run()
  The above call will return summary of the media content.
  ```

  **Tool Schemas**

  Vertex AI uses the same Gemini `Schema` type as
  `LangChain.ChatModels.ChatGoogleAI`, which accepts a select subset of an
  OpenAPI 3.0 schema object rather than full JSON Schema. Tool parameters and
  the JSON response schema are passed through
  `LangChain.Utils.GoogleSchema.sanitize/1` before being sent, which removes
  the keywords the API has no field for. Note that dropping
  `additionalProperties` means Vertex does not enforce closed objects, so
  validate a tool call's arguments inside the tool's own function where that
  matters.
  """
  use Ecto.Schema
  require Logger
  import Ecto.Changeset
  alias __MODULE__
  alias LangChain.Config
  alias LangChain.ChatModels.ChatModel
  alias LangChain.ChatModels.ChatOpenAI
  alias LangChain.Message
  alias LangChain.MessageDelta
  alias LangChain.Message.ContentPart
  alias LangChain.Message.ToolCall
  alias LangChain.Message.ToolResult
  alias LangChain.LangChainError
  alias LangChain.Utils
  alias LangChain.Utils.GoogleSchema
  alias LangChain.Callbacks
  alias LangChain.TokenUsage
  alias LangChain.NativeTool
  alias LangChain.Function

  @behaviour ChatModel

  @current_config_version 1

  # allow up to 2 minutes for response.
  @receive_timeout 60_000

  @primary_key false
  embedded_schema do
    field :endpoint, :string

    field :model, :string, default: "gemini-pro"
    field :api_key, :string, redact: true

    # What sampling temperature to use, between 0 and 2. Higher values like 0.8
    # will make the output more random, while lower values like 0.2 will make it
    # more focused and deterministic.
    field :temperature, :float, default: 0.9

    # The topP parameter changes how the model selects tokens for output. Tokens
    # are selected from the most to least probable until the sum of their
    # probabilities equals the topP value. For example, if tokens A, B, and C have
    # a probability of 0.3, 0.2, and 0.1 and the topP value is 0.5, then the model
    # will select either A or B as the next token by using the temperature and exclude
    # C as a candidate. The default topP value is 0.95.
    field :top_p, :float, default: 1.0

    # The topK parameter changes how the model selects tokens for output. A topK of
    # 1 means the selected token is the most probable among all the tokens in the
    # model's vocabulary (also called greedy decoding), while a topK of 3 means that
    # the next token is selected from among the 3 most probable using the temperature.
    # For each token selection step, the topK tokens with the highest probabilities
    # are sampled. Tokens are then further filtered based on topP with the final token
    # selected using temperature sampling.
    field :top_k, :float, default: 1.0

    # Configure thinking budget and whether to include thought summaries (content type `:thinking`).
    # See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/thinking
    #
    # Config reference: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.evaluationRuns#ThinkingConfig
    field :thinking_config, :map, default: nil

    # Duration in seconds for the response to be received. When streaming a very
    # lengthy response, a longer time limit may be required. However, when it
    # goes on too long by itself, it tends to hallucinate more.
    field :receive_timeout, :integer, default: @receive_timeout
    field :json_response, :boolean, default: false
    field :json_schema, :map, default: nil
    field :stream, :boolean, default: false

    # A list of maps for callback handlers (treated as internal)
    field :callbacks, {:array, :map}, default: []

    # Additional level of raw api request and response data
    field :verbose_api, :boolean, default: false

    # Req options to merge into the request.
    # Refer to `https://hexdocs.pm/req/Req.html#new/1-options` for
    # `Req.new` supported set of options.
    field :req_config, :map, default: %{}

    # The safety settings for the model, specified as a list of maps. Each map
    # should contain a `category` and a `threshold` for that category.
    # e.g. [%{"category" => "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold" => "BLOCK_ONLY_HIGH"}]
    # see https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters
    # for the list of categories and thresholds. Defaults to `[]` (the
    # provider's own defaults apply).
    field :safety_settings, {:array, :map}, default: []
  end

  @type t :: %ChatVertexAI{}

  @create_fields [
    :endpoint,
    :model,
    :api_key,
    :temperature,
    :top_p,
    :top_k,
    :thinking_config,
    :receive_timeout,
    :json_response,
    :json_schema,
    :stream,
    :req_config,
    :safety_settings,
    :callbacks
  ]
  @required_fields [
    :endpoint,
    :model
  ]

  @spec get_api_key(t) :: String.t()
  defp get_api_key(%ChatVertexAI{api_key: api_key}) do
    # if no API key is set default to `""` which will raise an API error
    api_key || Config.resolve(:vertex_ai_key, "")
  end

  @doc """
  Setup a ChatVertexAI client configuration.
  """
  @spec new(attrs :: map()) :: {:ok, t} | {:error, Ecto.Changeset.t()}
  def new(%{} = attrs \\ %{}) do
    %ChatVertexAI{}
    |> cast(attrs, @create_fields)
    |> common_validation()
    |> apply_action(:insert)
  end

  @doc """
  Setup a ChatVertexAI client configuration and return it or raise an error if invalid.
  """
  @spec new!(attrs :: map()) :: t() | no_return()
  def new!(attrs \\ %{}) do
    case new(attrs) do
      {:ok, chain} ->
        chain

      {:error, changeset} ->
        raise LangChainError, changeset
    end
  end

  defp common_validation(changeset) do
    changeset
    |> validate_required(@required_fields)
  end

  def for_api(%ChatVertexAI{} = vertex_ai, messages, functions) do
    {sys_instructions, other_messages} = Utils.split_system_message(messages)

    messages_for_api =
      other_messages
      |> Enum.map(&for_api/1)
      |> List.flatten()
      |> List.wrap()

    {response_mime_type, response_schema} =
      case vertex_ai.json_response do
        true ->
          {"application/json", GoogleSchema.sanitize(vertex_ai.json_schema)}

        false ->
          {nil, nil}
      end

    generation_config_params =
      %{
        "temperature" => vertex_ai.temperature,
        "topP" => vertex_ai.top_p,
        "topK" => vertex_ai.top_k
      }
      |> Utils.conditionally_add_to_map("thinkingConfig", vertex_ai.thinking_config)
      |> Utils.conditionally_add_to_map("response_mime_type", response_mime_type)
      |> Utils.conditionally_add_to_map("response_schema", response_schema)

    req =
      %{
        "contents" => messages_for_api,
        "generationConfig" => generation_config_params
      }
      |> Utils.conditionally_add_to_map("system_instruction", for_api(sys_instructions))
      |> Utils.conditionally_add_to_map("safetySettings", vertex_ai.safety_settings)

    if functions && not Enum.empty?(functions) do
      req
      |> Map.put("tools", [
        %{
          # Google AI functions use an OpenAI compatible format.
          # See: https://ai.google.dev/docs/function_calling#how_it_works
          # Note: We strip the "strict" field as it's OpenAI-specific and not supported by Vertex AI
          "functionDeclarations" =>
            functions
            |> Enum.map(&ChatOpenAI.for_api(vertex_ai, &1))
            |> Enum.map(&Map.delete(&1, "strict"))
            |> Enum.map(&function_declaration_for_api/1)
        }
      ])
    else
      req
    end
  end

  # Vertex AI reaches the same Gemini `Schema` type as `ChatGoogleAI`, so a
  # declaration carrying a keyword outside Google's subset is rejected the same
  # way. Sanitize the parameters, then drop them entirely when nothing is left
  # but an empty object, which Google also errors on with
  # "parameters.properties: should be non-empty for OBJECT type".
  defp function_declaration_for_api(%{"parameters" => parameters} = declaration) do
    sanitized = GoogleSchema.sanitize(parameters)

    if GoogleSchema.empty_object?(sanitized) do
      Map.delete(declaration, "parameters")
    else
      Map.put(declaration, "parameters", sanitized)
    end
  end

  defp function_declaration_for_api(declaration), do: declaration

  defp for_api(%Message{role: :assistant} = message) do
    content_parts = get_message_contents(message) || []
    tool_calls = Enum.map(message.tool_calls || [], &for_api/1)

    %{
      "role" => map_role(:assistant),
      "parts" => content_parts ++ tool_calls
    }
  end

  defp for_api(%Message{role: :tool} = message) do
    %{
      "role" => map_role(:tool),
      "parts" => Enum.map(message.tool_results, &for_api/1)
    }
  end

  defp for_api(%Message{role: :system} = message) do
    # System messages should return a single text part, not a list
    case get_message_contents(message) do
      [%{"text" => text}] -> %{"parts" => %{"text" => text}}
      _ -> %{"parts" => %{"text" => message.content}}
    end
  end

  defp for_api(%Message{role: :user, content: content}) when is_list(content) do
    %{
      "role" => map_role(:user),
      "parts" => Enum.map(content, &for_api(&1))
    }
  end

  defp for_api(%Message{} = message) do
    content_parts = get_message_contents(message) || []

    %{
      "role" => map_role(message.role),
      "parts" => content_parts
    }
  end

  defp for_api(%ContentPart{type: :text} = part) do
    %{"text" => part.content}
  end

  defp for_api(%ContentPart{type: :thinking}) do
    # The thinking parts are only thought summaries and are not meant to be
    # included in future generation requests.
    # See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/thinking#thought-summaries
    []
  end

  defp for_api(%ContentPart{type: :image} = part) do
    %{
      "inlineData" => %{
        "mimeType" => Keyword.fetch!(part.options, :media),
        "data" => part.content
      }
    }
  end

  defp for_api(%ContentPart{type: :image_url} = part) do
    %{
      "fileData" => %{
        "mimeType" => Keyword.fetch!(part.options, :media),
        "fileUri" => part.content
      }
    }
  end

  defp for_api(%ContentPart{type: :file_url} = part) do
    %{
      "fileData" => %{
        "mimeType" => Keyword.fetch!(part.options, :media),
        "fileUri" => part.content
      }
    }
  end

  defp for_api(%ToolCall{metadata: %{thought_signature: signature}} = call)
       when is_binary(signature) do
    %{
      "functionCall" => %{
        "args" => ToolCall.arguments_as_map(call),
        "name" => call.name
      },
      "thoughtSignature" => signature
    }
  end

  defp for_api(%ToolCall{} = call) do
    %{
      "functionCall" => %{
        "args" => ToolCall.arguments_as_map(call),
        "name" => call.name
      }
    }
  end

  defp for_api(%ToolResult{} = result) do
    response = tool_result_response_for_api(result.content)
    response_parts = tool_result_parts_for_api(result.content)

    %{
      "functionResponse" => %{
        "name" => result.name,
        "response" => response
      }
    }
    |> maybe_add_function_response_parts(response_parts)
  end

  defp for_api(%NativeTool{name: name, configuration: %{} = config}) do
    %{name => config}
  end

  defp for_api(%NativeTool{name: name, configuration: nil}) do
    name
  end

  defp for_api(nil), do: nil

  defp maybe_add_function_response_parts(
         %{"functionResponse" => function_response} = data,
         [_ | _] = response_parts
       ) do
    put_in(data, ["functionResponse"], Map.put(function_response, "parts", response_parts))
  end

  defp maybe_add_function_response_parts(data, _response_parts), do: data

  defp tool_result_response_for_api(nil), do: %{}

  defp tool_result_response_for_api(content) when is_binary(content) do
    case Jason.decode(content) do
      {:ok, data} ->
        data

      {:error, %Jason.DecodeError{}} ->
        %{"result" => content}
    end
  end

  defp tool_result_response_for_api(content_parts) when is_list(content_parts) do
    text_parts = Enum.filter(content_parts, &(&1.type == :text))

    case ContentPart.parts_to_string(text_parts) do
      nil ->
        %{}

      text_content ->
        case Jason.decode(text_content) do
          {:ok, data} ->
            data

          {:error, %Jason.DecodeError{}} ->
            %{"result" => text_content}
        end
    end
  end

  defp tool_result_parts_for_api(nil), do: []

  defp tool_result_parts_for_api(content) when is_binary(content), do: []

  defp tool_result_parts_for_api(content_parts) when is_list(content_parts) do
    content_parts
    |> Enum.filter(&tool_result_media_part?/1)
    |> Enum.map(&tool_result_media_part_for_api/1)
  end

  defp tool_result_media_part_for_api(%ContentPart{type: :image} = part) do
    %{
      "inlineData" =>
        %{
          "mimeType" => Keyword.fetch!(part.options, :media),
          "data" => part.content
        }
        |> maybe_add_display_name(part.options)
    }
  end

  defp tool_result_media_part_for_api(%ContentPart{type: :file} = part) do
    %{
      "inlineData" =>
        %{
          "mimeType" => Keyword.fetch!(part.options, :media),
          "data" => part.content
        }
        |> maybe_add_display_name(part.options)
    }
  end

  defp tool_result_media_part_for_api(%ContentPart{type: :image_url} = part) do
    %{
      "fileData" =>
        %{
          "mimeType" => Keyword.fetch!(part.options, :media),
          "fileUri" => part.content
        }
        |> maybe_add_display_name(part.options)
    }
  end

  defp tool_result_media_part_for_api(%ContentPart{type: :file_url} = part) do
    %{
      "fileData" =>
        %{
          "mimeType" => Keyword.fetch!(part.options, :media),
          "fileUri" => part.content
        }
        |> maybe_add_display_name(part.options)
    }
  end

  defp maybe_add_display_name(data, options) do
    case Keyword.get(options || [], :display_name) do
      nil -> data
      display_name -> Map.put(data, "displayName", display_name)
    end
  end

  defp tool_result_media_part?(%ContentPart{type: type})
       when type in [:image, :file, :image_url, :file_url],
       do: true

  defp tool_result_media_part?(%ContentPart{}), do: false

  @doc """
  Calls the Google AI API passing the ChatVertexAI struct with configuration,
  plus either a simple message or the list of messages to act as the prompt.

  Optionally pass in a list of tools available to the LLM for requesting
  execution in response.

  **NOTE:** This function *can* be used directly, but the primary interface
  should be through `LangChain.Chains.LLMChain`. The `ChatVertexAI` module is
  more focused on translating the `LangChain` data structures to and from the
  OpenAI API.

  Another benefit of using `LangChain.Chains.LLMChain` is that it combines the
  storage of messages, adding tools, adding custom context that should be passed
  to tools, and automatically applying `LangChain.MessageDelta` structs as they
  are are received, then converting those to the full `LangChain.Message` once
  fully complete.
  """
  @impl ChatModel
  def call(openai, prompt, tools \\ [])

  def call(%ChatVertexAI{} = vertex_ai, prompt, tools) when is_binary(prompt) do
    messages = [
      Message.new_system!(),
      Message.new_user!(prompt)
    ]

    call(vertex_ai, messages, tools)
  end

  def call(%ChatVertexAI{} = vertex_ai, messages, tools)
      when is_list(messages) do
    metadata = %{
      model: vertex_ai.model,
      provider: provider(),
      message_count: length(messages),
      tools_count: length(tools)
    }

    ChatModel.llm_telemetry_span(vertex_ai, metadata, fn ->
      try do
        # Track the prompt being sent
        LangChain.Telemetry.llm_prompt(
          %{system_time: System.system_time()},
          %{model: vertex_ai.model, messages: messages}
        )

        case do_api_request(vertex_ai, messages, tools) do
          {:error, reason} ->
            {:error, reason}

          parsed_data ->
            # Track the response being received
            LangChain.Telemetry.llm_response(
              %{system_time: System.system_time()},
              %{model: vertex_ai.model, response: parsed_data}
            )

            {:ok, parsed_data}
        end
      rescue
        err in LangChainError ->
          {:error, err}
      end
    end)
  end

  @doc false
  @spec do_api_request(t(), [Message.t()], [Function.t()]) ::
          list() | struct() | {:error, LangChainError.t()}
  def do_api_request(%ChatVertexAI{stream: false} = vertex_ai, messages, tools) do
    req =
      Req.new(
        url: build_url(vertex_ai),
        json: for_api(vertex_ai, messages, tools),
        receive_timeout: vertex_ai.receive_timeout,
        # Disable Req-level retry to prevent compounding with LangChain's own
        # :closed retry. See https://github.com/brainlid/langchain/issues/503
        retry: false,
        auth: {:bearer, get_api_key(vertex_ai)}
      )
      |> Req.merge(vertex_ai.req_config |> Keyword.new())

    req
    |> Req.post()
    |> case do
      {:ok, %Req.Response{body: data} = response} ->
        Callbacks.fire(vertex_ai.callbacks, :on_llm_response_headers, [response.headers])

        case do_process_response(vertex_ai, data) do
          {:error, reason} ->
            {:error, reason}

          result ->
            Callbacks.fire(vertex_ai.callbacks, :on_llm_new_message, [result])

            # Track non-streaming response completion
            LangChain.Telemetry.emit_event(
              [:langchain, :llm, :response, :non_streaming],
              %{system_time: System.system_time()},
              %{
                model: vertex_ai.model,
                response_size: byte_size(inspect(result))
              }
            )

            result
        end

      {:error, %Req.TransportError{reason: :timeout} = err} ->
        {:error,
         LangChainError.exception(type: "timeout", message: "Request timed out", original: err)}

      other ->
        Logger.warning(fn -> "Unexpected and unhandled API response! #{inspect(other)}" end)
        other
    end
  end

  def do_api_request(%ChatVertexAI{stream: true} = vertex_ai, messages, tools) do
    Req.new(
      url: build_url(vertex_ai),
      json: for_api(vertex_ai, messages, tools),
      auth: {:bearer, get_api_key(vertex_ai)},
      receive_timeout: vertex_ai.receive_timeout
    )
    |> Req.Request.put_header("accept-encoding", "utf-8")
    |> Req.merge(vertex_ai.req_config |> Keyword.new())
    |> Req.post(
      into:
        Utils.handle_stream_fn(
          vertex_ai,
          &ChatOpenAI.decode_stream/1,
          &do_process_response(vertex_ai, &1, MessageDelta)
        )
    )
    |> case do
      {:ok, %Req.Response{body: data} = response} ->
        Callbacks.fire(vertex_ai.callbacks, :on_llm_response_headers, [response.headers])

        # Google AI uses `finishReason: "STOP` for all messages in the stream.
        # This field can't be used to terminate the list of deltas, so simulate
        # this behavior by forcing the final delta to have `status: :complete`.
        complete_final_delta(data)

      {:error, %LangChainError{} = error} ->
        {:error, error}

      {:error, %Req.TransportError{reason: :timeout} = err} ->
        {:error,
         LangChainError.exception(type: "timeout", message: "Request timed out", original: err)}

      other ->
        Logger.warning(fn ->
          "Unhandled and unexpected response from streamed post call. #{inspect(other)}"
        end)

        {:error,
         LangChainError.exception(
           type: "unexpected_response",
           message: "Unexpected response",
           original: other
         )}
    end
  end

  @doc false
  @spec build_url(t()) :: String.t()
  def build_url(%ChatVertexAI{endpoint: endpoint, model: model} = vertex_ai) do
    "#{endpoint}/models/#{model}:#{get_action(vertex_ai)}"
    |> use_sse(vertex_ai)
  end

  @spec use_sse(String.t(), t()) :: String.t()
  defp use_sse(url, %ChatVertexAI{stream: true}), do: url <> "?alt=sse"
  defp use_sse(url, _model), do: url

  @spec get_action(t()) :: String.t()
  defp get_action(%ChatVertexAI{stream: false}), do: "generateContent"
  defp get_action(%ChatVertexAI{stream: true}), do: "streamGenerateContent"

  def complete_final_delta(data) when is_list(data) do
    update_in(data, [Access.at(-1), Access.at(-1)], &%{&1 | status: :complete})
  end

  def do_process_response(model, response, message_type \\ Message)

  def do_process_response(model, %{"candidates" => candidates} = data, message_type)
      when is_list(candidates) do
    token_usage = get_token_usage(data)

    case token_usage do
      %TokenUsage{} = usage ->
        Callbacks.fire(model.callbacks, :on_llm_token_usage, [usage])
        :ok

      nil ->
        :ok
    end

    candidates
    |> Enum.map(&do_process_response(model, &1, message_type))
    |> Enum.map(&TokenUsage.set(&1, token_usage))
  end

  def do_process_response(
        model,
        %{"content" => %{"parts" => parts} = content_data} = data,
        Message
      ) do
    text_part =
      parts
      |> filter_parts_for_types(["text"])
      |> filter_text_parts()
      |> Enum.map(fn part ->
        type =
          case part["thought"] do
            true -> :thinking
            _ -> :text
          end

        ContentPart.new!(%{type: type, content: part["text"]})
      end)

    tool_calls_from_parts =
      parts
      |> filter_parts_for_types(["functionCall"])
      |> Enum.map(fn part ->
        do_process_response(model, part, nil)
      end)

    tool_result_from_parts =
      parts
      |> filter_parts_for_types(["functionResponse"])
      |> Enum.map(fn part ->
        do_process_response(model, part, nil)
      end)

    %{
      role: unmap_role(content_data["role"]),
      content: text_part,
      complete: false,
      index: data["index"]
    }
    |> Utils.conditionally_add_to_map(:tool_calls, tool_calls_from_parts)
    |> Utils.conditionally_add_to_map(:tool_results, tool_result_from_parts)
    |> Message.new()
    |> case do
      {:ok, message} ->
        message

      {:error, %Ecto.Changeset{} = changeset} ->
        {:error, LangChainError.exception(changeset)}
    end
  end

  def do_process_response(
        model,
        %{"content" => %{"parts" => parts} = content_data} = data,
        MessageDelta
      ) do
    text_content =
      case parts do
        [%{"text" => text}] ->
          text

        _other ->
          nil
      end

    parts
    |> filter_parts_for_types(["text"])
    |> Enum.map(fn part ->
      ContentPart.new!(%{type: :text, content: part["text"]})
    end)

    tool_calls_from_parts =
      parts
      |> filter_parts_for_types(["functionCall"])
      |> Enum.map(fn part ->
        do_process_response(model, part, nil)
      end)

    %{
      role: unmap_role(content_data["role"]),
      content: text_content,
      complete: true,
      index: data["index"]
    }
    |> Utils.conditionally_add_to_map(:tool_calls, tool_calls_from_parts)
    |> MessageDelta.new()
    |> case do
      {:ok, message} ->
        message

      {:error, %Ecto.Changeset{} = changeset} ->
        {:error, LangChainError.exception(changeset)}
    end
  end

  def do_process_response(
        _model,
        %{"functionCall" => %{"args" => raw_args, "name" => name}} = data,
        _
      ) do
    %{
      call_id: "call-#{name}",
      name: name,
      arguments: raw_args,
      complete: true,
      index: data["index"],
      metadata:
        if(data["thoughtSignature"],
          do: %{thought_signature: data["thoughtSignature"]},
          else: nil
        )
    }
    |> ToolCall.new()
    |> case do
      {:ok, message} ->
        message

      {:error, %Ecto.Changeset{} = changeset} ->
        {:error, LangChainError.exception(changeset)}
    end
  end

  def do_process_response(
        _model,
        %{
          "finishReason" => finish,
          "content" => %{"parts" => parts, "role" => role},
          "index" => index
        },
        message_type
      )
      when is_list(parts) do
    status =
      case message_type do
        MessageDelta ->
          :incomplete

        Message ->
          case finish do
            "STOP" ->
              :complete

            "SAFETY" ->
              :complete

            other ->
              Logger.warning("Unsupported finishReason in response. Reason: #{inspect(other)}")
              nil
          end
      end

    content = Enum.map_join(parts, & &1["text"])

    case message_type.new(%{
           "content" => content,
           "role" => unmap_role(role),
           "status" => status,
           "index" => index
         }) do
      {:ok, message} ->
        message

      {:error, %Ecto.Changeset{} = changeset} ->
        {:error, LangChainError.exception(changeset)}
    end
  end

  def do_process_response(_model, %{"error" => %{"message" => reason}} = response, _) do
    {:error, LangChainError.exception(message: reason, original: response)}
  end

  def do_process_response(_model, {:error, %Jason.DecodeError{} = response}, _) do
    error_message = "Received invalid JSON: #{inspect(response)}"

    {:error,
     LangChainError.exception(type: "invalid_json", message: error_message, original: response)}
  end

  def do_process_response(_model, other, _) do
    {:error,
     LangChainError.exception(
       type: "unexpected_response",
       message: "Unexpected response",
       original: other
     )}
  end

  @doc false
  def filter_parts_for_types(parts, types) when is_list(parts) and is_list(types) do
    Enum.filter(parts, fn p ->
      Enum.any?(types, &Map.has_key?(p, &1))
    end)
  end

  @doc false
  def filter_text_parts(parts) when is_list(parts) do
    Enum.filter(parts, fn p ->
      case p do
        %{"text" => text} -> text && text != ""
        _ -> false
      end
    end)
  end

  @doc """
  Return the content parts for the message.
  """
  @spec get_message_contents(MessageDelta.t() | Message.t()) ::
          [%{String.t() => any()}] | nil
  def get_message_contents(%{content: content} = _message) when is_binary(content) do
    [%{"text" => content}]
  end

  def get_message_contents(%{content: contents} = _message) when is_list(contents) do
    Enum.map(contents, &for_api/1)
  end

  def get_message_contents(%{content: nil} = _message) do
    nil
  end

  defp map_role(role) do
    case role do
      :assistant -> :model
      :tool -> :function
      # System prompts are not supported yet. Google recommends using user prompt.
      :system -> :user
      role -> role
    end
  end

  defp get_token_usage(%{"usageMetadata" => usage} = _response_body) do
    # extract out the reported response token usage
    TokenUsage.new!(%{
      input: Map.get(usage, "promptTokenCount", 0),
      output: Map.get(usage, "candidatesTokenCount", 0),
      raw: usage
    })
  end

  defp get_token_usage(_response_body), do: nil

  defp unmap_role("model"), do: "assistant"
  defp unmap_role("function"), do: "tool"
  defp unmap_role(role), do: role

  @impl ChatModel
  def provider, do: "vertex_ai"

  @doc """
  Determine if an error should be retried. If `true`, a fallback LLM may be
  used. If `false`, the error is understood to be more fundamental with the
  request rather than a service issue and it should not be retried or fallback
  to another service.
  """
  @impl ChatModel
  @spec retry_on_fallback?(LangChainError.t()) :: boolean()
  def retry_on_fallback?(%LangChainError{type: "rate_limited"}), do: true
  def retry_on_fallback?(%LangChainError{type: "rate_limit_exceeded"}), do: true
  def retry_on_fallback?(%LangChainError{type: "timeout"}), do: true
  def retry_on_fallback?(%LangChainError{type: "too_many_requests"}), do: true
  def retry_on_fallback?(_), do: false

  @doc """
  Generate a config map that can later restore the model's configuration.
  """
  @impl ChatModel
  @spec serialize_config(t()) :: %{String.t() => any()}
  def serialize_config(%ChatVertexAI{} = model) do
    Utils.to_serializable_map(
      model,
      [
        :endpoint,
        :model,
        :temperature,
        :top_p,
        :top_k,
        :thinking_config,
        :receive_timeout,
        :json_response,
        :json_schema,
        :stream,
        :safety_settings
      ],
      @current_config_version
    )
  end

  @doc """
  Restores the model from the config.
  """
  @impl ChatModel
  def restore_from_map(%{"version" => 1} = data) do
    ChatVertexAI.new(data)
  end
end