puck

if Code.ensure_loaded?(ReqLLM) do
  defmodule Puck.Backends.ReqLLM do
    @moduledoc """
    Backend implementation using the ReqLLM library.

    ReqLLM provides a unified interface to multiple LLM providers.

    ## Model Specification

    The model must include the provider in `"provider:model"` format:

        model: "anthropic:claude-sonnet-4-5"

    ## Configuration

    Backend options are passed in the client tuple and forwarded to ReqLLM:

    - `:temperature` - Sampling temperature (0.0 to 2.0)
    - `:max_tokens` - Maximum tokens in response
    - `:top_p` - Nucleus sampling parameter
    - `:stop` - Stop sequences

    ## Examples

        # Basic usage
        client = Puck.Client.new({Puck.Backends.ReqLLM, "anthropic:claude-sonnet-4-5"})

        # With options
        client = Puck.Client.new(
          {Puck.Backends.ReqLLM, model: "anthropic:claude-sonnet-4-5", temperature: 0.7}
        )

    See ReqLLM documentation for supported providers and additional options.
    """

    @behaviour Puck.Backend

    alias Puck.Content.Part
    alias Puck.{Message, Response}
    alias ReqLLM.Message.ContentPart

    @impl true
    def call(config, messages, opts) do
      model = Map.fetch!(config, :model)
      output_schema = Keyword.get(opts, :output_schema)
      backend_opts = Keyword.get(opts, :backend_opts, [])
      options = build_options(config, backend_opts)
      req_messages = Enum.map(messages, &to_req_llm_message/1)

      if output_schema do
        call_with_schema(model, req_messages, output_schema, options)
      else
        call_text(model, req_messages, options)
      end
    end

    defp call_text(model, req_messages, options) do
      case ReqLLM.generate_text(model, req_messages, options) do
        {:ok, response} -> {:ok, normalize_response(response, model)}
        {:error, reason} -> {:error, reason}
      end
    end

    defp call_with_schema(model, messages, output_schema, options) do
      # Convert struct schemas to object schemas for LLM (struct schemas can't be JSON encoded)
      llm_schema = to_llm_schema(output_schema)

      case ReqLLM.generate_object(model, messages, llm_schema, options) do
        {:ok, response} -> {:ok, normalize_object_response(response, model, output_schema)}
        {:error, reason} -> {:error, reason}
      end
    end

    defp to_llm_schema(%Zoi.Types.Struct{fields: fields}), do: Zoi.object(fields, strict: true)

    defp to_llm_schema(%Zoi.Types.Union{schemas: schemas} = union) do
      converted_schemas = Enum.map(schemas, &to_llm_schema/1)
      %{union | schemas: converted_schemas}
    end

    defp to_llm_schema(schema), do: schema

    @impl true
    def stream(config, messages, opts) do
      model = Map.fetch!(config, :model)
      output_schema = Keyword.get(opts, :output_schema)
      backend_opts = Keyword.get(opts, :backend_opts, [])
      options = build_options(config, backend_opts)
      req_messages = Enum.map(messages, &to_req_llm_message/1)

      if output_schema do
        stream_with_schema(model, req_messages, output_schema, options)
      else
        stream_text(model, req_messages, options)
      end
    end

    defp stream_text(model, req_messages, options) do
      case ReqLLM.stream_text(model, req_messages, options) do
        {:ok, stream_response} ->
          chunk_stream =
            stream_response.stream
            |> Stream.flat_map(fn chunk ->
              case chunk.type do
                :content ->
                  [%{type: :content, content: chunk.text, metadata: %{backend: :req_llm}}]

                :thinking ->
                  [%{type: :thinking, content: chunk.text, metadata: %{backend: :req_llm}}]

                _ ->
                  []
              end
            end)

          {:ok, chunk_stream}

        {:error, reason} ->
          {:error, reason}
      end
    end

    defp stream_with_schema(model, messages, output_schema, options) do
      llm_schema = to_llm_schema(output_schema)

      case ReqLLM.stream_object(model, messages, llm_schema, options) do
        {:ok, stream_response} ->
          chunk_stream =
            stream_response.stream
            |> Stream.transform("", &accumulate_and_parse_json(&1, &2, output_schema))

          {:ok, chunk_stream}

        {:error, reason} ->
          {:error, reason}
      end
    end

    defp accumulate_and_parse_json(%{type: :content, text: text}, acc, schema) do
      new_acc = acc <> (text || "")

      case try_parse_json(new_acc) do
        {:ok, object} ->
          parsed = maybe_parse_struct(schema, object)

          chunk = %{
            type: :content,
            content: parsed,
            metadata: %{partial: true, backend: :req_llm}
          }

          {[chunk], new_acc}

        :error ->
          {[], new_acc}
      end
    end

    defp accumulate_and_parse_json(%{type: :thinking, text: text}, acc, _schema) do
      chunk = %{type: :thinking, content: text, metadata: %{backend: :req_llm}}
      {[chunk], acc}
    end

    defp accumulate_and_parse_json(_chunk, acc, _schema), do: {[], acc}

    defp try_parse_json(text) do
      case Jason.decode(text) do
        {:ok, object} when is_map(object) -> {:ok, object}
        _ -> :error
      end
    end

    @impl true
    def introspect(config) do
      model = Map.get(config, :model, "unknown")
      {provider, _model_name} = parse_model_string(model)

      %{
        provider: provider,
        model: model,
        operation: :chat,
        capabilities: [:streaming, :multi_provider]
      }
    end

    # Private helpers

    defp build_options(config, backend_opts) do
      config
      |> Map.drop([:model])
      |> Map.to_list()
      |> Keyword.merge(backend_opts)
      |> Enum.reject(fn {_k, v} -> is_nil(v) end)
    end

    defp normalize_response(response, model) do
      content = ReqLLM.Response.text(response)
      thinking = ReqLLM.Response.thinking(response)
      usage = ReqLLM.Response.usage(response) || %{}
      finish_reason = ReqLLM.Response.finish_reason(response)

      Response.new(
        content: content,
        thinking: thinking,
        finish_reason: normalize_finish_reason(finish_reason),
        usage: normalize_usage(usage, response),
        metadata: build_metadata(response, model, finish_reason)
      )
    end

    defp normalize_object_response(response, model, output_schema) do
      object = ReqLLM.Response.object(response)
      content = maybe_parse_struct(output_schema, object)
      thinking = ReqLLM.Response.thinking(response)
      usage = ReqLLM.Response.usage(response) || %{}
      finish_reason = ReqLLM.Response.finish_reason(response)

      Response.new(
        content: content,
        thinking: thinking,
        finish_reason: normalize_finish_reason(finish_reason),
        usage: normalize_usage(usage, response),
        metadata: build_metadata(response, model, finish_reason)
      )
    end

    defp maybe_parse_struct(%Zoi.Types.Struct{} = schema, object) when is_map(object) do
      case Zoi.parse(schema, object) do
        {:ok, struct} -> struct
        {:error, _} -> object
      end
    end

    defp maybe_parse_struct(%Zoi.Types.Union{} = schema, object) when is_map(object) do
      case Zoi.parse(schema, object) do
        {:ok, parsed} -> parsed
        {:error, _} -> object
      end
    end

    defp maybe_parse_struct(_schema, object), do: object

    defp build_metadata(response, model, raw_finish_reason) do
      {provider, _model_name} = parse_model_string(model)

      %{
        provider: provider,
        model: response.model,
        response_id: response.id,
        raw_finish_reason: raw_finish_reason
      }
      |> Enum.reject(fn {_k, v} -> is_nil(v) end)
      |> Map.new()
    end

    defp parse_model_string(model) when is_binary(model) do
      case String.split(model, ":", parts: 2) do
        [provider, model_name] -> {provider, model_name}
        [model_name] -> {"unknown", model_name}
      end
    end

    defp parse_model_string(%LLMDB.Model{provider: provider, model: model_name}) do
      {to_string(provider), model_name}
    end

    defp normalize_finish_reason(:stop), do: :stop
    defp normalize_finish_reason(:length), do: :max_tokens
    defp normalize_finish_reason(:content_filter), do: :content_filter
    defp normalize_finish_reason(:error), do: :error
    defp normalize_finish_reason(nil), do: nil

    defp normalize_usage(usage, response) when is_map(usage) do
      thinking_tokens = ReqLLM.Response.reasoning_tokens(response)

      base = %{
        input_tokens: Map.get(usage, :input_tokens, 0),
        output_tokens: Map.get(usage, :output_tokens, 0)
      }

      if thinking_tokens > 0 do
        Map.put(base, :thinking_tokens, thinking_tokens)
      else
        base
      end
    end

    defp normalize_usage(_, _response), do: %{}

    # Convert Puck.Message → ReqLLM message format
    # Single text → use plain map with string content (loose map)
    # Multi-part → use ReqLLM.Message struct with ContentPart list
    defp to_req_llm_message(%Message{role: role, content: [%Part{type: :text, text: text}]}) do
      %{role: role, content: text}
    end

    defp to_req_llm_message(%Message{role: role, content: parts}) do
      content_parts = Enum.map(parts, &to_req_llm_part/1)
      %ReqLLM.Message{role: role, content: content_parts}
    end

    defp to_req_llm_part(%Part{type: :text, text: text}) do
      ContentPart.text(text)
    end

    defp to_req_llm_part(%Part{type: :image_url, url: url}) do
      ContentPart.image_url(url)
    end

    defp to_req_llm_part(%Part{type: :image, data: data, media_type: media_type}) do
      ContentPart.image(data, media_type)
    end

    defp to_req_llm_part(%Part{
           type: :file,
           data: data,
           filename: filename,
           media_type: media_type
         }) do
      ContentPart.file(data, filename, media_type)
    end

    defp to_req_llm_part(%Part{type: type} = part) do
      # Fallback for other types (audio, video, custom) - pass as map
      part |> Map.from_struct() |> Map.delete(:metadata) |> Map.put(:type, type)
    end
  end
end