defmodule LangChain.ChatModels.ChatPerplexity do
  @moduledoc """
  Represents the [Perplexity Chat model](https://docs.perplexity.ai/api-reference/chat-completions).

  This module implements a client for the Perplexity Chat API, providing functions to validate input parameters,
  format API requests, and parse API responses into LangChain's structured data types.

  Perplexity does not natively support tool calling in the same manner as some other chat models.
  To overcome this limitation, this module employs a workaround using structured outputs via a JSON schema.
  When tools are provided, the API request is augmented with a JSON schema that defines the expected format
  for tool calls. The response processing logic then detects and decodes these tool call details, converting them
  into corresponding ToolCall structs. This approach allows LangChain to seamlessly emulate tool calling functionality
  and integrate it with its standard workflow, similar to how ChatOpenAI handles function calls.

  In addition, this module supports various configuration options such as temperature, top_p, top_k,
  and streaming, as well as callbacks for token usage and new message events.

  Overall, this implementation provides a unified interface for interacting with the Perplexity Chat API
  while working around its limitations regarding tool calling.

  ## Tool Calls

  In order to use tool calls, you need specifically prompt Perplexity as outlined in their
  [Prompt Guide](https://docs.perplexity.ai/guides/prompt-guide) as well as the
  [Structured Outputs Guide](https://docs.perplexity.ai/guides/structured-outputs).

  Provide it additional prompting like:

  ```
  Rules:
  1. Provide only the final answer. It is important that you do not include any explanation on the steps below.
  2. Do not show the intermediate steps information.

  Output a JSON object with the following fields:
  - title: The article title
  - keywords: An array of SEO keywords
  - meta_description: The SEO meta description
  ```
  """
  use Ecto.Schema
  require Logger
  import Ecto.Changeset
  alias __MODULE__
  alias LangChain.Config
  alias LangChain.ChatModels.ChatModel
  alias LangChain.Message
  alias LangChain.MessageDelta
  alias LangChain.Message.Citation
  alias LangChain.Message.ContentPart
  alias LangChain.Message.ToolCall
  alias LangChain.TokenUsage
  alias LangChain.LangChainError
  alias LangChain.Utils
  alias LangChain.Callbacks
  alias LangChain.Telemetry

  @behaviour ChatModel

  @current_config_version 1

  # Default endpoint for Perplexity API
  @default_endpoint "https://api.perplexity.ai/chat/completions"

  # Default timeout of 1 minute
  @receive_timeout 60_000

  @primary_key false
  embedded_schema do
    field :endpoint, :string, default: @default_endpoint
    field :model, :string, default: "sonar-reasoning-pro"
    field :api_key, :string

    # What sampling temperature to use, between 0 and 2.
    # Higher values make output more random, lower values more deterministic.
    field :temperature, :float, default: 0.2

    # The nucleus sampling threshold, between 0 and 1.
    field :top_p, :float, default: 0.9

    # The number of tokens for highest top-k filtering (0-2048).
    field :top_k, :integer, default: 0

    # Maximum number of tokens to generate
    field :max_tokens, :integer

    # Whether to stream the response
    field :stream, :boolean, default: false

    # Presence penalty between -2.0 and 2.0
    field :presence_penalty, :float, default: 0.0

    # Frequency penalty greater than 0
    field :frequency_penalty, :float, default: 1.0

    # Search domain filter for limiting citations
    field :search_domain_filter, {:array, :string}

    # Whether to return images in response
    field :return_images, :boolean, default: false

    # Whether to return related questions
    field :return_related_questions, :boolean, default: false

    # Time interval for search recency
    field :search_recency_filter, :string

    # Response format for structured outputs
    field :response_format, :map

    # Duration in seconds for response timeout
    field :receive_timeout, :integer, default: @receive_timeout

    # A list of maps for callback handlers
    field :callbacks, {:array, :map}, default: []

    # For help with debugging. It outputs the RAW Req response received and the
    # RAW Elixir map being submitted to the API.
    field :verbose_api, :boolean, default: false
  end

  @type t :: %ChatPerplexity{}

  @create_fields [
    :endpoint,
    :model,
    :api_key,
    :temperature,
    :top_p,
    :top_k,
    :max_tokens,
    :stream,
    :presence_penalty,
    :frequency_penalty,
    :search_domain_filter,
    :return_images,
    :return_related_questions,
    :search_recency_filter,
    :response_format,
    :receive_timeout,
    :verbose_api
  ]

  @required_fields [:model]

  @spec get_api_key(t()) :: String.t()
  defp get_api_key(%ChatPerplexity{api_key: api_key}) do
    api_key || Config.resolve(:perplexity_key, "")
  end

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

  @doc """
  Setup a ChatPerplexity 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)
    |> validate_number(:temperature, greater_than_or_equal_to: 0, less_than: 2)
    |> validate_number(:top_p, greater_than: 0, less_than_or_equal_to: 1)
    |> validate_number(:top_k, greater_than_or_equal_to: 0, less_than_or_equal_to: 2048)
    |> validate_number(:presence_penalty, greater_than_or_equal_to: -2, less_than_or_equal_to: 2)
    |> validate_number(:frequency_penalty, greater_than: 0)
    |> validate_number(:receive_timeout, greater_than_or_equal_to: 0)
  end

  @doc """
  Return the params formatted for an API request.
  """
  @spec for_api(t(), [Message.t()], ChatModel.tools()) :: %{atom() => any()}
  def for_api(%ChatPerplexity{} = perplexity, messages, tools) do
    response_format =
      if tools && !Enum.empty?(tools) do
        %{
          "type" => "json_schema",
          "json_schema" => %{
            "schema" => %{
              "type" => "object",
              "required" => ["tool_calls"],
              "properties" => %{
                "tool_calls" => %{
                  "type" => "array",
                  "items" => %{
                    "type" => "object",
                    "required" => ["name", "arguments"],
                    "properties" => %{
                      "name" => %{
                        "type" => "string",
                        "enum" => Enum.map(tools, & &1.name)
                      },
                      "arguments" => build_arguments_schema(tools)
                    }
                  }
                }
              }
            }
          }
        }
      else
        perplexity.response_format
      end

    %{
      model: perplexity.model,
      messages: Enum.map(messages, &for_api(perplexity, &1)),
      temperature: perplexity.temperature,
      top_p: perplexity.top_p,
      top_k: perplexity.top_k,
      stream: perplexity.stream
    }
    |> Utils.conditionally_add_to_map(:max_tokens, perplexity.max_tokens)
    |> Utils.conditionally_add_to_map(:presence_penalty, perplexity.presence_penalty)
    |> Utils.conditionally_add_to_map(:frequency_penalty, perplexity.frequency_penalty)
    |> Utils.conditionally_add_to_map(:search_domain_filter, perplexity.search_domain_filter)
    |> Utils.conditionally_add_to_map(:return_images, perplexity.return_images)
    |> Utils.conditionally_add_to_map(
      :return_related_questions,
      perplexity.return_related_questions
    )
    |> Utils.conditionally_add_to_map(:search_recency_filter, perplexity.search_recency_filter)
    |> Utils.conditionally_add_to_map(:response_format, response_format)
  end

  defp build_arguments_schema([tool | _]) do
    properties =
      tool.parameters
      |> Enum.map(fn param ->
        {param.name, param_to_json_schema(param)}
      end)
      |> Map.new()

    required =
      tool.parameters
      |> Enum.filter(& &1.required)
      |> Enum.map(& &1.name)

    %{
      "type" => "object",
      "required" => required,
      "properties" => properties
    }
  end

  defp param_to_json_schema(param) do
    base = %{"type" => atom_to_json_type(param.type)}

    base
    |> add_enum(param)
    |> add_items(param)
  end

  defp atom_to_json_type(:string), do: "string"
  defp atom_to_json_type(:number), do: "number"
  defp atom_to_json_type(:integer), do: "integer"
  defp atom_to_json_type(:boolean), do: "boolean"
  defp atom_to_json_type(:array), do: "array"
  defp atom_to_json_type(:object), do: "object"

  defp add_enum(schema, %{enum: enum}) when enum in [nil, []], do: schema
  defp add_enum(schema, %{enum: enum}), do: Map.put(schema, "enum", enum)

  defp add_items(schema, %{type: :array, item_type: item_type}) do
    Map.put(schema, "items", %{"type" => atom_to_json_type(String.to_existing_atom(item_type))})
  end

  defp add_items(schema, _), do: schema

  @doc """
  Convert a LangChain Message-based structure to the expected map of data for
  the Perplexity API.
  """
  @spec for_api(t(), Message.t()) :: %{String.t() => any()}
  def for_api(%ChatPerplexity{}, %Message{} = msg) do
    content =
      case msg.content do
        content when is_binary(content) -> content
        content when is_list(content) -> ContentPart.parts_to_string(content)
        nil -> nil
      end

    %{
      "role" => msg.role,
      "content" => content
    }
  end

  @impl ChatModel
  def call(perplexity, prompt, tools \\ [])

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

    call(perplexity, messages, tools)
  end

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

    Telemetry.span([:langchain, :llm, :call], metadata, fn ->
      try do
        # Track the prompt being sent
        Telemetry.llm_prompt(
          %{system_time: System.system_time()},
          %{model: perplexity.model, messages: messages}
        )

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

          {:ok, %Req.Response{body: body}} when is_list(body) ->
            # Streaming response - body contains accumulated MessageDelta lists
            Telemetry.llm_response(
              %{system_time: System.system_time()},
              %{model: perplexity.model, response: body}
            )

            {:ok, body}

          parsed_data ->
            # Non-streaming response
            Telemetry.llm_response(
              %{system_time: System.system_time()},
              %{model: perplexity.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()], ChatModel.tools(), integer()) ::
          list() | struct() | {:error, LangChainError.t()}
  def do_api_request(perplexity, messages, tools, retry_count \\ 3)

  def do_api_request(_perplexity, _messages, _tools, 0) do
    raise LangChainError, "Retries exceeded. Connection failed."
  end

  def do_api_request(
        %ChatPerplexity{stream: false} = perplexity,
        messages,
        tools,
        retry_count
      ) do
    req =
      Req.new(
        url: perplexity.endpoint,
        json: for_api(perplexity, messages, tools),
        auth: {:bearer, get_api_key(perplexity)},
        receive_timeout: perplexity.receive_timeout
      )

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

        Callbacks.fire(perplexity.callbacks, :on_llm_token_usage, [
          get_token_usage(data)
        ])

        case do_process_response(perplexity, data, tools) do
          {:error, %LangChainError{} = reason} ->
            Logger.error("Error processing response: #{inspect(reason)}")
            {:error, reason}

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

            # Track non-streaming response completion
            Telemetry.emit_event(
              [:langchain, :llm, :response, :non_streaming],
              %{system_time: System.system_time()},
              %{
                model: perplexity.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)}

      {:error, %Req.TransportError{reason: :closed}} ->
        Logger.debug(fn -> "Connection closed: retry count = #{inspect(retry_count)}" end)
        do_api_request(perplexity, messages, tools, retry_count - 1)

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

  def do_api_request(
        %ChatPerplexity{stream: true} = perplexity,
        messages,
        tools,
        retry_count
      ) do
    Req.new(
      url: perplexity.endpoint,
      json: for_api(perplexity, messages, tools),
      auth: {:bearer, get_api_key(perplexity)},
      receive_timeout: perplexity.receive_timeout
    )
    |> Req.post(
      into:
        Utils.handle_stream_fn(
          perplexity,
          &decode_stream/1,
          &process_stream_chunk(perplexity, &1)
        )
    )
    |> case do
      {:ok, response} ->
        Callbacks.fire(perplexity.callbacks, :on_llm_response_headers, [response.headers])

        {:ok, response}

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

      {:error, %Req.TransportError{reason: :closed}} ->
        Logger.debug(fn -> "Connection closed: retry count = #{inspect(retry_count)}" end)
        do_api_request(perplexity, messages, tools, retry_count - 1)

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

  @doc """
  Decode a streamed response from the Perplexity API.
  """
  @spec decode_stream({String.t(), String.t()}) :: {%{String.t() => any()}}
  def decode_stream({raw_data, buffer}, done \\ []) do
    raw_data
    |> String.split("data: ")
    |> Enum.reduce({done, buffer}, fn str, {done, incomplete} = acc ->
      str
      |> String.trim()
      |> case do
        "" ->
          acc

        "[DONE]" ->
          acc

        json ->
          parse_combined_data(incomplete, json, done)
      end
    end)
  end

  defp parse_combined_data("", json, done) do
    json
    |> Jason.decode()
    |> case do
      {:ok, parsed} ->
        {done ++ [parsed], ""}

      {:error, _reason} ->
        {done, json}
    end
  end

  defp parse_combined_data(incomplete, json, done) do
    starting_json = incomplete <> json
    decode_stream({starting_json, ""}, done)
  end

  @doc false
  @spec do_process_response(t(), map(), ChatModel.tools()) ::
          Message.t() | {:error, LangChainError.t()}
  def do_process_response(model, %{"choices" => [choice | _]} = data, tools) do
    # Fire token usage callback if present
    if usage = Map.get(data, "usage") do
      case get_token_usage(%{"usage" => usage}) do
        %TokenUsage{} = token_usage ->
          Callbacks.fire(model.callbacks, :on_llm_token_usage, [token_usage])

        nil ->
          :ok
      end
    end

    # Process the first choice and attach citations from top-level response
    model
    |> do_process_response(choice, tools)
    |> attach_perplexity_citations(data)
  end

  def do_process_response(
        _model,
        %{"finish_reason" => finish_reason, "message" => %{"content" => content, "role" => role}} =
          data,
        tools
      )
      when tools != [] do
    status = finish_reason_to_status(finish_reason)

    # Try to parse content as JSON since we expect structured output.
    # Perplexity returns structured output matching the tool_calls JSON schema
    # we send in for_api: {"tool_calls": [{"name": "...", "arguments": {...}}]}
    case Jason.decode(content) do
      {:ok, %{"tool_calls" => tool_calls}} when is_list(tool_calls) ->
        parsed_tool_calls =
          Enum.map(tool_calls, fn call ->
            args_json = Jason.encode!(call["arguments"])

            tool_call =
              ToolCall.new!(%{
                type: :function,
                status: :complete,
                name: call["name"],
                arguments: args_json,
                call_id: Ecto.UUID.generate()
              })

            # Force arguments back to JSON string (validate_and_parse_arguments
            # decodes it to a map, but we want to keep the string form)
            %{tool_call | arguments: args_json}
          end)

        case Message.new(%{
               "role" => role,
               "content" => nil,
               "status" => status,
               "index" => data["index"],
               "tool_calls" => parsed_tool_calls
             }) do
          {:ok, message} -> message
          {:error, changeset} -> {:error, LangChainError.exception(changeset)}
        end

      {:ok, _parsed} ->
        # Valid JSON but not in tool_calls format — treat as single tool call
        tool_call =
          ToolCall.new!(%{
            type: :function,
            status: :complete,
            name: List.first(tools).name,
            arguments: content,
            call_id: Ecto.UUID.generate()
          })

        case Message.new(%{
               "role" => role,
               "content" => nil,
               "status" => status,
               "index" => data["index"],
               "tool_calls" => [tool_call]
             }) do
          {:ok, message} -> message
          {:error, changeset} -> {:error, LangChainError.exception(changeset)}
        end

      {:error, _} ->
        {:error,
         LangChainError.exception(
           type: "bad_request",
           message: "response_format: expected JSON structured output but got: #{content}"
         )}
    end
  end

  def do_process_response(
        _model,
        %{"finish_reason" => finish_reason, "message" => %{"content" => content, "role" => role}} =
          data,
        tools
      )
      when tools == [] do
    status = finish_reason_to_status(finish_reason)

    case Message.new(%{
           "role" => role,
           "content" => content,
           "status" => status,
           "index" => data["index"]
         }) do
      {:ok, message} -> message
      {:error, changeset} -> {:error, LangChainError.exception(changeset)}
    end
  end

  def do_process_response(_model, %{"error" => %{"message" => reason, "type" => type}}) do
    Logger.error("Received error from API: #{inspect(reason)}")
    {:error, LangChainError.exception(type: type, message: reason)}
  end

  def do_process_response(_model, %{"error" => %{"message" => reason}}) do
    Logger.error("Received error from API: #{inspect(reason)}")
    {:error, LangChainError.exception(message: reason)}
  end

  def do_process_response(model, %{"choices" => %{} = usage} = _data) do
    case get_token_usage(%{"usage" => usage}) do
      %TokenUsage{} = token_usage ->
        Callbacks.fire(model.callbacks, :on_llm_token_usage, [token_usage])
        :skip

      nil ->
        :skip
    end
  end

  def do_process_response(_model, %{"choices" => []}), do: :skip

  def do_process_response(model, %{"choices" => choices} = data) when is_list(choices) do
    # Fire token usage callback if present
    if usage = Map.get(data, "usage") do
      case get_token_usage(%{"usage" => usage}) do
        %TokenUsage{} = token_usage ->
          Callbacks.fire(model.callbacks, :on_llm_token_usage, [token_usage])

        nil ->
          :ok
      end
    end

    # Process each response individually
    for choice <- choices do
      do_process_response(model, choice)
    end
  end

  def do_process_response(
        _model,
        %{"finish_reason" => finish_reason, "message" => %{"content" => content}} = data
      ) do
    status = finish_reason_to_status(finish_reason)

    # Try to parse content as JSON for potential tool calls
    case Jason.decode(content) do
      {:ok, %{"tool_calls" => tool_calls}} when is_list(tool_calls) ->
        # Convert JSON tool calls to Message struct with tool calls
        case Message.new(%{
               "role" => :assistant,
               "content" => nil,
               "status" => status,
               "index" => data["index"],
               "tool_calls" =>
                 Enum.map(tool_calls, fn call ->
                   tool_call =
                     ToolCall.new!(%{
                       type: :function,
                       status: :complete,
                       name: call["name"],
                       arguments: Jason.encode!(call["arguments"]),
                       call_id: Ecto.UUID.generate()
                     })

                   # Force the arguments field to be a JSON string even if the ToolCall schema casts it
                   %{tool_call | arguments: Jason.encode!(call["arguments"])}
                 end)
             }) do
          {:ok, message} -> message
          {:error, changeset} -> {:error, LangChainError.exception(changeset)}
        end

      _ ->
        # Regular message processing
        case Message.new(%{
               "role" => :assistant,
               "content" => content,
               "status" => status,
               "index" => data["index"]
             }) do
          {:ok, message} -> message
          {:error, changeset} -> {:error, LangChainError.exception(changeset)}
        end
    end
  end

  def do_process_response(_model, %{
        "choices" => [
          %{
            "delta" => %{"role" => role, "content" => content},
            "finish_reason" => finish,
            "index" => index
          } = _choice
        ]
      }) do
    status = finish_reason_to_status(finish)

    data =
      %{}
      |> Map.put("role", role)
      |> Map.put("content", content)
      |> Map.put("index", index)
      |> Map.put("status", status)

    case MessageDelta.new(data) do
      {:ok, message} ->
        send(self(), {:message_delta, message})
        message

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

  def do_process_response(
        _model,
        %{
          "choices" => [
            %{
              "delta" => %{"content" => content},
              "finish_reason" => finish,
              "index" => index
            }
            | _
          ]
        }
      ) do
    status = finish_reason_to_status(finish)

    data =
      %{}
      |> Map.put("role", "assistant")
      |> Map.put("content", content)
      |> Map.put("index", index)
      |> Map.put("status", status)

    case MessageDelta.new(data) do
      {:ok, message} ->
        send(self(), {:message_delta, message})
        message

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

  def do_process_response(_model, %{"choices" => []} = _msg), do: :skip

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

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

  def do_process_response(_model, other) do
    Logger.error("Trying to process an unexpected response. #{inspect(other)}")

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

  defp finish_reason_to_status(nil), do: :incomplete
  defp finish_reason_to_status("stop"), do: :complete
  defp finish_reason_to_status("length"), do: :length

  defp finish_reason_to_status(other) do
    Logger.warning("Unsupported finish_reason in message. Reason: #{inspect(other)}")
    nil
  end

  defp get_token_usage(%{"usage" => usage} = _response_body) do
    TokenUsage.new!(%{
      input: Map.get(usage, "prompt_tokens"),
      output: Map.get(usage, "completion_tokens"),
      raw: usage
    })
  end

  defp get_token_usage(_response_body), do: nil

  @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(%ChatPerplexity{} = model) do
    model
    |> Utils.to_serializable_map(
      [
        :endpoint,
        :model,
        :temperature,
        :top_p,
        :top_k,
        :max_tokens,
        :stream,
        :presence_penalty,
        :frequency_penalty,
        :search_domain_filter,
        :return_images,
        :return_related_questions,
        :search_recency_filter,
        :response_format,
        :receive_timeout,
        :verbose_api
      ],
      @current_config_version
    )
    |> Map.delete("module")
  end

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

  # Content delta on the final chunk (finish_reason: "stop"). Also extract
  # citations from the chunk since Perplexity includes them on every chunk
  # but we only need to process them once at completion. Returns a list
  # [citation_delta, content_delta] so both flow through the pipeline.
  defp process_stream_chunk(
         _perplexity,
         %{
           "choices" => [
             %{
               "delta" => %{"content" => content},
               "finish_reason" => "stop"
             }
             | _
           ]
         } = chunk
       ) do
    content_delta = %MessageDelta{content: content, role: :assistant, status: :complete}

    case build_streaming_citation_delta(chunk) do
      nil -> content_delta
      citation_delta -> [citation_delta, content_delta]
    end
  end

  # Content delta on intermediate chunks (no finish_reason or nil).
  # Ignore citations here since they repeat on every chunk.
  defp process_stream_chunk(_perplexity, %{
         "choices" => [
           %{"delta" => %{"content" => content}} | _
         ]
       }) do
    %MessageDelta{content: content, role: :assistant}
  end

  # Completion chunk with no content delta.
  defp process_stream_chunk(
         _perplexity,
         %{
           "choices" => [
             %{"finish_reason" => "stop"} | _
           ]
         } = chunk
       ) do
    content_delta = %MessageDelta{status: :complete, role: :assistant}

    case build_streaming_citation_delta(chunk) do
      nil -> content_delta
      citation_delta -> [citation_delta, content_delta]
    end
  end

  defp process_stream_chunk(_perplexity, _chunk) do
    %MessageDelta{}
  end

  # --- Citation Support ---

  # Attach citations from top-level Perplexity response data to a Message.
  @doc false
  @spec attach_perplexity_citations(Message.t() | any(), map()) :: Message.t() | any()
  def attach_perplexity_citations(%Message{} = msg, data) do
    citation_urls = Map.get(data, "citations", [])
    search_results = Map.get(data, "search_results", [])

    sources = build_perplexity_sources(citation_urls, search_results)

    if sources == [] do
      msg
    else
      full_text = get_message_text(msg)
      citations = find_citation_markers(full_text, sources)

      if citations == [] do
        msg
      else
        updated_content = attach_citations_to_first_text_part(msg.content, citations)
        %{msg | content: updated_content}
      end
    end
  end

  def attach_perplexity_citations(other, _data), do: other

  @doc false
  @spec build_perplexity_sources(list(), list()) :: [map()]
  def build_perplexity_sources(citation_urls, search_results) do
    cond do
      is_list(search_results) and search_results != [] ->
        search_results
        |> Enum.with_index()
        |> Enum.map(fn {result, idx} ->
          url = Map.get(result, "url") || Enum.at(citation_urls || [], idx)

          metadata =
            %{}
            |> maybe_put_metadata("date", Map.get(result, "date"))
            |> maybe_put_metadata("snippet", Map.get(result, "snippet"))
            |> maybe_put_metadata("source_domain", Map.get(result, "source"))

          %{
            type: :web,
            title: Map.get(result, "title"),
            url: url,
            metadata: metadata
          }
        end)

      is_list(citation_urls) and citation_urls != [] ->
        Enum.map(citation_urls, fn url ->
          %{type: :web, url: url}
        end)

      true ->
        []
    end
  end

  @doc false
  @spec find_citation_markers(String.t(), [map()]) :: [Citation.t()]
  def find_citation_markers(text, sources) when is_binary(text) and text != "" do
    ~r/\[(\d+)\]/
    |> Regex.scan(text, return: :index)
    |> Enum.flat_map(fn [{full_start, full_len}, {num_start, num_len}] ->
      number_str = binary_part(text, num_start, num_len)
      index = String.to_integer(number_str) - 1

      case Enum.at(sources, index) do
        nil ->
          []

        source ->
          [
            Citation.new!(%{
              start_index: full_start,
              end_index: full_start + full_len,
              source: source,
              metadata: %{"provider_type" => "perplexity_citation", "citation_index" => index}
            })
          ]
      end
    end)
  end

  def find_citation_markers(_, _), do: []

  defp get_message_text(%Message{content: content}) when is_list(content) do
    content
    |> Enum.filter(&(&1.type == :text))
    |> Enum.map_join("", & &1.content)
  end

  defp get_message_text(%Message{content: content}) when is_binary(content), do: content
  defp get_message_text(_), do: ""

  defp attach_citations_to_first_text_part(content, citations) when is_list(content) do
    case Enum.find_index(content, &(&1.type == :text)) do
      nil ->
        content

      idx ->
        List.update_at(content, idx, fn part ->
          %{part | citations: citations}
        end)
    end
  end

  defp attach_citations_to_first_text_part(content, _citations), do: content

  # Build a citation-only MessageDelta from a streaming chunk's citation data.
  # Returns nil if no citations are present.
  defp build_streaming_citation_delta(chunk) do
    citation_urls = Map.get(chunk, "citations", [])
    search_results = Map.get(chunk, "search_results", [])
    sources = build_perplexity_sources(citation_urls, search_results)

    if sources == [] do
      nil
    else
      citation_structs =
        sources
        |> Enum.with_index()
        |> Enum.map(fn {source, idx} ->
          Citation.new!(%{
            source: source,
            metadata: %{"provider_type" => "perplexity_citation", "citation_index" => idx}
          })
        end)

      citation_part = %ContentPart{type: :text, content: nil, citations: citation_structs}
      %MessageDelta{content: citation_part, role: :assistant, status: :incomplete}
    end
  end

  defp maybe_put_metadata(map, _key, nil), do: map
  defp maybe_put_metadata(map, key, value), do: Map.put(map, key, value)
end