nous

defmodule Nous.StreamNormalizer.Gemini do
  @moduledoc """
  Stream normalizer for Google Gemini API.

  Handles Gemini's SSE streaming format where events arrive as JSON maps
  with a `candidates` array.

  ## Streaming Event Types

  | Gemini Event | Normalized Output |
  |-------------|-------------------|
  | `candidates[0].content.parts[0].text` | `{:text_delta, text}` |
  | `candidates[0].content.parts` with `functionCall` | `{:tool_call_delta, call}` |
  | `candidates[0].finishReason` present | `{:finish, reason}` |
  | `{:stream_done, reason}` | `{:finish, reason}` |
  | Error with `"error"` key | `{:error, message}` |
  """

  @behaviour Nous.StreamNormalizer

  @impl true
  def normalize_chunk({:stream_done, reason}) do
    [{:finish, reason}]
  end

  def normalize_chunk(%{"error" => error}) when is_map(error) do
    message = Map.get(error, "message", inspect(error))
    [{:error, message}]
  end

  def normalize_chunk(%{"candidates" => candidates} = chunk) when is_list(candidates) do
    usage_events = maybe_usage_event(chunk)

    case candidates do
      [candidate | _] -> parse_candidate(candidate, chunk) ++ usage_events
      [] -> if usage_events == [], do: [{:unknown, chunk}], else: usage_events
    end
  end

  def normalize_chunk(chunk) do
    [{:unknown, chunk}]
  end

  @impl true
  def complete_response?(%{"candidates" => [candidate | _]}) when is_map(candidate) do
    Map.has_key?(candidate, "finishReason")
  end

  def complete_response?(_), do: false

  @impl true
  def convert_complete_response(%{"candidates" => [candidate | _]}) do
    parts = get_in(candidate, ["content", "parts"]) || []
    finish_reason = normalize_finish_reason(Map.get(candidate, "finishReason"))

    events = Enum.flat_map(parts, &parse_part/1)
    events ++ [{:finish, finish_reason}]
  end

  def convert_complete_response(chunk) do
    [{:unknown, chunk}]
  end

  defp maybe_usage_event(chunk) do
    case Map.get(chunk, "usageMetadata") do
      nil -> []
      usage -> [{:usage, Nous.Messages.Gemini.parse_usage(usage)}]
    end
  end

  defp parse_candidate(candidate, _chunk) do
    parts = get_in(candidate, ["content", "parts"]) || []
    finish_reason = Map.get(candidate, "finishReason")

    part_events = Enum.flat_map(parts, &parse_part/1)

    if finish_reason do
      part_events ++ [{:finish, normalize_finish_reason(finish_reason)}]
    else
      case part_events do
        [] -> [{:unknown, candidate}]
        events -> events
      end
    end
  end

  defp parse_part(%{"text" => text} = part) when text != "" do
    if Map.get(part, "thought") do
      [{:thinking_delta, text}]
    else
      [{:text_delta, text}]
    end
  end

  defp parse_part(%{"functionCall" => %{"name" => name, "args" => args}} = part) do
    [{:tool_call_delta, attach_thought_signature(%{"name" => name, "arguments" => args}, part)}]
  end

  defp parse_part(%{"functionCall" => %{"name" => name}} = part) do
    [{:tool_call_delta, attach_thought_signature(%{"name" => name, "arguments" => %{}}, part)}]
  end

  defp parse_part(_), do: []

  defp attach_thought_signature(delta, part) do
    case Map.get(part, "thoughtSignature") do
      sig when is_binary(sig) and sig != "" ->
        Map.put(delta, "metadata", %{"thought_signature" => sig})

      _ ->
        delta
    end
  end

  defp normalize_finish_reason("STOP"), do: "stop"
  defp normalize_finish_reason("MAX_TOKENS"), do: "length"
  defp normalize_finish_reason("SAFETY"), do: "safety"
  defp normalize_finish_reason(reason) when is_binary(reason), do: String.downcase(reason)
  defp normalize_finish_reason(_), do: "stop"
end